Write a multislice and s_matrix ptychography program

CDTools/models/__init__.py

@@ -52,7 +52,7 @@ from matplotlib.widgets import Slider
 from matplotlib import ticker
 import numpy as np
 
-__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho']
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho']
 
 
 class CDIModel(t.nn.Module):
@@ -349,6 +349,7 @@ class CDIModel(t.nn.Module):
                         plt.title(name)
                     except (IndexError, KeyError, AttributeError) as e:
                         pass
+
             except (IndexError, KeyError, AttributeError) as e:
                 pass
 
@@ -471,3 +472,5 @@ from CDTools.models.simple_ptycho import SimplePtycho
 from CDTools.models.fancy_ptycho import FancyPtycho
 from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
 from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
+from CDTools.models.s_matrix_ptycho import SMatrixPtycho
+from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho

CDTools/models/bragg_2d_ptycho.py

@@ -42,6 +42,18 @@ from copy import copy
 # expected numerical aperture of the probe.
 #
 
+#
+# I'm worried that the propagation doesn't happen along the correct
+# direction, if a phase ramp is expected to be baked in to the
+# retrieved focal spot. Unclear if this is the case though.
+#
+# Retrieved focal spot should have the implicit phase ramp subtracted
+# (that is, it should be the focal spot along the sample plane, but with
+# the e^ikz dependence removed). Therefore, the original e^ikz dependence
+# should be easy to re-add jusy by using the propagate_along feature
+# in ggasp. So I believe this should not be a problem
+#
+
 class Bragg2DPtycho(CDIModel):
 
     def __init__(self, wavelength, detector_geometry,
@@ -52,7 +64,7 @@ class Bragg2DPtycho(CDIModel):
                  background = None, translation_offsets=None, mask=None,
                  weights = None, translation_scale = 1, saturation=None,
                  probe_support = None, obj_support=None, oversampling=1,
-                 propagate_probe=True, correct_tilt=True):
+                 propagate_probe=True, correct_tilt=True, lens=False):
 
         
         # We need the detector geometry
@@ -148,6 +160,7 @@ class Bragg2DPtycho(CDIModel):
         else:
             self.obj_support = t.ones_like(self.obj)
 
+
         self.oversampling = oversampling
 
         self.propagate_probe = propagate_probe
@@ -163,16 +176,27 @@ class Bragg2DPtycho(CDIModel):
                     self.probe_basis, self.detector_geometry['basis'],
                     det_shape,
                     self.detector_geometry['distance'],
-                    self.wavelength,dtype=t.float32)
+                    self.wavelength,dtype=t.float32,
+                    lens=lens)
         else:
             self.k_map = None
             self.intensity_map = None
 
         self.prop_dir = t.Tensor([0,0,1]).to(dtype=t.float32)
-        
+
+        # This propagator should be able to be multiplied by the propagation
+        # distance each time to get a propagator
+        self.universal_propagator = cmath.cphase(ggasp(self.probe.shape[1:],
+                                    self.probe_basis, self.wavelength,
+                                    t.Tensor([0,0,self.wavelength/(2*np.pi)]),
+                                    propagation_vector=self.prop_dir,
+                                    dtype=t.float32,
+                                    propagate_along_offset=True))
+
+
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -201,7 +225,6 @@ class Bragg2DPtycho(CDIModel):
                                                    padding=padding,
                                                    opt_for_fft=False,
                                                    oversampling=oversampling)
-
         # now we grab the sample surface normal
         if hasattr(dataset, 'sample_info') and \
            dataset.sample_info is not None and \
@@ -332,13 +355,15 @@ class Bragg2DPtycho(CDIModel):
                    obj_support=obj_support,
                    oversampling=oversampling,
                    propagate_probe=propagate_probe,
-                   correct_tilt=correct_tilt)
+                   correct_tilt=correct_tilt,
+                   lens=lens)
                    
     
     def interaction(self, index, translations):
         pix_trans, props = tools.interactions.project_translations_to_sample(
             self.probe_basis, translations)
 
+        
         pix_trans -= self.min_translation
         props -= self.median_propagation
 
