diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 0166734..9e28501 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -39,7 +39,7 @@ import numpy as np import threading import queue import time -import pytorch_warmup +#import pytorch_warmup __all__ = ['CDIModel'] @@ -282,7 +282,8 @@ class CDIModel(t.nn.Module): scheduler = None if warmup: - warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer) + print('Warmup is not currently implemented, sorry!') + #warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer) else: warmup_scheduler = None diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 310c3e9..5d0ada9 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -122,7 +122,7 @@ class FancyPtycho(CDIModel): @classmethod - def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, loss='amplitude mse', units='um'): + def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -496,7 +496,7 @@ class FancyPtycho(CDIModel): weights = cmath.torch_to_complex(self.weights.detach().cpu()) - + probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)* weights.conj(),axis=2) diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 8b449c2..db82260 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -329,7 +329,10 @@ class Multislice2DPtycho(CDIModel): exit_waves = tools.interactions.ptycho_2D_round( exit_waves, obj, pix_trans, multiple_modes=True) - + #if self.iteration_count >= 10: + # plt.figure() + # plt.imshow(np.abs(cmath.torch_to_complex(exit_waves[0,0].detach().cpu()))) + # plt.show() elif self.obj.dim() == 4: @@ -346,17 +349,19 @@ class Multislice2DPtycho(CDIModel): if i < self.nz-1: #on all but the last iteration exit_waves = tools.propagators.near_field( exit_waves,self.as_prop) - return exit_waves def forward_propagator(self, wavefields): - left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2] - right = [self.probe.shape[-3]//2+self.probe.shape[-3], - self.probe.shape[-2]//2+self.probe.shape[-2]] - - return tools.propagators.far_field(wavefields)[...,left[0]:right[0], - left[1]:right[1],:] + if self.prevent_aliasing: + left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2] + right = [self.probe.shape[-3]//2+self.probe.shape[-3], + self.probe.shape[-2]//2+self.probe.shape[-2]] + + return tools.propagators.far_field(wavefields)[...,left[0]:right[0], + left[1]:right[1],:] + else: + return tools.propagators.far_field(wavefields) def measurement(self, wavefields): diff --git a/CDTools/models/rpi.py b/CDTools/models/rpi.py index 7bc35bf..6fe4a84 100644 --- a/CDTools/models/rpi.py +++ b/CDTools/models/rpi.py @@ -266,7 +266,7 @@ class RPI(CDIModel): try: # will fail if index has no length, for example when index # is just an int. In this case, we just do nothing instead - output = output.unsqueeze(1).repeat(1,len(index),1,1,1) + output = output.unsqueeze(0).repeat(1,len(index),1,1,1) except TypeError: pass diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index 0dfe6ed..65abcef 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -354,7 +354,11 @@ def get_data(cxi_file, cut_zeroes = True): data[data < 0] = 0 if 'axes' in cxi_file[pull_from].attrs: - axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':') + try: + axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':') + except AttributeError as e: # Weird string vs bytes thing, ehhh + axes = str(cxi_file[pull_from].attrs['axes']).split(':') + axes = [axis.strip().lower() for axis in axes] else: axes = None diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 9c49978..7538a1a 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -34,7 +34,7 @@ dataset.get_as(device='cuda') # The regularization is an L2 regularizer that empirically helps accelerate # convergence for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)): - #model.inspect(dataset) + model.inspect(dataset) print(i,loss) #model.inspect(dataset)