From 4896d7354ada6483c053efd9df61d197e4d13608 Mon Sep 17 00:00:00 2001 From: Maddie Cain Date: Sun, 12 May 2019 22:58:37 -0400 Subject: [PATCH 1/2] fixed ePIE bug and shortened inspect function --- CDTools/models/simple_ptycho.py | 65 ++++++++++++++++++++------------- CDTools/tools/projectors.py | 7 ++-- examples/simple_ptycho.py | 17 +++++---- 3 files changed, 53 insertions(+), 36 deletions(-) diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index de6c5a5..48d9388 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -7,7 +7,7 @@ from CDTools.tools import plotting as p from copy import copy from torch.utils import data as torchdata from matplotlib import pyplot as plt - +import numpy as np class SimplePtycho(CDIModel): @@ -45,6 +45,7 @@ class SimplePtycho(CDIModel): self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + @classmethod def from_dataset(cls, dataset): wavelength = dataset.wavelength @@ -138,14 +139,10 @@ class SimplePtycho(CDIModel): def inspect(self): - p.plot_amplitude(self.probe, basis=self.probe_basis) - plt.title('Probe Amplitude') - p.plot_phase(self.probe, basis=self.probe_basis) - plt.title('Probe Phase') - p.plot_amplitude(self.obj, basis=self.probe_basis) - plt.title('Object Amplitude') - p.plot_phase(self.obj, basis=self.probe_basis) - plt.title('Object Phase') + p.plot_amplitude(self.probe, basis=self.probe_basis, title = 'Probe Amplitude') + p.plot_phase(self.probe, basis=self.probe_basis, title = 'Probe Phase') + p.plot_amplitude(self.obj, basis=self.probe_basis, title = 'Object Amplitude') + p.plot_phase(self.obj, basis=self.probe_basis, title = 'Object Phase') def save_results(self): @@ -153,8 +150,8 @@ class SimplePtycho(CDIModel): probe = probe * self.probe_norm.detach().cpu().numpy() obj = tools.cmath.torch_to_complex(self.obj.detach().cpu()) return {'probe':probe,'obj':obj} - - + + def ePIE(self, iterations, dataset, beta = 1.0): """Runs an ePIE reconstruction as described in `Maiden et al. (2017) `_. Optional parameters are: @@ -172,13 +169,13 @@ class SimplePtycho(CDIModel): 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) + new_probe = probe + tools.cmath.cmult(beta * tools.cmath.cconj(object[translation])/(self.probe_norm*t.max(tools.cmath.cabssq(object))), exit_wave_corrected-exit_wave) + return new_probe 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 + new_object = object.clone() + new_object[translation] = object[translation] + tools.cmath.cmult(beta * tools.cmath.cconj(probe)/(self.probe_norm*t.max(tools.cmath.cabssq(probe))), exit_wave_corrected-exit_wave) + return new_object with t.no_grad(): data_loader = torchdata.DataLoader(dataset, shuffle=True) @@ -186,15 +183,33 @@ class SimplePtycho(CDIModel): 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) + probe = self.probe.data.clone() + object = self.obj.data.clone() - 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) + exit_wave = self.interaction(i, translations).clone() + # Apply modulus constraint + exit_wave_corrected = exit_wave.clone() + exit_wave_corrected = self.forward_propagator(exit_wave_corrected.clone()) + exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(exit_wave_corrected.clone()[self.detector_slice], patterns, mask = mask) + exit_wave_corrected = self.backward_propagator(exit_wave_corrected.clone()) + + # Calculate the section of the object wavefunction to be modified + pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, + translations) + pix_trans -= self.min_translation + + pix_trans = t.round(pix_trans).to(dtype=t.int32).numpy() + + object_slice = np.s_[pix_trans[0]: + pix_trans[0]+probe_shape[0], + pix_trans[1]: + pix_trans[1]+probe_shape[1]] + + # Apply probe and object updates + self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, object_slice) + self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, object_slice) + + # Calculate loss loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns)) - yield t.mean(t.Tensor(loss)).cpu().numpy() + yield it, t.mean(t.Tensor(loss)).cpu().numpy() diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors.py index f4f9f8b..b4b2a04 100644 --- a/CDTools/tools/projectors.py +++ b/CDTools/tools/projectors.py @@ -22,7 +22,7 @@ def modulus(wavefront, intensities, mask = None): wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1 - + Returns: torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities """ @@ -55,10 +55,9 @@ def support(wavefront, support): Args: wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts - support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside - + support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside + Returns: torch.Tensor : The JxNxMx2 wavefield with the support mask applied """ return wavefront * support.to(wavefront.dtype)[...,None] - diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 5a97b4a..c604dfd 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -12,8 +12,8 @@ import pickle #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' +filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi' +#filename = '../../Desktop/Reconstructions/114429_p.cxi' with h5py.File(filename,'r') as f: @@ -24,16 +24,19 @@ 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='cpu') #dataset.to(device='cuda') -dataset.get_as(device='cuda') +dataset.get_as(device='cpu') +for i, loss in enumerate(model.ePIE(10, dataset)): + print(i, loss) -for loss in model.Adam_optimize(1, dataset): - print(loss) +for i, loss in enumerate(model.Adam_optimize(10, dataset)): + print(i, loss) + +model.inspect() #with open('test_results.pickle', 'wb') as f: # pickle.dump(model.save_results(),f) -model.inspect() plt.show() From e7c0e4771c0f092dd1283eb74dce9e1cd8a15d86 Mon Sep 17 00:00:00 2001 From: Maddie Cain Date: Tue, 14 May 2019 01:38:15 -0400 Subject: [PATCH 2/2] generate diffraction data with errors added for debugging on gpu --- examples/generate_data.py | 86 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 examples/generate_data.py diff --git a/examples/generate_data.py b/examples/generate_data.py new file mode 100644 index 0000000..ce7716e --- /dev/null +++ b/examples/generate_data.py @@ -0,0 +1,86 @@ +from __future__ import division, print_function, absolute_import + +import CDTools +from CDTools.tools.initializers import gaussian +from CDTools.tools.propagators import far_field +from CDTools.tools.measurements import intensity + +from CDTools.tools.cmath import * +from CDTools.tools import interactions +import h5py +import numpy as np +from matplotlib import pyplot as plt +import torch as t +import scipy.misc + + + +def simulate_pattern(pixel_translation_step = (10, 10), probe = None, intensity_err = True, + random_noise = True, translation_err = True, background_noise = True): + """ + Generates diffraction data from scipy's face image with a gaussian phase. + Can simulate different experimental errors, including: + * Varying probe intensities over time + * Nanopositioner location uncertainties + * Background detector noise + * Random noise + * Uncentered probe + + We can do this by simulating the data collection process with a large + probe array, and changing the translation by a random (integer) amount. + Then, we can scale the intensity by a random amount (close to 1), + replace pixels with some scaled version of their current value to simulate + random noise, then add in horizontal and vertical detector background noise. + Finally, we slice the simulated diffraction patterns to uncenter the probe + on the final patterns [TO DO]. + + Args: + pixel_translation_step (array-like) : a tuple containing (y pixel translation, x pixel translation) + probe (np.ndarray) : a two-dimensional array representing the probe wavefunction. Defaults to a gaussian probe + intensity_err (Boolean) : Defaults to True, to simulate varying probe intensities over time + random_noise (Boolean) : Defaults to True, indicating adding random noise to the diffraction patterns + translation_err (Boolean) : Defaults to True, to simulate nanopositioner errors + background_noise (Boolean) : [TO DO] + Returns: + patterns (t.Tensor) : NxMx2 array of simulated diffraction patterns with error + translations (t.Tensor) : N*Mx2 array of instrument translations + """ + obj = scipy.misc.face() + # Convert to grayscale + obj = 0.2989 * obj[:,:,0] + 0.5870 * obj[:,:,1] + 0.1140 * obj[:,:,2] + phase = gaussian(obj.shape, [500,500], amplitude=1, center = None) + obj = obj*np.exp(1j*torch_to_complex(phase)) + + if probe is None: + probe = torch_to_complex(gaussian([512,512], [5,5], amplitude=1, center = None)) + + obj_shape = np.array(obj.shape) + probe_shape = np.array(probe.shape) + + translation_range = (obj_shape-probe_shape)//np.array(pixel_translation_step)*np.array(pixel_translation_step) + translations = np.mgrid[0:translation_range[0]+pixel_translation_step[0]:pixel_translation_step[0], 0:translation_range[1]+pixel_translation_step[1]:pixel_translation_step[1]] + translations = translations.T.reshape((translations.T.shape[0]*translations.T.shape[1], 2)) + translations = t.tensor(translations, dtype = t.float32) + ideal_translations = t.tensor(translations, dtype = t.float32) + obj = complex_to_torch(obj) + probe = complex_to_torch(probe) + + if translation_err: + err = t.normal(mean=t.zeros(translations.shape), std=t.ones(translations.shape)/2) + translations += err + # Ensure that the new translations don't go off the object area + translations[translations<0] = 0 + translation_range = t.tensor(translation_range, dtype = t.float32) + translations[...,0][translations[...,0] > translation_range[0]] = translation_range[0] + translations[...,1][translations[...,1] > translation_range[1]] = translation_range[1] + + patterns = intensity(far_field(interactions.ptycho_2D_round(probe, obj, translations))) + + if intensity_err: + patterns *= t.normal(mean = t.ones(patterns.shape), std = t.ones(patterns.shape)/32) + + if background_noise: + patterns += t.normal(mean = t.zeeros(patterns.shape), std = t.ones(patterns.shape)*t.max(patterns)/32) + + + return patterns, ideal_translations