@@ -348,23 +373,31 @@ class Bragg2DPtycho(CDIModel):
         single_translation = False
         if translations.dim() == 1:
             translations = translations[None,:]
+            pix_trans = pix_trans[None,:]
             single_translation = True
             
         all_exit_waves = []
+
         for i in range(self.probe.shape[0]):
             pr = self.probe[i] * self.probe_support
 
             exit_waves = []
             for j in range(translations.size()[0]):
                 if self.propagate_probe:
-                    propagator = ggasp(pr.shape, self.probe_basis, self.wavelength,
-                                       t.Tensor([0,0,props[j]]),
-                                       propagation_vector=self.prop_dir,
-                                       dtype=pr.dtype,device=pr.device, propagate_along_offset=True)
+                    #propagator = ggasp(pr.shape, self.probe_basis, self.wavelength,
+                    #                   t.Tensor([0,0,props[j]]),
+                    #                   propagation_vector=self.prop_dir,
+                    #                   dtype=pr.dtype,device=pr.device, propagate_along_offset=True)
+                    # Minus sign is empirical
+                    propagator = cmath.expi((-props[j]*(2*np.pi)/self.wavelength)
+                        * self.universal_propagator)
                     prop_pr = tools.propagators.near_field(pr, propagator)
-
+                    #plt.close('all')
+                    #plt.imshow(np.abs(cmath.torch_to_complex(prop_pr.detach().cpu())))
+                    #plt.show()
                 else:
                     prop_pr = pr
