From 2a717d913305728b9d31cda6b236f731f6de9b82 Mon Sep 17 00:00:00 2001 From: Maddie Cain Date: Thu, 18 Apr 2019 01:01:37 -0400 Subject: [PATCH 1/2] fixed bug in centroid and began adding to simple_ptycho --- CDTools/models/__init__.py | 52 ++++++++++++++-------------- CDTools/models/simple_ptycho.py | 56 +++++++++++++++++++++++-------- CDTools/tools/image_processing.py | 10 +++--- examples/simple_ptycho.py | 11 +++--- 4 files changed, 80 insertions(+), 49 deletions(-) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 7befb04..d135d03 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): @@ -129,13 +129,13 @@ class CDIModel(t.nn.Module): data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, shuffle=True) - + # Define the optimizer optimizer = t.optim.Adam(self.parameters(), lr = lr) - - return self.AD_optimize(iterations, data_loader, optimizer) - - + 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.fancy_ptycho import FancyPtycho diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 2004704..ff68caf 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -10,7 +10,7 @@ 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 +36,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 +56,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 +64,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 +72,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 +80,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 +104,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 +124,41 @@ 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(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. + """ + """def probe_update(exit_wave, exit_wave_corrected, probe, object, i): + return probe+tools.cmult(beta*(tools.cconj(object)/ \ + t.max(tools.cabssq(object)))[self.translations[i][0]:self.translations[i][0]+self.shape[0],self.translations[i][1]:self.translations[i][1]+self.shape[1]], \ + exit_wave_corrected-exit_wave) + def object_update(exit_wave, exit_wave_corrected, probe, object, i): + object[self.translations[i][0]:self.translations[i][0]+self.shape[0],self.translations[i][1]:self.translations[i][1]+self.shape[1]]\ + +=tools.cmult(tools.cconj(probe)/t.max(tools.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): + for translations, patterns in data_loader: + translation = translations[1] + probe = self.probe.clone() + object = self.obj.clone() + exit_wave = tools.cmult(probe, object[self.translations[i][0]:self.translations[i][0]+self.shape[0],self.translations[i][1]:self.translations[i][1]+self.shape[1]]).clone() + exit_wave_corrected = projectors.modulus(exit_wave, self.data[i], self.propagator).clone() + self.probe = probe_update(exit_wave, exit_wave_corrected, probe, object, i).clone() + self.obj = object_update(exit_wave, exit_wave_corrected, probe, object, i).clone() diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 6d0174c..b01b386 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -21,7 +21,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)] @@ -73,7 +74,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): """ pass - + def find_pixel_shift(im1, im2): """Calculates the integer pixel shift between two images by maximizing the autocorrelation @@ -94,10 +95,10 @@ def find_pixel_shift(im1, im2): if im2.shape[-1] != 2: im2 = t.stack((im2,t.zeros_like(im2)),dim=-1) - + cor = cmath.cabs(t.ifft(cmath.cmult(t.fft(im1,2), cmath.cconj(t.fft(im2,2))),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) @@ -122,4 +123,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 7e65f43..34e9b71 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -7,8 +7,9 @@ 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/114429_p.cxi' +#filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi' +#filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi' @@ -26,7 +27,10 @@ dataset.get_as(device='cuda') for loss in model.Adam_optimize(100, dataset): print(loss) - + +for loss in model.ePIE(100, dataset): + print(loss) + from matplotlib import pyplot as plt plot_amplitude(model.probe) @@ -34,4 +38,3 @@ plot_phase(model.probe) plot_amplitude(model.obj) plot_phase(model.obj) plt.show() - From 84ac2cc0589176787815c39b9a2abfc4dc1115bc Mon Sep 17 00:00:00 2001 From: Maddie Cain Date: Wed, 24 Apr 2019 12:26:32 -0400 Subject: [PATCH 2/2] added ePIE --- CDTools/models/simple_ptycho.py | 34 +++++++++++++++++++++++---------- examples/simple_ptycho.py | 14 ++++++++------ 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index b0ac9ba..dc9d27f 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -4,6 +4,7 @@ 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): @@ -134,7 +135,7 @@ class SimplePtycho(CDIModel): pass - def ePIE(iterations, dataset, beta = 1.0): + def ePIE(self, iterations, dataset, beta = 1.0): """Runs an ePIE reconstruction as described in `Maiden et al. (2017) `_. Optional parameters are: @@ -144,21 +145,34 @@ class SimplePtycho(CDIModel): :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.cmult(beta*(tools.cconj(object)/ \ - t.max(tools.cabssq(object)))[translations[0]:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]], \ + 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[translations[0]:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]]\ - +=tools.cmult(tools.cconj(probe)/t.max(tools.cabssq(probe)),exit_wave_corrected-exit_wave) + 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): - for translations, patterns in data_loader: + loss = [] + for (i, [translations]), [patterns] in data_loader: probe = self.probe.clone() object = self.obj.clone() - exit_wave = tools.cmult(probe, object[translations:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]]).clone() - exit_wave_corrected = projectors.modulus(self.forward_propagator(exit_wave), patterns, mask = self.mask) - self.probe = probe_update(exit_wave, exit_wave_corrected, probe, object, translations) - self.obj = object_update(exit_wave, exit_wave_corrected, probe, object, translations) + 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/examples/simple_ptycho.py b/examples/simple_ptycho.py index 72ef0c6..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,12 +23,12 @@ 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') -for loss in model.ePIE(100, dataset): +for loss in model.ePIE(1, dataset): print(loss) from matplotlib import pyplot as plt