From 91f11643e91cbcaaf7ace189a1ec71454491a53c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 10 Apr 2019 10:25:05 -0400 Subject: [PATCH] Make some tweaks to the simple ptycho model --- CDTools/models/__init__.py | 3 +- CDTools/models/simple_ptycho.py | 24 ++++++++--- CDTools/tools/cmath.py | 16 ++++++- CDTools/tools/image_processing.py | 28 ++++++++++++- CDTools/tools/interactions.py | 62 ++++++++++++++++++++++++++++ CDTools/tools/measurements.py | 22 +++++----- CDTools/tools/plotting.py | 3 ++ examples/simple_ptycho.py | 27 ++++++------ tests/tools/test_image_processing.py | 30 ++++++++++++++ tests/tools/test_measurements.py | 27 +++++++----- 10 files changed, 199 insertions(+), 43 deletions(-) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 95760a1..7befb04 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -101,7 +101,7 @@ class CDIModel(t.nn.Module): loss = 0 N = 0 for inputs, patterns in data_loader: - N += patterns.shape[0] + N += 1 def closure(): optimizer.zero_grad() @@ -138,3 +138,4 @@ class CDIModel(t.nn.Module): from CDTools.models.simple_ptycho import SimplePtycho +from CDTools.models.fancy_ptycho import FancyPtycho diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 958e37a..2004704 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -11,7 +11,7 @@ class SimplePtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, detector_slice, probe_guess, obj_guess, min_translation = t.Tensor([0,0]), - background = None): + mask=None): super(SimplePtycho,self).__init__() self.wavelength = t.Tensor([wavelength]) @@ -28,6 +28,10 @@ class SimplePtycho(CDIModel): self.probe_basis = t.Tensor(probe_basis) self.detector_slice = detector_slice + if mask is None: + self.mask = None + else: + self.mask = t.ByteTensor(mask) # We rescale the probe here so it learns at the same rate as the # object @@ -68,10 +72,16 @@ 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 - - return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation) + + if hasattr(dataset, 'mask') and dataset.mask is not None: + 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): @@ -82,6 +92,7 @@ class SimplePtycho(CDIModel): self.obj, pix_trans) + def forward_propagator(self, wavefields): return tools.propagators.far_field(wavefields) @@ -95,8 +106,8 @@ class SimplePtycho(CDIModel): detector_slice=self.detector_slice) - def loss(self, sim_data, real_data): - return tools.losses.amplitude_mse(real_data, sim_data) + def loss(self, sim_data, real_data, mask=None): + return tools.losses.amplitude_mse(real_data, sim_data,mask=mask) def to(self, *args, **kwargs): @@ -111,6 +122,9 @@ class SimplePtycho(CDIModel): if hasattr(det_geo, 'corner'): det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) + 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) diff --git a/CDTools/tools/cmath.py b/CDTools/tools/cmath.py index 2335033..7a89d4e 100644 --- a/CDTools/tools/cmath.py +++ b/CDTools/tools/cmath.py @@ -13,7 +13,7 @@ import torch as t __all__ = ['complex_to_torch','torch_to_complex','cabssq','cabs','cconj', - 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift'] + 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift','expi'] # @@ -239,3 +239,17 @@ def ifftshift(array,dims=None): array = t.cat((array.narrow(dim,cut_to,cut_len), array.narrow(dim,0,cut_to)), dim) return array + + +def expi(array): + """Returns a complex-format array for exp(i* (real array)) + + Expects the input to be in the form of a real-valued array + + Args: + array (torch.Tensor) : An array to be exponentiated + + Returns: + torch.Tensor : A complex-format array + """ + return t.stack((t.cos(array),t.sin(array)),dim=-1) diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 228c72c..6d0174c 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -71,12 +71,16 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): search_around (array_like) : Default (0,0), the shift to search in the vicinity of resolution (int): Default is 10, the resolution to calculate to in units of 1/n """ + pass + def find_pixel_shift(im1, im2): """Calculates the integer pixel shift between two images by maximizing the autocorrelation This function simply takes the circular correlation with an FFT and - returns the position of the maximum of that correlation + returns the position of the maximum of that correlation. This corresponds + to the amount that im1 would have to be shifted by to line up best with + im2 Args: im1 (t.Tensor): The first real or complex-valued torch tensor @@ -84,7 +88,21 @@ def find_pixel_shift(im1, im2): search_around (array_like) : Default (0,0), the shift to search in the vicinity of resolution (int): Default is 10, the resolution to calculate to in units of 1/n """ - pass + # If last dimension is not 2, then convert to a complex tensor now + if im1.shape[-1] != 2: + im1 = t.stack((im1,t.zeros_like(im1)),dim=-1) + 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) + return (cormax + sh // 2) % sh - sh//2 + def find_shift(im1, im2, resolution=10): @@ -99,3 +117,9 @@ def find_shift(im1, im2, resolution=10): im2 (t.Tensor): The second real or complex-valued torch tensor resolution (int): Default is 10, the resolution to calculate to in units of 1/n """ + integer_shift = find_pixel_shift(im1,im2) + subpixel_shift = find_subpixel_shift(im1, im2, search_around=integer_shift, + resolution=resolution) + + return subpixel_shift + diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py index 4100d3a..2ef06fd 100644 --- a/CDTools/tools/interactions.py +++ b/CDTools/tools/interactions.py @@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import from CDTools.tools.cmath import * import torch as t +import numpy as np @@ -178,3 +179,64 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): #TODO: Implement a sinc-interpolated shift using a fourier space shifting op +def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10): + """Returns a stack of exit waves accounting for subpixel shifts + + This function returns a collection of exit waves, with the first + dimension as the translation index and the final dimensions + corresponding to the detector. The exit waves are calculated by + shifting the probe with each translation in turn, using sinc + interpolation (done via multiplication with a complex exponential + in Fourier space) + + If shift_probe is True, it applies the subpixel shift to the probe, + otherwise the subpixel shift is applied to the object + + Args: + probe (torch.Tensor) : An MxL probe function for the exit waves + object (torch.Tensor) : The object function to be probed + translations (torch.Tensor) : The Nx2 array of translations to simulate + shift_probe (bool) : Whether to subpixel shift the probe or object + padding (int) : Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects + Returns: + torch.Tensor : An NxMxL tensor of the calculated exit waves + """ + single_translation = False + if translations.dim() == 1: + translations = translations[None,:] + single_translation = True + + # Separate the translations into a part that chooses the window + # And a part that defines the windowing function + integer_translations = t.floor(translations) + subpixel_translations = translations - integer_translations + integer_translations = integer_translations.to(dtype=t.int32) + + exit_waves = [] + if shift_probe: + i = t.arange(probe.shape[0]) - probe.shape[0]//2 + j = t.arange(probe.shape[1]) - probe.shape[1]//2 + I,J = t.meshgrid(i,j) + I = 2 * np.pi * I.to(t.float32) / probe.shape[0] + J = 2 * np.pi * I.to(t.float32) / probe.shape[1] + I = I.to(dtype=probe.dtype,device=probe.device) + J = J.to(dtype=probe.dtype,device=probe.device) + + for tr, sp in zip(integer_translations, + subpixel_translations): + fft_probe = fftshift(t.fft(probe, 2)) + shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J)) + shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2) + + obj_slice = obj[tr[0]:tr[0]+probe.shape[0], + tr[1]:tr[1]+probe.shape[1]] + + exit_waves.append(cmult(shifted_probe, obj_slice)) + + else: + raise NotImplementedError('Object shift not yet implemented') + + if single_translation: + return exit_waves[0] + else: + return t.stack(exit_waves) diff --git a/CDTools/tools/measurements.py b/CDTools/tools/measurements.py index 54dc2ef..d5add69 100644 --- a/CDTools/tools/measurements.py +++ b/CDTools/tools/measurements.py @@ -12,7 +12,7 @@ import numpy as np __all__ = ['intensity', 'incoherent sum', 'quadratic_background'] -def intensity(wavefield, detector_slice=None): +def intensity(wavefield, detector_slice=None, epsilon=1e-7): """Returns the intensity of a wavefield The intensity is defined as the magnitude squared of the @@ -27,15 +27,15 @@ def intensity(wavefield, detector_slice=None): torch.Tensor : A real MxN array storing the wavefield's intensities """ if detector_slice is None: - return cmath.cabssq(wavefield) + return cmath.cabssq(wavefield) + epsilon else: if wavefield.dim() == 3: - return cmath.cabssq(wavefield[detector_slice]) + return cmath.cabssq(wavefield[detector_slice]) + epsilon else: - return cmath.cabssq(wavefield[(np.s_[:],) + detector_slice]) + return cmath.cabssq(wavefield[(np.s_[:],) + detector_slice]) + epsilon -def incoherent_sum(wavefields, detector_slice=None): +def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7): """Returns the incoherent sum of the intensities of the wavefields The intensity is defined as the sum of the magnitudes squared of @@ -56,16 +56,16 @@ def incoherent_sum(wavefields, detector_slice=None): """ # This syntax just adds an axis to the slice to preserve the J direction if detector_slice is None: - return t.sum(cmath.cabssq(wavefields),dim=-3) + return t.sum(cmath.cabssq(wavefields),dim=-3) + epsilon else: if wavefields.dim() == 4: - return t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=-3) + return t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=-3) + epsilon else: - return t.sum(cmath.cabssq(wavefields[(np.s_[:],np.s_[:])+detector_slice]),dim=-3) + return t.sum(cmath.cabssq(wavefields[(np.s_[:],np.s_[:])+detector_slice]),dim=-3) + epsilon -def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity): +def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7): """Returns the intensity of a wavefield plus a background The intensity is calculated via the given measurment function @@ -83,8 +83,8 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement torch.Tensor : A real MxN array storing the wavefield's intensities """ if detector_slice is None: - return measurement(wavefield) + background**2 + return measurement(wavefield, epsilon=epsilon) + background**2 else: - return measurement(wavefield, detector_slice) \ + return measurement(wavefield, detector_slice, epsilon=epsilon) \ + background**2 diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting.py index a7b07b2..130c05d 100644 --- a/CDTools/tools/plotting.py +++ b/CDTools/tools/plotting.py @@ -49,6 +49,7 @@ def plot_1d(arr, fig = None, **kwargs): ax = fig.add_subplot(111, **kwargs) plt.scatter(np.arange(arr.shape[-1]), arr) + def plot_amplitude(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): """ Plots the amplitude of a complex Tensor or numpy array with dimensions NxMx2. Args: @@ -72,6 +73,7 @@ def plot_amplitude(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), ** plt.colorbar() return fig + def plot_phase(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): """ Plots the phase of a complex Tensor or numpy array with dimensions NxMx2. Args: @@ -97,6 +99,7 @@ def plot_phase(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwa plt.colorbar() return fig + def plot_colorized(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): """ Plots the colorized version of a complex Tensor or numpy array with dimensions NxMx2. The darkness corresponds to the intensity of the image, and the color corresponds diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index e7af484..7e65f43 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -2,33 +2,36 @@ from __future__ import division, print_function, absolute_import import CDTools from CDTools.tools import cmath +from CDTools.tools.plotting import * import h5py import torch as t import numpy as np -with h5py.File('../example_data/NiCr.cxi','r') as f: +filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi' +filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi' + + + +with h5py.File(filename,'r') as f: dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f) 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') -dataset.to(device='cuda') +#dataset.to(device='cuda') dataset.get_as(device='cuda') -for loss in model.Adam_optimize(500, dataset): +for loss in model.Adam_optimize(100, dataset): print(loss) - + from matplotlib import pyplot as plt -probe = cmath.torch_to_complex(model.probe.detach().cpu()) -obj = cmath.torch_to_complex(model.obj.detach().cpu()) - -plt.imshow(np.abs(probe)) -plt.colorbar() -plt.figure() -plt.imshow(np.abs(obj)) -plt.colorbar() +plot_amplitude(model.probe) +plot_phase(model.probe) +plot_amplitude(model.obj) +plot_phase(model.obj) plt.show() diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 01859cc..d7a21f9 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -38,3 +38,33 @@ def test_centroid_sq(): for im in np_ims] centroids = image_processing.centroid_sq(ims, comp=True) assert t.allclose(centroids, t.Tensor(np.array(sp_centroids))) + + +def test_find_pixel_shift(): + + # Test two real ims + big_im = t.rand((30,70)) + im1 = big_im[3:,:-20] + im2 = big_im[:-3,20:] + assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([-3,20])) + + # Test a real and complex im + big_im = t.rand((30,70)) + im1 = t.stack((big_im[:-5,10:],t.zeros_like(big_im[:-5,10:])),dim=-1) + im2 = big_im[5:,:-10] + assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10])) + assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10])) + + # Test two complex ims + big_im = t.rand((45,45,2)) + im1 = big_im[:-5,:-4] + im2 = big_im[5:,4:] + assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4])) + + +def test_find_subpixel_shift(): + pass + + +def test_find_shift(): + pass diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index eccfdef..fe94f50 100644 --- a/tests/tools/test_measurements.py +++ b/tests/tools/test_measurements.py @@ -9,51 +9,55 @@ import pytest def test_intensity(): wavefields = t.rand((5,10,10,2)) - np_result = np.abs(cmath.torch_to_complex(wavefields))**2 - assert t.allclose(measurements.intensity(wavefields), + epsilon=1e-6 + np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + epsilon + assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon), t.tensor(np_result)) # Test single field case - assert t.allclose(measurements.intensity(wavefields[0]), + assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon), t.tensor(np_result[0])) det_slice = np.s_[3:,5:8] - assert t.allclose(measurements.intensity(wavefields,det_slice), + assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon), t.tensor(np_result[(np.s_[:],)+det_slice])) # Test single field case - assert t.allclose(measurements.intensity(wavefields[0],det_slice), + assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon), t.tensor(np_result[0][det_slice])) def test_incoherent_sum(): wavefields = t.rand((5,4,10,10,2)) - np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=1) - assert t.allclose(measurements.incoherent_sum(wavefields), + epsilon=1e-6 + np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=1) + epsilon + assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon), t.tensor(np_result)) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[0]), + assert t.allclose(measurements.incoherent_sum(wavefields[0],epsilon=epsilon), t.tensor(np_result[0])) det_slice = np.s_[3:,5:8] - assert t.allclose(measurements.incoherent_sum(wavefields,det_slice), + assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon), t.tensor(np_result[(np.s_[:],)+det_slice])) # Test single field case - assert t.allclose(measurements.incoherent_sum(wavefields[0],det_slice), + assert t.allclose(measurements.incoherent_sum(wavefields[0],det_slice,epsilon=epsilon), t.tensor(np_result[0][det_slice])) def test_quadratic_background(): # test with intensity wavefields = t.rand((5,10,10,2)) + epsilon=1e-6 background = t.rand((10,10)) - np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2 + np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2 + epsilon det_slice = np.s_[3:,5:8] result = measurements.quadratic_background(wavefields,background[det_slice], detector_slice=det_slice, + epsilon=epsilon, measurement=measurements.intensity) assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice])) @@ -63,5 +67,6 @@ def test_quadratic_background(): np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) np_result += background.numpy()**2 result = measurements.quadratic_background(wavefields, background, + epsilon=epsilon, measurement=measurements.incoherent_sum) assert t.allclose(result, t.tensor(np_result))