+ 
                 exit_waves.append(self.probe_norm *
                                   tools.interactions.ptycho_2D_sinc(prop_pr,
                                             self.obj_support * self.obj,
@@ -439,6 +472,7 @@ class Bragg2DPtycho(CDIModel):
         self.obj_support = self.obj_support.to(*args,**kwargs)
         self.surface_normal = self.surface_normal.to(*args, **kwargs)
         self.prop_dir = self.prop_dir.to(*args, **kwargs)
+        self.universal_propagator = self.universal_propagator.to(*args,**kwargs)
 
 
         
@@ -487,21 +521,41 @@ class Bragg2DPtycho(CDIModel):
 
 
     # Needs to be updated to allow for plotting to an existing figure
+    # plot_list = [
+    #     ('Dominant Probe Amplitude',
+    #      lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+    #     ('Dominant Probe Phase',
+    #      lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+    #     ('Subdominant Probe Amplitude',
+    #      lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+    #      lambda self: len(self.probe) >=2),
+    #     ('Subdominant Probe Phase',
+    #      lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+    #      lambda self: len(self.probe) >=2),
+    #     ('Object Amplitude', 
+    #      lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+    #     ('Object Phase',
+    #      lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
+    #     ('Corrected Translations',
+    #      lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
+    #     ('Background',
+    #      lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    # ]
     plot_list = [
         ('Dominant Probe Amplitude',
-         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)),
         ('Dominant Probe Phase',
-         lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+         lambda self, fig: p.plot_phase(self.probe[0], fig=fig)),
         ('Subdominant Probe Amplitude',
-         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig),
          lambda self: len(self.probe) >=2),
         ('Subdominant Probe Phase',
-         lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self, fig: p.plot_phase(self.probe[1], fig=fig),
          lambda self: len(self.probe) >=2),
         ('Object Amplitude', 
-         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig)),
         ('Object Phase',
-         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
+         lambda self, fig: p.plot_phase(self.obj, fig=fig)),
         ('Corrected Translations',
          lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
         ('Background',

CDTools/models/fancy_ptycho.py

@@ -168,6 +168,7 @@ class FancyPtycho(CDIModel):
         probe_max = t.max(cmath.cabs(probe))
         probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
         probe = t.stack([probe,] + probe_stack)
+        #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
 
         obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
                               
@@ -188,6 +189,7 @@ class FancyPtycho(CDIModel):
             psr = int(probe_support_radius)
             probe_support[p_cent[0]-psr:p_cent[0]+psr,
                           p_cent[1]-psr:p_cent[1]+psr] = 1
+            probe = probe * probe_support[None,:,:]
         else:
             probe_support = None;
 
@@ -225,6 +227,8 @@ class FancyPtycho(CDIModel):
             
         all_exit_waves = []
         for i in range(self.probe.shape[0]):
+            # from storing the probe in Fourier space
+            #pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support
             pr = self.probe[i] * self.probe_support
             exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
                                                            self.obj_support * self.obj,
@@ -263,6 +267,7 @@ class FancyPtycho(CDIModel):
     
     def loss(self, sim_data, real_data, mask=None):
         return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
 
     
     def to(self, *args, **kwargs):

CDTools/models/multislice_2d_ptycho.py

@@ -0,0 +1,419 @@
+from __future__ import division, print_function, absolute_import
+
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import cmath
+from CDTools.tools import plotting as p
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+
+class Multislice2DPtycho(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_guess, dz, nz,
+                 detector_slice=None,
+                 surface_normal=np.array([0.,0.,1.]),
+                 min_translation = t.Tensor([0,0]),
+                 background = None, translation_offsets=None, mask=None,
+                 weights = None, translation_scale = 1, saturation=None,
+                 probe_support = None, obj_support=None, oversampling=1):
+        
+        super(Multislice2DPtycho,self).__init__()
+        self.wavelength = t.Tensor([wavelength])
+        self.detector_geometry = copy(detector_geometry)
+        self.dz = dz
+        self.nz = nz
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = t.Tensor(det_geo['distance'])
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = t.Tensor(det_geo['basis'])
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = t.Tensor(det_geo['corner'])
+
+        self.min_translation = t.Tensor(min_translation)
+
+        self.probe_basis = t.Tensor(probe_basis)
+        self.detector_slice = detector_slice
+        self.surface_normal = t.Tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.BoolTensor(mask)
+        
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 3:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
+        else:
+            self.probe_norm = 1 * 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))
+        
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape[:-1])
+
+                
+        self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
+
+        if weights is None:
+            self.weights = None
+        else:
+            self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
+        
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) 
+
+        self.translation_scale = translation_scale
+
+        if probe_support is not None:
+            self.probe_support = probe_support
+        else:
+            self.probe_support = t.ones_like(self.probe[0])
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support
+        else:
+            self.obj_support = t.ones_like(self.obj)
+
+        self.oversampling = oversampling
+
+        spacing = np.linalg.norm(self.probe_basis,axis=0)
+        shape = np.array(self.probe.shape[1:-1])
+
+        self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz)
+
+        
+    @classmethod
+    def from_dataset(cls, dataset, dz, nz, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True):
+        
+        wavelength = dataset.wavelength
+        det_basis = dataset.detector_geometry['basis']
+        det_shape = dataset[0][1].shape
+        distance = dataset.detector_geometry['distance']
+
+        # always do this on the cpu
+        get_as_args = dataset.get_as_args
+        dataset.get_as(device='cpu')
+        (indices, translations), patterns = dataset[:]
+        dataset.get_as(*get_as_args[0],**get_as_args[1])
+
+        # Set to none to avoid issues with things outside the detector
+        if auto_center:
+            center = tools.image_processing.centroid(t.sum(patterns,dim=0))
+        else:
+            center = None
+            
+        # Then, generate the probe geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center,
+                                                   padding=padding,
+                                                   opt_for_fft=False,
+                                                   oversampling=oversampling)
+
+
+        if hasattr(dataset, 'sample_info') and \
+           dataset.sample_info is not None and \
+           'orientation' in dataset.sample_info:
+            surface_normal = dataset.sample_info['orientation'][2]
+        else:
+            surface_normal = np.array([0.,0.,1.])
+
+
+        # If this information is supplied when the function is called,
+        # then we override the information in the .cxi file
+        if scattering_mode in {'t', 'transmission'}:
+            surface_normal = np.array([0.,0.,1.])
+        elif scattering_mode in {'r', 'reflection'}:
+            outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
+            outgoing_dir /= np.linalg.norm(outgoing_dir)
+            surface_normal = outgoing_dir + np.array([0.,0.,1.])
+            surface_normal /= np.linalg.norm(outgoing_dir)
+
+
+        # Next generate the object geometry from the probe geometry and
+        # the translations
+        pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
+
+        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
+
+        if hasattr(dataset, 'background') and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        else:
+            background = None
+
+        # Finally, initialize the probe and  object using this information
+        if probe_size is None:
+            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        else:
+            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(cmath.cabs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
+        probe = t.stack([probe,] + probe_stack)
+
+        obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
+                              
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
+
+        weights = t.ones(len(dataset))
+        
+        if hasattr(dataset, 'mask') and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+        else:
+            mask = None
+
+        if probe_support_radius is not None:
+            probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
+            p_cent = np.array(probe.shape[1:3]).astype(int) // 2
+            psr = int(probe_support_radius)
+            probe_support[p_cent[0]-psr:p_cent[0]+psr,
+                          p_cent[1]-psr:p_cent[1]+psr] = 1
+            probe = probe * probe_support[None,:,:]
+        else:
+            probe_support = None;
+
+        if restrict_obj != -1:
+            ro = restrict_obj
+            os = np.array(obj_size)
+            ps = np.array(probe_shape)
+            obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
+                        ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
+        else:
+            obj_support = None
+
+        return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets = translation_offsets,
+                   weights=weights, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   obj_support=obj_support,
+                   oversampling=oversampling)
+                   
+    
+    def interaction(self, index, translations):
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations,
+                                                             surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+            
+        all_exit_waves = []
+        for i in range(self.probe.shape[0]):
+            pr = self.probe[i] * self.probe_support
+            #exit_waves =  pr
+            #print(self.probe_norm)
+            #for i in range(self.nz):
+            exit_waves = []
+            if len(pix_trans.shape) == 1:
+                pix_trans = [pix_trans]
+            for trans in pix_trans:
+                exit_wave = self.probe_norm * pr
+                for i in range(self.nz-1):
+                    
+                    #exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
+                    #                        self.obj_support * self.obj,
+                    #                        trans,
+                    #                        shift_probe=True)
+                    exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
+                                            self.obj_support * cmath.cexpi(self.obj.data/self.nz),
+                                            trans)
+                    exit_wave = tools.propagators.near_field(exit_wave,self.as_prop)
+                    #tools.plotting.plot_amplitude(exit_wave)
+                    #plt.show()
+                    
+                # only final layer gets a derivative
+                exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
+                                        self.obj_support * cmath.cexpi(0.1*self.obj),
+                                        trans)
+                exit_waves.append(exit_wave)
+                
+            exit_waves = t.stack(exit_waves)
+            
+
+            if np.array(index).size == 1:
+                index = [index]
+                strip_first_index = True
+            else:
+                strip_first_index = False
+            
+            if exit_waves.dim() == 4:
+                exit_waves =  self.weights[index][:,None,None,None] * exit_waves
+            else:
+                exit_waves =  self.weights[index] * exit_waves
+
+            if strip_first_index:
+                exit_waves = exit_waves[0,...]
+                
+            all_exit_waves.append(exit_waves)
+
+        
+        return t.stack(all_exit_waves)
+    
+        
+    def forward_propagator(self, wavefields):
+        return tools.propagators.far_field(wavefields)
+
+
+    def backward_propagator(self, wavefields):
+        return tools.propagators.inverse_far_field(wavefields)
+
+    
+    def measurement(self, wavefields):
+        return tools.measurements.quadratic_background(wavefields,
+                            self.background,
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+
+    
+    def loss(self, sim_data, real_data, mask=None):
+        regularizer = t.mean(t.abs(self.obj))
+        loss = tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
+        lambd = 1
+        return loss + lambd * regularizer
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
+
+    
+    def to(self, *args, **kwargs):
+        super(Multislice2DPtycho, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args,**kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
+        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)
+        self.probe_support = self.probe_support.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.surface_normal = self.surface_normal.to(*args, **kwargs)
+        self.as_prop = self.as_prop.to(*args, **kwargs)
+        
+    def sim_to_dataset(self, args_list):
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        surface_normal = self.surface_normal.detach().cpu().numpy()
+        xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
+        xsurfacevec /= np.linalg.norm(xsurfacevec)
+        ysurfacevec = np.cross(surface_normal, xsurfacevec)
+        ysurfacevec /= np.linalg.norm(ysurfacevec)
+        orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
+        
+        sample_info = {'description': 'A simulated sample',
+                       'orientation': orientation}
+
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho2DDataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
+
+    
+    def corrected_translations(self,dataset):
+        translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
+        t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
+        return translations + t_offset
+
+
+    # Needs to be updated to allow for plotting to an existing figure
+    plot_list = [
+        ('Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Subdominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Subdominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Object Amplitude', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = cmath.torch_to_complex(self.obj.detach().cpu())
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+        dz = self.dz
+        nz = self.nz
+        prop = self.as_prop
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'obj':obj,
+                'background':background,
+                'weights':weights, 'dz':dz, 'nz':nz,
+                'interlayer propagator': prop}

CDTools/models/s_matrix_ptycho.py

@@ -0,0 +1,365 @@
+from __future__ import division, print_function, absolute_import
+
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import cmath
+from CDTools.tools import plotting as p
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+
+class SMatrixPtycho(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis, probe_guess, probe_fourier_support,
+                 s_matrix_guess,
+                 detector_slice=None,
+                 surface_normal=np.array([0.,0.,1.]),
+                 min_translation = t.Tensor([0,0]),
+                 background = None, translation_offsets=None, mask=None,
+                 weights = None, translation_scale = 1, saturation=None,
+                 oversampling=1):
+        
+        super(SMatrixPtycho,self).__init__()
+        self.wavelength = t.Tensor([wavelength])
+        self.detector_geometry = copy(detector_geometry)
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = t.Tensor(det_geo['distance'])
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = t.Tensor(det_geo['basis'])
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = t.Tensor(det_geo['corner'])
+
+        self.min_translation = t.Tensor(min_translation)
+
+        self.probe_basis = t.Tensor(probe_basis)
+        self.detector_slice = detector_slice
+        self.surface_normal = t.Tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.BoolTensor(mask)
+        
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 3:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
+        else:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))         
+            
+        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+                                    / self.probe_norm)
+        
+        self.s_matrix = t.nn.Parameter(s_matrix_guess.to(t.float32))
+        
+        if background is None:
+            ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[1],
+                        s_matrix_guess.shape[1] - 1 + probe_guess.shape[2]]
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(t.ones(ew_shape)[self.detector_slice].shape).to(t.float32)
+            else:
+                background = 1e-6 * t.ones(ew_shape).to(t.float32)
+                
+        self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
+
+        if weights is None:
+            self.weights = None
+        else:
+            self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
+        
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) 
+
+        self.translation_scale = translation_scale
+
+        self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
+        
+        self.oversampling = oversampling
+        
+        
+    @classmethod
+    def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True):
+        
+        wavelength = dataset.wavelength
+        det_basis = dataset.detector_geometry['basis']
+        det_shape = dataset[0][1].shape
+        distance = dataset.detector_geometry['distance']
+
+        # always do this on the cpu
+        get_as_args = dataset.get_as_args
+        dataset.get_as(device='cpu')
+        (indices, translations), patterns = dataset[:]
+        dataset.get_as(*get_as_args[0],**get_as_args[1])
+
+        # Set to none to avoid issues with things outside the detector
+        if auto_center:
+            center = tools.image_processing.centroid(t.sum(patterns,dim=0))
+        else:
+            center = None
+            
+        # Then, generate the probe geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, ew_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center,
+                                                   padding=padding,
+                                                   opt_for_fft=False,
+                                                   oversampling=oversampling)
+
+        # This shrinks the probe to ensure that the output wavefield
+        # is the correct shape
+        probe_shape = t.Size(np.array(ew_shape) - (2*locality_radius))
+        
+        if hasattr(dataset, 'sample_info') and \
+           dataset.sample_info is not None and \
+           'orientation' in dataset.sample_info:
+            surface_normal = dataset.sample_info['orientation'][2]
+        else:
+            surface_normal = np.array([0.,0.,1.])
+
+
+        # If this information is supplied when the function is called,
+        # then we override the information in the .cxi file
+        if scattering_mode in {'t', 'transmission'}:
+            surface_normal = np.array([0.,0.,1.])
+        elif scattering_mode in {'r', 'reflection'}:
+            outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
+            outgoing_dir /= np.linalg.norm(outgoing_dir)
+            surface_normal = outgoing_dir + np.array([0.,0.,1.])
+            surface_normal /= np.linalg.norm(outgoing_dir)
+
+
+        # Next generate the object geometry from the probe geometry and
+        # the translations
+        pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
+
+        # The locality radius correction is probably not needed because
+        # it will always be way less than 200, but it ensures that there
+        # is no wrapping in the s-matrix
+        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200+2*locality_radius)
+
+        if hasattr(dataset, 'background') and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        else:
+            background = None
+
+        # Finally, initialize the probe and  object using this information
+        if probe_size is None:
+            if locality_radius != 0:
+                probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius]
+            else:
+                probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        else:
+            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(cmath.cabs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
+        probe = t.stack([tools.propagators.inverse_far_field(probe),] + probe_stack)
+
+        s_matrix = t.zeros([2*locality_radius+1,2*locality_radius+1,obj_size[0],
+                            obj_size[1],2])
+        s_matrix[locality_radius,locality_radius,:,:,:] = \
+            tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
+
+
+        
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
+
+        weights = t.ones(len(dataset))
+        
+        if hasattr(dataset, 'mask') and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+        else:
+            mask = None
+
+        
+        probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
+        xs, ys = np.mgrid[:probe.shape[1],:probe.shape[2]]
+        xs = xs - np.mean(xs)
+        ys = ys - np.mean(ys)
+        Rs = np.sqrt(xs**2 + ys**2)
+        probe_support[Rs<probe_convergence_radius] = 1
+        probe = probe * probe_support[None,:,:]
+        
+
+        return cls(wavelength, det_geo, probe_basis, probe, probe_support,
+                   s_matrix,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets = translation_offsets,
+                   weights=weights, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   oversampling=oversampling)
+                   
+    
+    def interaction(self, index, translations):
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations,
+                                                             surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+            
+        all_exit_waves = []
+        for i in range(self.probe.shape[0]):
+            pr = tools.propagators.inverse_far_field(self.probe[i] * self.probe_fourier_support)
+            exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(
+                pr, self.s_matrix, pix_trans, shift_probe=True)
+            
+            exit_waves = exit_waves 
+
+
+            if exit_waves.dim() == 4:
+                exit_waves =  self.weights[index][:,None,None,None] * exit_waves
+            else:
+                exit_waves =  self.weights[index] * exit_waves
+                
+            all_exit_waves.append(exit_waves)
+
+        
+        return t.stack(all_exit_waves)
+    
+        
+    def forward_propagator(self, wavefields):
+        return tools.propagators.far_field(wavefields)
+
+
+    def backward_propagator(self, wavefields):
+        return tools.propagators.inverse_far_field(wavefields)
+
+    
+    def measurement(self, wavefields):
+        return tools.measurements.quadratic_background(wavefields,
+                            self.background,
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+
+    
+    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):
+        super(SMatrixPtycho, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args,**kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
+        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)
+        self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs)
+        self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+        
+    def sim_to_dataset(self, args_list):
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        surface_normal = self.surface_normal.detach().cpu().numpy()
+        xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
+        xsurfacevec /= np.linalg.norm(xsurfacevec)
+        ysurfacevec = np.cross(surface_normal, xsurfacevec)
+        ysurfacevec /= np.linalg.norm(ysurfacevec)
+        orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
+        
+        sample_info = {'description': 'A simulated sample',
+                       'orientation': orientation}
+
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho2DDataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
+
+    
+    def corrected_translations(self,dataset):
+        translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
+        t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
+        return translations + t_offset
+
+
+    # Needs to be updated to allow for plotting to an existing figure
+    plot_list = [
+        ('Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Subdominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Subdominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Exit Wave Amplitude under Uniform Illumination', 
+         lambda self, fig: p.plot_amplitude(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
+        ('Exit Wave Phase under Uniform Illumination',
+         lambda self, fig: p.plot_phase(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        s_matrix = cmath.torch_to_complex(self.s_matrix.detach().cpu())
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'s_matrix':s_matrix,
+                'background':background,
+                'weights':weights}

CDTools/tools/cmath.py

@@ -304,3 +304,24 @@ def expi(x):
 
     """
     return t.stack((t.cos(x),t.sin(x)),dim=-1)
+
+
+def cexpi(z):
+    """Returns a complex-format tensor for exp(i* (z))
+    
+    Expects the input to be in the form of a complex-valued tensor
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        An array to be exponentiated
+    
+    Returns
+    -------
+    torch.Tensor
+        A complex-format tensor
+
+    """
+    real = t.cos(z[...,0]) * t.exp(-z[...,1])
+    imag = t.sin(z[...,0]) * t.exp(-z[...,1])
+    return t.stack((real, imag),dim=-1)

CDTools/tools/interactions.py

@@ -423,4 +423,97 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
         return t.stack(exit_waves)
 
 
+def ptycho_2D_sinc_s_matrix(probe, s_matrix, 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
+
+    This needs to be edited to use a slightly different meaning of the s-wave
+    format. Currently, each pixel in the latter two dimensions index a location
+    on the input wavefield, and the first two indexes index differences from
+    that pixel. It is easier to interpret the resulting matrix though if the
+    latter two indices index locations in the output plane.
+
+    Parameters
+    ----------
+    probe : torch.Tensor
+        An MxL probe function for the exit waves
+    s_matrix : torch.Tensor
+        The 4D S-Matrix tensor (2B+1x2B+1xObject Shape) to be probed
+    translations : torch.Tensor
+        The Nx2 array of translations to simulate
+    shift_probe : bool
+        Default True, 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
+    -------
+    exit_waves : 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 = []
+
+    B = s_matrix.shape[0]//2
+    
+    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 * J.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)
+
+            s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0],
+                                      tr[1]:tr[1]+probe.shape[1]]
+
+
+            output = t.zeros([s_matrix_slice.shape[2]+2*B,
+                              s_matrix_slice.shape[3]+2*B,2]).to(
+                                  device=s_matrix_slice.device,
+                                  dtype=s_matrix_slice.dtype)
+            
+            for i in range(s_matrix.shape[0]):
+                for j in range(s_matrix.shape[1]):
+                    output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
+                        cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
+
+            exit_waves.append(output)
+            #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/losses.py

@@ -100,7 +100,7 @@ def intensity_mse(intensities, sim_intensities, mask=None):
 
 
     
-def poisson_nll(intensities, sim_intensities, mask=None):
+def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-4):
     """ Returns the Poisson negative log likelihood for a simulated dataset's intensities
 
     Calculates the overall Poisson maximum likelihood metric using
@@ -133,12 +133,12 @@ def poisson_nll(intensities, sim_intensities, mask=None):
 
     """
     if mask is None:
-        return t.sum(sim_intensities -
-                     intensities * t.log(sim_intensities)) \
+        return t.sum(sim_intensities+eps -
+                     intensities * t.log(sim_intensities+eps)) \
                      / intensities.view(-1).shape[0]
     
     else:
         masked_intensities = intensities.masked_select(mask)
         masked_sims = sim_intensities.masked_select(mask)
         return t.sum(masked_sims - masked_intensities *
-                     t.log(masked_sims)) / masked_intensities.shape[0]
+                     t.log(masked_sims+eps)) / masked_intensities.shape[0]

CDTools/tools/plotting.py

@@ -127,6 +127,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
     if basis is not None:
         if isinstance(basis,t.Tensor):
             basis = basis.detach().cpu().numpy()
+        # This fails if the 
         basis_norm = np.linalg.norm(basis, axis = 0)
         basis_norm = basis_norm * get_units_factor(units)
 
@@ -207,7 +208,7 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
         except:
             plt.imshow(phase, cmap = 'hsv', extent=extent)
     else:
-        plt.imshow(phase)#, cmap = cmap, extent=extent)
+        plt.imshow(phase, cmap = cmap, extent=extent)
 
     cbar = plt.colorbar()
     cbar.set_label('Phase (rad)')
@@ -222,6 +223,9 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
     return fig
 
 
+def plot_amplitude_surfacenorm():
+    pass
+
 def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
     """ Plots the colorized version of a complex array with dimensions NxM
 

CDTools/tools/propagators.py

@@ -74,7 +74,7 @@ def inverse_far_field(wavefront):
     return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
 
 
-def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, **kwargs):
+def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, lens=False, **kwargs):
     """Generates k-space and intensity maps to allow for high-NA far-field propagation of light
 
     At high numerical apertures or for very tilted samples, the simple
@@ -97,6 +97,16 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     The intensity map is simply an object, the shape of the detector, which
     encodes intensity corrections between 0 and 1 per pixel.
 
+    If the optional "lens" parameter is set to True, the intensity map will
+    be set to a uniform map, and the distortion of Fourier space due to the
+    flat nature of the detector (that is, the portion of the distortion
+    that exists even if the sample is not tilted) will be disabled. This is
+    to account for the fact that a good, infinity-conjugate imaging lens
+    will do it's best to correct for these abberations in the lens. Of course,
+    the lens will not be perfect, but in such a case it is a better
+    approximation to assume that the lens is perfect than to assume that it
+    is not there at all.
+
     Parameters
     ----------
     sample_basis: array
@@ -109,7 +119,8 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
         The sample-to-detector distance
     wavelength: float
         The wavelength of light being propagated
-
+    lens: bool
+        Whether the diffraction pattern is formed by a lens or not.
 
     Returns
     -------
@@ -153,15 +164,39 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     samp_det_vec = np.cross(det_basis[:,0],det_basis[:,1])
     samp_det_vec *= distance / np.linalg.norm(samp_det_vec)
 
-    Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \
-        + samp_det_vec[:,None,None]
+    # This could potentially correct for a mistake in the implied
+    # propagation direction (e.g. choosing e^ikx instead of e^-ikx)
+    #samp_det_vec *= -1
 
+    if lens == False:
+        # This correctly reproduces the sample-to-each-pixel vectors
+        # in the case where the diffraction pattern is actually formed
+        # by Fraunhoffer diffraction
+        Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \
+            + samp_det_vec[:,None,None]
+    else:
+        # This forms a distorted set of vectors designed to produce the
+        # correct Fourier space map in the case where an imaging lens is
+        # used in the 2f geometry. One should not read too much meaning
+        # into these vectors, they are simply set up to produce the
+        # correct final K-map
+        Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1)#
+        Rs += (samp_det_vec / np.linalg.norm(samp_det_vec))[:,None,None] * \
+            np.sqrt(np.sum((samp_det_vec)**2)-np.sum(Rs**2,axis=0))[None,:,:]
+        
     k0 = 2*np.pi/wavelength
     
     Ks = k0 * Rs / np.linalg.norm(Rs, axis=0)
 
+    # My attempt at seeing what happens if I flip the Ks
+    #Ks *= -1
+    
     # This is the cosine of the angle with the detector normal
     intensity_map = np.tensordot(samp_det_vec/(k0*distance),Ks,axes=1)
+    if lens:
+        # Set the intensity map to be uniform if a lens is being used
+        intensity_map = np.ones_like(intensity_map)
+    
     intensity_map = t.Tensor(intensity_map).to(*args, **kwargs)
 
 
@@ -171,6 +206,10 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     # uniform phase.
     Ks -= k0 * samp_det_vec[:,None,None] / distance
 
+    # A potential alternative when Ks are flipped
+    #Ks += k0 * samp_det_vec[:,None,None] / distance
+
+    
     # Now we move on to finding the conversion into k-space
     # for the sample grid. It turns out we can do this by multiplying
     # them with the real space basis (dual of the reciprocal space
@@ -233,9 +272,13 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
     low_NA_wavefield = far_field(wavefront)
     # I'm going to need to separately interpolate the real and complex parts
     # This can be done
-
+    
     k_map = k_map[None,:,:,:]
     # Will only work for a 4D wavefile stack.
+    #plt.figure()
+    #plt.pcolormesh(k_map[0,:,:,0].cpu().numpy(),k_map[0,:,:,1].cpu().numpy(),
+    #               np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy()))
+    #plt.show()
     def process_wavefield_stack(low_NA_wavefield):
         real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
         imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)