Make some tweaks to the simple ptycho model

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

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)

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)

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
+    

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)

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
 

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

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()
 

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

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))