diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 4e7a76e..72034f8 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -13,8 +13,8 @@ from torch.utils import data as torchdata # I gain in it being easy to reload the non-learned aspects of the model, # like the wavelength and sample geometry. Remember that it's important # that the final outputs of the reconstructions are transferrable to other -# places -# +# places +# # @@ -26,14 +26,14 @@ from torch.utils import data as torchdata class CDIModel(t.nn.Module): """This base model defines all the functions that must be exposed for a valid CDIModel subclass - + Most of the functions only raise a NotImplementedError at this level and must be explicitly defined by any subclass. The functions required can be split into several subsections: - + Creation: from_dataset : a function to create a CDIModel from an appropriate CDataset - + Simulation: interaction : a function to simulate exit waves from experimental parameters forward_propagator : the propagator from the experiment plane to the detector plane @@ -47,33 +47,33 @@ class CDIModel(t.nn.Module): Reconstruction: AD_optimize : predefined, a generic automatic differentiation reconstruction Adam_optimize : predefined, sensible automatic differentiation reconstruction using ADAM - + The work of defining the various subclasses boils down to creating an appropriate implementation for this set of functions. """ - - - + + + def from_dataset(self, dataset): raise NotImplementedError() - + def interaction(self, *args): raise NotImplementedError() - + def forward_propagator(self, exit_wave): raise NotImplementedError() def backward_propagator(self, detector_wave): raise NotImplementedError() - - + + def measurement(self, detector_wave): raise NotImplementedError() - + def forward(self, *args): return self.measurement(self.forward_propagator(self.interaction(*args))) @@ -86,15 +86,15 @@ class CDIModel(t.nn.Module): def to(self, *args, **kwargs): super(CDIModel,self).to(*args,**kwargs) - + def simulate(self, args_list): return t.Tensor([self.forward(*args) for args in args_list]) - + def simulate_to_dataset(self, args_list): raise NotImplementedError() - + def AD_optimize(self, iterations, data_loader, optimizer, scheduler=None): for it in range(iterations): @@ -102,7 +102,7 @@ class CDIModel(t.nn.Module): N = 0 for inputs, patterns in data_loader: N += 1 - + def closure(): optimizer.zero_grad() sim_patterns = self.forward(*inputs) @@ -113,7 +113,7 @@ class CDIModel(t.nn.Module): loss.backward() return loss - + loss += optimizer.step(closure).detach().cpu().numpy() loss /= N @@ -121,7 +121,7 @@ class CDIModel(t.nn.Module): scheduler.step(loss) yield loss - + def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False): @@ -132,13 +132,13 @@ class CDIModel(t.nn.Module): # Define the optimizer optimizer = t.optim.Adam(self.parameters(), lr = lr) - + # Define the scheduler if schedule: scheduler = t.optim.ReduceLROnPlateau(optimizer, factor=0.2) else: scheduler = None - + return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler) @@ -152,15 +152,16 @@ class CDIModel(t.nn.Module): else: data_loader = torchdata.DataLoader(dataset) - + # Define the optimizer optimizer = t.optim.LBFGS(self.parameters(), lr = lr, history_size=history_size) - + return self.AD_optimize(iterations, data_loader, optimizer) - + + from CDTools.models.simple_ptycho import SimplePtycho from CDTools.models.fancy_ptycho import FancyPtycho from CDTools.models.incoherent_ptycho import IncoherentPtycho diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 2004704..dc9d27f 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -4,13 +4,14 @@ import torch as t from CDTools.models import CDIModel from CDTools import tools from copy import copy +from torch.utils import data as torchdata class SimplePtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, detector_slice, - probe_guess, obj_guess, min_translation = t.Tensor([0,0]), + probe_guess, obj_guess, min_translation = [0,0], mask=None): super(SimplePtycho,self).__init__() @@ -36,7 +37,7 @@ class SimplePtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32))) - + self.probe = t.nn.Parameter(probe_guess.to(t.float32) / self.probe_norm) self.obj = t.nn.Parameter(obj_guess.to(t.float32)) @@ -56,7 +57,7 @@ class SimplePtycho(CDIModel): dataset.get_as(*get_as_args[0],**get_as_args[1]) center = tools.image_processing.centroid(t.sum(patterns,dim=0)) - + # Then, generate the probe geometry from the dataset ewg = tools.initializers.exit_wave_geometry probe_basis, probe_shape, det_slice = ewg(det_basis, @@ -64,7 +65,7 @@ class SimplePtycho(CDIModel): wavelength, distance, center=center) - + # Next generate the object geometry from the probe geometry and # the translations pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations) @@ -72,7 +73,7 @@ class SimplePtycho(CDIModel): # Finally, initialize the probe and object using this information probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice) - + obj = t.ones(obj_size+(2,)) det_geo = dataset.detector_geometry @@ -80,10 +81,10 @@ class SimplePtycho(CDIModel): mask = dataset.mask.to(t.uint8) else: mask = None - + return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask) - - + + def interaction(self, index, translations): pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, translations) @@ -104,12 +105,12 @@ class SimplePtycho(CDIModel): def measurement(self, wavefields): return tools.measurements.intensity(wavefields, detector_slice=self.detector_slice) - + def loss(self, sim_data, real_data, mask=None): - return tools.losses.amplitude_mse(real_data, sim_data,mask=mask) + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + - def to(self, *args, **kwargs): super(SimplePtycho, self).to(*args, **kwargs) self.wavelength = self.wavelength.to(*args,**kwargs) @@ -124,13 +125,54 @@ class SimplePtycho(CDIModel): if self.mask is not None: self.mask = self.mask.to(*args, **kwargs) - + self.min_translation = self.min_translation.to(*args,**kwargs) self.probe_basis = self.probe_basis.to(*args,**kwargs) self.probe_norm = self.probe_norm.to(*args,**kwargs) - + def sim_to_dataset(self, args_list): pass - + + def ePIE(self, iterations, dataset, beta = 1.0): + """Runs an ePIE reconstruction as described in `Maiden et al. (2017) `_. + Optional parameters are: + + :arg ``iterations``: Controls the number of iterations run, defaults to 1. + :arg ``beta``: Algorithmic parameter described in Maiden's implementation of rPIE. Defaults to 0.15. + :arg ``probe``: Initial probe wavefunction. + :arg ``object``: Initial object wavefunction. + """ + probe_shape = self.probe.shape + + if self.mask is not None: + mask = self.mask[...,None] + + def probe_update(exit_wave, exit_wave_corrected, probe, object, translation): + return probe+tools.cmath.cmult(beta*(tools.cmath.cconj(object)/t.max(tools.cmath.cabssq(object)))[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]], \ + exit_wave_corrected-exit_wave) + + def object_update(exit_wave, exit_wave_corrected, probe, object, translation): + object[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]]\ + +=tools.cmath.cmult(tools.cmath.cconj(probe)/t.max(tools.cmath.cabssq(probe)),exit_wave_corrected-exit_wave) + return object + + with t.no_grad(): + data_loader = torchdata.DataLoader(dataset, shuffle=True) + + for it in range(iterations): + loss = [] + for (i, [translations]), [patterns] in data_loader: + probe = self.probe.clone() + object = self.obj.clone() + exit_wave = self.interaction(i, translations) + exit_wave_corrected = exit_wave.clone() + exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(self.forward_propagator(exit_wave)[self.detector_slice], patterns, mask = mask) + + integer_translations = t.round(translations).to(dtype=t.int32) + self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, integer_translations) + self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, integer_translations) + loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns)) + + yield t.mean(t.Tensor(loss)).cpu().numpy() diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 322889f..5eea286 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -23,7 +23,8 @@ def centroid(im, dims=2): Returns: t.Tensor : An (i,j) index or stack of indices """ - indices = (t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)) + # For some reason this needs to be a list + indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)] indices = t.meshgrid(*indices) use_dims = [-dims+i for i in range(dims)] @@ -61,7 +62,7 @@ def centroid_sq(im, dims=2, comp=False): def sinc_subpixel_shift(im, shift): """Performs a subpixel shift with sinc interpolation on the given tensor - + The subpixel shift is done circularly via a multiplication with a linear phase mask in Fourier space. @@ -120,15 +121,15 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): # Not sure if this is more or less stable than just the correlation # maximum - requires some testing cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2) - - + + # Now, I need to shift the array to pull out a contiguous window # around the correlation maximum try: search_around = search_around.cpu() except: search_around = t.tensor(search_around) - + window_size = 15 shift_zero = tuple(-search_around + t.tensor([window_size,window_size])) cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size] @@ -148,11 +149,11 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): cormax = t.tensor([t.argmax(upsampled) // sh[1], t.argmax(upsampled) % sh[1]]).to(device=upsampled.device) subpixel_shift = ((cormax + sh // 2) % sh - sh//2).to(dtype=upsampled.dtype) - + return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \ - subpixel_shift / resolution - - + subpixel_shift / resolution + + def find_pixel_shift(im1, im2): """Calculates the integer pixel shift between two images by maximizing the autocorrelation @@ -179,8 +180,8 @@ def find_pixel_shift(im1, im2): # Not sure if this is more or less stable than just the correlation # maximum - requires some testing cor = cmath.cabs(t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)) - - + + sh = t.tensor(cor.shape).to(device=im1.device) cormax = t.tensor([t.argmax(cor) // sh[1], t.argmax(cor) % sh[1]]).to(device=im1.device) @@ -205,4 +206,3 @@ def find_shift(im1, im2, resolution=10): resolution=resolution) return subpixel_shift - diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 54daa27..f5a0f19 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -7,9 +7,11 @@ import h5py import torch as t import numpy as np -filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi' -filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi' -filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi' + +filename = '../../Downloads/114429_p.cxi' +#filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi' +#filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi' +#filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi' with h5py.File(filename,'r') as f: @@ -21,17 +23,14 @@ model = CDTools.models.SimplePtycho.from_dataset(dataset) # Uncomment these to use on the CPU # default is CPU with 32-bit floats -model.to(device='cuda') +#model.to(device='cuda') #dataset.to(device='cuda') -dataset.get_as(device='cuda') +#dataset.get_as(device='cuda') -#model.probe.requires_grad = False -for loss in model.Adam_optimize(10, dataset): + +for loss in model.ePIE(1, dataset): print(loss) -#for loss in model.Adam_optimize(20, dataset, lr=0.0005): -# print(loss) - from matplotlib import pyplot as plt plot_amplitude(model.probe) @@ -39,4 +38,3 @@ plot_phase(model.probe) plot_amplitude(model.obj) plot_phase(model.obj) plt.show() -