move to a src directory layout

src/CDTools/__init__.py

@@ -0,0 +1,12 @@
+# This is needed to allow us to use torch.tensor in the module without it
+# constantly complaining.
+import warnings
+warnings.filterwarnings("ignore",
+                        message='To copy construct from a tensor, ')
+
+__all__ = ['tools', 'datasets', 'models']
+
+from CDTools import tools
+from CDTools import datasets
+from CDTools import models
+

src/CDTools/datasets/__init__.py

@@ -0,0 +1,34 @@
+""" This module contains all the datasets for interacting with ptychography data
+
+All the access to data from standard ptychography and CDI experiments is 
+coordinated through the various datasets defined in this module. They make use
+of the lower-level data reading and writing functions defined in tools.data,
+but critically all of these datasets subclass torch.Dataset. This allows
+them to be used as standard torch datasets during reconstructions, which
+helps make it easy to use the various data-handling strategies that are
+implemented by default in pytorch (such as drawing data in a random order,
+drawing minibatches, etc.)
+
+New Datasets can be defined a subclass of the main CDataset class defined
+in the base.py file. Example implementations of all these functions
+can be found in the code for the Ptycho2DDataset class. In addition, it is
+recommended to read through the tutorial section on defining a new CDI
+dataset before attempting to do so
+
+* __init__
+* __len__
+* _load
+* to
+* from_cxi
+* to_cxi
+* inspect
+
+"""
+
+# I don't believe that __all__ really needed, but it's nice to define it
+# to be explicit that import * is safe
+__all__ = ['CDataset','Ptycho2DDataset','PolarizedPtycho2DDataset']
+
+from CDTools.datasets.base import CDataset
+from CDTools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+from CDTools.datasets.polarized_ptycho_2d_dataset import PolarizedPtycho2DDataset

src/CDTools/datasets/base.py

@@ -0,0 +1,252 @@
+""" This module contains the base CDataset class for handling CDI data
+
+Subclasses of CDataset are required to define their own implementations
+of the following functions:
+
+* __init__
+* __len__
+* _load
+* to
+* from_cxi
+* to_cxi
+* inspect
+
+"""
+
+import torch as t
+from copy import copy
+import h5py
+import pathlib
+from CDTools.tools import data as cdtdata
+from torch.utils import data as torchdata
+
+__all__ = ['CDataset']
+
+
+class CDataset(torchdata.Dataset):
+    """ The base dataset class which all other datasets subclass
+
+    Subclasses torch.utils.data.Dataset
+
+    This base dataset class defines the functionality which should be
+    common to all subclassed datasets. This includes the loading and
+    storage of the metadata portions of .cxi files, as well as the tools
+    needed to allow for easy mixing of data on the CPU and GPU.
+    """
+
+    def __init__(self, entry_info=None, sample_info=None,
+                 wavelength=None,
+                 detector_geometry=None, mask=None,
+                 background=None):
+
+        """The __init__ function allows construction from python objects.
+
+        The detector_geometry dictionary is defined to have the
+        entries defined by the outputs of data.get_detector_geometry.
+
+
+        Parameters
+        ----------
+        entry_info : dict
+            A dictionary containing the entry_info metadata
+        sample_info : dict
+            A dictionary containing the sample_info metadata
+        wavelength : float
+            The wavelength of light used in the experiment
+        detector_geometry : dict
+            A dictionary containing the various detector geometry
+            parameters
+        mask : array
+            A mask for the detector, defined as 1 for live pixels, 0
+            for dead
+        background : array
+            An initial guess for the not-previously-subtracted
+            detector background
+        """
+
+        # Force pass-by-value-like behavior to stop strangeness
+        self.entry_info = copy(entry_info)
+        self.sample_info = copy(sample_info)
+        self.wavelength = wavelength
+        self.detector_geometry = copy(detector_geometry)
+        if mask is not None:
+            self.mask = t.tensor(mask, dtype=t.bool)
+        else:
+            self.mask = None
+        if background is not None:
+            self.background = t.tensor(background, dtype=t.float32)
+        else:
+            self.background = None
+
+        self.get_as(device='cpu')
+
+
+    def to(self, *args, **kwargs):
+        """Sends the relevant data to the given device and dtype
+
+        This function sends the stored mask and background to the
+        specified device and dtype
+
+        Accepts the same parameters as torch.Tensor.to
+        """
+        # The mask should always stay a uint8, but it should switch devices
+        mask_kwargs = copy(kwargs)
+        try:
+            mask_kwargs.pop('dtype')
+        except KeyError:
+            pass
+
+        if self.mask is not None:
+            self.mask = self.mask.to(*args,**mask_kwargs)
+        if self.background is not None:
+            self.background = self.background.to(*args,**kwargs)
+
+
+    def get_as(self, *args, **kwargs):
+        """Sets the dataset to return data on the given device and dtype
+
+        Oftentimes there isn't room to store an entire dataset on a GPU,
+        but it is still worth running the calculation on the GPU even with
+        the overhead incurred by transferring data back and forth. In that
+        case, get_as can be used instead of to, to declare a set of
+        device and dtype that the data should be returned as, whenever it
+        is accessed through the __getitem__ function (as it would be in
+        any reconstructions).
+
+        Parameters
+        ----------
+        Accepts the same parameters as torch.Tensor.to
+        """
+        self.get_as_args = (args, kwargs)
+
+    def __len__(self):
+        raise NotImplementedError()
+    
+    def __getitem__(self, index):
+        # Deals with loading to appropriate device/dtype, if
+        # specified via a call to get_as
+        inputs, outputs = self._load(index)
+        if hasattr(self, 'get_as_args'):
+            outputs = outputs.to(*self.get_as_args[0],**self.get_as_args[1])
+            moved_inputs = []
+            for inp in inputs:
+                try:
+                    moved_inputs.append(inp.to(*self.get_as_args[0],**self.get_as_args[1]) )
+                except:
+                    moved_inputs.append(inp)
+        else:
+            moved_inputs = inputs
+        return moved_inputs, outputs
+
+    
+    def _load(self, index):
+        """ Internal function to load data
+        
+        In all subclasses of CDataset, a _load function should be defined.
+        This function is used internally by the global __getitem__ function
+        defined in the base class, which handles moving data around when
+        the dataset is (for example) storing the data on the CPU but 
+        getting data as GPU tensors.
+
+        It should accept an index or slice, and return output as a tuple.
+        The first item of the tuple is a tuple containing the inputs to
+        the forward model for the related ptychography model. The second
+        item of the tuple should be the set of diffraction patterns
+        associated with the returned inputs.
+
+        Since there is no kind of data stored in a CDataset, this
+        function is defined as returing a NotImplemented Error
+        """
+        raise NotImplementedError()
+        
+            
+    @classmethod
+    def from_cxi(cls, cxi_file):
+        """Generates a new CDataset from a .cxi file directly
+
+        This is the most commonly used constructor for CDatasets and
+        subclasses thereof. It populates the dataset using the information
+        in a .cxi file. It can either take an h5py.File object directly,
+        or a filename or pathlib object pointing to the file
+
+        Parameters
+        ----------
+        file : str, pathlib.Path, or h5py.File
+            The .cxi file to load from
+
+        Returns
+        -------
+        dataset : CDataset
+            The constructed dataset object
+        """
+        
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with h5py.File(cxi_file,'r') as f:
+                return cls.from_cxi(f)
+        
+        entry_info = cdtdata.get_entry_info(cxi_file)
+        sample_info = cdtdata.get_sample_info(cxi_file)
+        wavelength = cdtdata.get_wavelength(cxi_file)
+        distance, basis, corner = cdtdata.get_detector_geometry(cxi_file)
+        detector_geometry = {'distance' : distance,
+                             'basis'    : basis,
+                             'corner'   : corner}
+        mask = cdtdata.get_mask(cxi_file)
+        dark = cdtdata.get_dark(cxi_file)
+        return cls(entry_info = entry_info,
+                   sample_info = sample_info,
+                   wavelength=wavelength,
+                   detector_geometry=detector_geometry,
+                   mask=mask, background=dark)
+    
+    
+    def to_cxi(self, cxi_file):
+        """Saves out a CDataset as a .cxi file 
+
+        This function saves all the compatible information in a CDataset
+        object into a .cxi file. This is useful for saving out modified
+        or simulated datasets
+        
+        Parameters
+        ----------
+        cxi_file : str, pathlib.Path, or h5py.File
+            The .cxi file to write to
+        """
+
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with h5py.File(cxi_file,'w') as f:
+                return self.to_cxi(f)
+
+        if self.entry_info is not None:
+            cdtdata.add_entry_info(cxi_file, self.entry_info)
+        if self.sample_info is not None:
+            cdtdata.add_sample_info(cxi_file, self.sample_info)
+        if self.wavelength is not None:
+            cdtdata.add_source(cxi_file, self.wavelength)
+        if self.detector_geometry is not None:
+            if 'corner' in self.detector_geometry:
+                corner = self.detector_geometry['corner']
+            else:
+                corner = None
+            cdtdata.add_detector(cxi_file,
+                               self.detector_geometry['distance'],
+                               self.detector_geometry['basis'],
+                               corner = corner)
+        if self.mask is not None:
+            cdtdata.add_mask(cxi_file, self.mask)
+        if self.background is not None:
+            cdtdata.add_dark(cxi_file, self.background)
+        
+    def inspect(self):
+        """The prototype for the inspect function
+
+        In all subclasses of CDataset, an inspect function should be
+        defined which opens a tool that shows the data in a natural
+        layout for that kind of experiment. In the base class, no actual
+        data is stored, so this is defined to raise a NotImplementedError
+        """
+        raise NotImplementedError
+
+

src/CDTools/datasets/polarized_ptycho_2d_dataset.py

@@ -0,0 +1,218 @@
+import numpy as np
+import torch as t
+import h5py
+import pathlib
+from CDTools.datasets import Ptycho2DDataset
+from CDTools.tools import data as cdtdata
+from CDTools.tools import plotting
+
+__all__ = ['PolarizedPtycho2DDataset']
+
+
+class PolarizedPtycho2DDataset(Ptycho2DDataset):
+    """The standard dataset for a 2D ptychography scan
+
+    Subclasses datasets.Ptycho2DDataset
+
+    This class loads and saves 2D ptychography scan data from .cxi files.
+    It should save and load files compatible with most reconstruction
+    programs, although it is only tested against SHARP.
+    """
+
+    def __init__(self, translations, polarizer_angles, analyzer_angles, patterns, axes=None, *args, **kwargs):
+        """The __init__ function allows construction from python objects.
+
+        The detector_geometry dictionary is defined to have the
+        entries defined by the outputs of data.get_detector_geometry.
+
+
+        Parameters
+        ----------
+        translations : array (nx3 t.tensor when polarized=True)
+            An nx3 array containing the probe translations at each scan point
+        patterns : array
+            An nxmxl array containing the full stack of measured diffraction patterns
+        axes : list(str)
+            A
+             of names for the axes of the probe translations
+        entry_info : dict
+            A dictionary containing the entry_info metadata
+        sample_info : dict
+            A dictionary containing the sample_info metadata
+        wavelength : float
+            The wavelength of light used in the experiment
+        detector_geometry : dict
+            A dictionary containing the various detector geometry
+            parameters
+        mask : array
+            A mask for the detector, defined as 1 for live pixels, 0
+            for dead
+        background : array
+            An initial guess for the not-previously-subtracted
+            detector background
+        """
+        
+        super(PolarizedPtycho2DDataset,self).__init__(translations, patterns,
+                                             *args, **kwargs)
+        
+        
+        self.polarizer = t.tensor(polarizer_angles, dtype=t.float32)
+        self.analyzer = t.tensor(analyzer_angles, dtype=t.float32)
+        
+        
+    def _load(self, index):
+        """ Internal function to load data
+
+        This function is used internally by the global __getitem__ function
+        defined in the base class, which handles moving data around when
+        the dataset is (for example) storing the data on the CPU but
+        getting data as GPU tensors.
+
+        It loads data in the format (inputs, output)
+
+        The inputs for a 2D ptychogaphy data set are:
+
+        1) The indices of the patterns to use
+        2) The recorded probe positions associated with those points
+        3) The angles of the polarizers if polarized=True
+
+        Parameters
+        ----------
+        index (polarized=False): int or slice
+             The index or indices of the scan points to use
+
+        index (polarized=True): 
+            tuple ((phi1, phi2), ind) 
+
+            ind - index or indices of the scan points to use (in a (phi1, phi2) polarization state), int or slice
+            (phi1, phi2) - angles of the 1st and 2nd polarizers, ints
+
+        Returns
+        -------
+        inputs : tuple
+             A tuple of the inputs to the related forward models
+             if polarized: inputs = ((phi1, phi2), ind, transl[(phi1, phi2)][ind])
+        outputs : tuple
+             The output pattern or stack of output patterns
+        """
+        return ((index, self.translations[index],
+                 self.polarizer[index], self.analyzer[index]),
+                self.patterns[index])
+
+    def to(self, *args, **kwargs):
+        """Sends the relevant data to the given device and dtype
+
+        This function sends the stored translations, patterns,
+        mask and background to the specified device and dtype
+
+        Accepts the same parameters as torch.Tensor.to
+        """
+        super(PolarizedPtycho2DDataset, self).to(*args, **kwargs)
+        self.polarizer = self.polarizer.to(*args, **kwargs)
+        self.analyzer = self.analyzer.to(*args, **kwargs)
+
+
+    # It sucks that I can't reuse the base factory method here,
+    # perhaps there is a way but I couldn't figure it out.
+    @classmethod
+    def from_cxi(cls, cxi_file, cut_zeros=True):
+        """Generates a new PolarizedPtycho2DDataset from a .cxi file directly
+
+        This generates a new PolarizedPtycho2DDataset from a .cxi file storing
+        a 2D ptychography scan.
+
+        Parameters
+        ----------
+        file : str, pathlib.Path, or h5py.File
+            The .cxi file to load from
+        cut_zeros : bool
+            Default True, whether to set all negative data to zero
+
+        Returns
+        -------
+        dataset : PolarizedPtycho2DDataset
+            The constructed dataset object
+        """
+
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with h5py.File(cxi_file, 'r') as f:
+                return cls.from_cxi(f, cut_zeros=cut_zeros)
+
+        # Generate a base dataset
+        dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros)
+
+        # Mutate the class to this subclass (PolarizedPtycho2DDataset)
+        dataset.__class__ = cls
+
+        # Now, we save out the polarizer and analyzer states
+        polarizer = cdtdata.get_shot_to_shot_info(cxi_file, 'polarizer_angle')
+        analyzer = cdtdata.get_shot_to_shot_info(cxi_file, 'analyzer_angle')
+
+        dataset.analyzer = t.tensor(analyzer, dtype=t.float32)
+        dataset.polarizer = t.tensor(polarizer, dtype=t.float32)
+        
+        return dataset
+
+    def to_cxi(self, cxi_file, polarized=False):
+        """Saves out a PolarizedPtycho2DDataset as a .cxi file
+
+        This function saves all the compatible information in a
+        PolarizedPtycho2DDataset object into a .cxi file. This saved .cxi file
+        should be compatible with any standard .cxi file based
+        reconstruction tool, such as SHARP.
+
+        Parameters
+        ----------
+        cxi_file : str, pathlib.Path, or h5py.File
+            The .cxi file to write to
+        """
+
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with cdtdata.create_cxi(cxi_file) as f:
+                return self.to_cxi(f)
+
+        # This saves the translations, patterns, etc.
+        super(PolarizedPtycho2DDataset, self).to_cxi(cxi_file)
+
+        # Now, we save out the polarizer and analyzer states
+        cdtdata.add_shot_to_shot_info(cxi_file, self.polarizer,
+                                      'polarizer_angle')
+        cdtdata.add_shot_to_shot_info(cxi_file, self.analyzer,
+                                      'analyzer_angle')
+
+    def inspect(self, logarithmic=True, units='um'):
+        """Launches an interactive plot for perusing the data
+
+        This launches an interactive plotting tool in matplotlib that
+        shows the spatial map constructed from the integrated intensity
+        at each position on the left, next to a panel on the right that
+        can display a base-10 log plot of the detector readout at each
+        position.
+        """
+
+
+        def get_images(idx):
+            inputs, output = self[idx]
+            meas_data = output.detach().cpu().numpy()
+            if hasattr(self, 'mask') and self.mask is not None:
+                mask = self.mask.detach().cpu().numpy()
+            else:
+                mask = 1
+                
+            if logarithmic:
+                return np.log(meas_data) / np.log(10) * mask
+            else:
+                return meas_data * mask
+
+        translations = self.translations.detach().cpu().numpy()
+        nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
+
+        if logarithmic:
+            cbar_title='Log Base 10 of Diffraction Intensity'
+        else:
+            cbar_title='Diffraction Intensity'
+        
+        plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
+

src/CDTools/datasets/ptycho_2d_dataset.py

@@ -0,0 +1,255 @@
+import numpy as np
+import torch as t
+from copy import copy
+import h5py
+import pathlib
+from CDTools.datasets import CDataset
+from CDTools.tools import data as cdtdata
+from CDTools.tools import plotting
+
+__all__ = ['Ptycho2DDataset']
+
+
+class Ptycho2DDataset(CDataset):
+    """The standard dataset for a 2D ptychography scan
+
+    Subclasses datasets.CDataset
+
+    This class loads and saves 2D ptychography scan data from .cxi files.
+    It should save and load files compatible with most reconstruction
+    programs, although it is only tested against SHARP.
+    """
+    def __init__(self, translations, patterns, intensities=None,
+                 axes=None, *args, **kwargs):
+        """The __init__ function allows construction from python objects.
+
+        The detector_geometry dictionary is defined to have the
+        entries defined by the outputs of data.get_detector_geometry.
+
+        Note that the created dataset object will not copy the data in the
+        patterns parameter in order to avoid doubling the memory requiement
+        for large datasets.
+
+        Parameters
+        ----------
+        translations : array
+            An nx3 array containing the probe translations at each scan point
+        patterns : array
+            An nxmxl array containing the full stack of measured diffraction patterns
+        axes : list(str)
+            A list of names for the axes of the probe translations
+        entry_info : dict
+            A dictionary containing the entry_info metadata
+        sample_info : dict
+            A dictionary containing the sample_info metadata
+        wavelength : float
+            The wavelength of light used in the experiment
+        detector_geometry : dict
+            A dictionary containing the various detector geometry
+            parameters
+        mask : array
+            A mask for the detector, defined as 1 for live pixels, 0
+            for dead
+        background : array
+            An initial guess for the not-previously-subtracted
+            detector background
+        intensities : array
+            A list of measured shot-to-shot intensities
+        """
+        
+
+        super(Ptycho2DDataset,self).__init__(*args, **kwargs)
+        self.axes = copy(axes)
+        self.translations = t.tensor(translations)
+        
+        self.patterns = t.as_tensor(patterns)
+
+        if self.mask is None:
+            self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)
+        self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
+        self.patterns.masked_fill_(t.isnan(self.patterns),0)
+
+        if intensities is not None:
+            self.intensities = t.as_tensor(intensities, dtype=t.float32)
+        else:
+            self.intensities = None
+            
+    def __len__(self):
+        return self.patterns.shape[0]
+
+    def _load(self, index):
+        """ Internal function to load data
+
+        This function is used internally by the global __getitem__ function
+        defined in the base class, which handles moving data around when
+        the dataset is (for example) storing the data on the CPU but
+        getting data as GPU tensors.
+
+        It loads data in the format (inputs, output)
+
+        The inputs for a 2D ptychogaphy data set are:
+
+        1) The indices of the patterns to use
+        2) The recorded probe positions associated with those points
+
+        Parameters
+        ----------
+        index : int or slice
+             The index or indices of the scan points to use
+
+        Returns
+        -------
+        inputs : tuple
+             A tuple of the inputs to the related forward models
+        outputs : tuple
+             The output pattern or stack of output patterns
+        """
+        return (index, self.translations[index]), self.patterns[index]
+
+
+    def to(self, *args, **kwargs):
+        """Sends the relevant data to the given device and dtype
+
+        This function sends the stored translations, patterns,
+        mask and background to the specified device and dtype
+
+        Accepts the same parameters as torch.Tensor.to
+        """
+        super(Ptycho2DDataset,self).to(*args,**kwargs)
+        self.translations = self.translations.to(*args, **kwargs)
+        self.patterns = self.patterns.to(*args, **kwargs)
+
+
+    # It sucks that I can't reuse the base factory method here,
+    # perhaps there is a way but I couldn't figure it out.
+    @classmethod
+    def from_cxi(cls, cxi_file, cut_zeros=True, load_patterns=True):
+        """Generates a new Ptycho2DDataset from a .cxi file directly
+
+        This generates a new Ptycho2DDataset from a .cxi file storing
+        a 2D ptychography scan.
+
+        Parameters
+        ----------
+        file : str, pathlib.Path, or h5py.File
+            The .cxi file to load from
+        cut_zeros : bool
+            Default True, whether to set all negative data to zero
+
+        Returns
+        -------
+        dataset : Ptycho2DDataset
+            The constructed dataset object
+        """
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with h5py.File(cxi_file,'r') as f:
+                return cls.from_cxi(f, cut_zeros=cut_zeros, load_patterns=load_patterns)
+
+        # Generate a base dataset
+        dataset = CDataset.from_cxi(cxi_file)
+        # Mutate the class to this subclass (BasicPtychoDataset)
+        dataset.__class__ = cls
+
+        # Load the data that is only relevant for this class
+        translations = cdtdata.get_ptycho_translations(cxi_file)
+        # And now re-do the stuff from __init__
+        dataset.translations = t.tensor(translations, dtype=t.float32)
+        if load_patterns:
+            patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
+            dataset.patterns = t.as_tensor(patterns)
+            if dataset.patterns.dtype == t.float64:
+                raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request')
+            
+            dataset.axes = axes
+            
+            if dataset.mask is None:
+                dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool)
+
+        try:
+            intensities = cdtdata.get_shot_to_shot_info(cxi_file, 'intensities')
+            dataset.intensities = t.as_tensor(intensities, dtype=t.float32)
+        except KeyError:
+            dataset.intensities = None
+            
+        return dataset
+
+
+    def to_cxi(self, cxi_file):
+        """Saves out a Ptycho2DDataset as a .cxi file
+
+        This function saves all the compatible information in a
+        Ptycho2DDataset object into a .cxi file. This saved .cxi file
+        should be compatible with any standard .cxi file based
+        reconstruction tool, such as SHARP.
+
+        Parameters
+        ----------
+        cxi_file : str, pathlib.Path, or h5py.File
+            The .cxi file to write to
+        """
+
+        # If a bare string is passed
+        if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
+            with cdtdata.create_cxi(cxi_file) as f:
+                return self.to_cxi(f)
+
+        super(Ptycho2DDataset,self).to_cxi(cxi_file)
+        if hasattr(self, 'axes'):
+            cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
+        else:
+            cdtdata.add_data(cxi_file, self.patterns)
+        cdtdata.add_ptycho_translations(cxi_file, self.translations)
+
+        if hasattr(self, 'intensities') and self.intensities is not None:
+            cdtdata.add_shot_to_shot_info(cxi_file, self.intensities, 'intensities')
+
+
+    def inspect(self, logarithmic=True, units='um', log_offset=1):
+        """Launches an interactive plot for perusing the data
+
+        This launches an interactive plotting tool in matplotlib that
+        shows the spatial map constructed from the integrated intensity
+        at each position on the left, next to a panel on the right that
+        can display a base-10 log plot of the detector readout at each
+        position.
+        """
+
+
+        def get_images(idx):
+            inputs, output = self[idx]
+            meas_data = output.detach().cpu().numpy()
+            if hasattr(self, 'mask') and self.mask is not None:
+                mask = self.mask.detach().cpu().numpy()
+            else:
+                mask = 1
+                
+            if logarithmic:
+                return np.log(meas_data + log_offset) / np.log(10) * mask
+            else:
+                return meas_data * mask
+
+        translations = self.translations.detach().cpu().numpy()
+        
+        # This takes about twice as long as it would to just do it all at
+        # once, but it avoids creating another self.patterns-sized array
+        # as an intermediate step. This can be super important because
+        # self.patterns can be more than half the available memory
+        nanomap_values = np.ones(self.translations.shape[0])
+
+        chunk_size = 10
+        for i in range(0, self.translations.shape[0], chunk_size):
+            nanomap_values[i:i+chunk_size] = \
+                t.sum(self.mask * self.patterns[i:i+chunk_size],dim=(1,2))
+
+        # This is the faster but more memory-intensive version
+        # nanomap_values = (self.mask * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
+    
+        if logarithmic:
+            cbar_title = ('Log Base 10 of Diffraction Intensity + %0.2f'
+                          % log_offset)
+        else:
+            cbar_title = 'Diffraction Intensity'
+        
+        return plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
+

src/CDTools/models/__init__.py

@@ -0,0 +1,40 @@
+"""This module contains all the models for different CDI Reconstructions
+
+All the reconstructions are coordinated through the ptychography models
+defined here. The models are, at their core, just subclasses of the 
+:code:`torch.nn.model` class, so they contain the same structure of
+parameters, etc. Their central functionality is as a simulation that maps
+some input (usually, the index number of a scan point) to an output that
+corresponds to the measured data (usually, a diffraction pattern). This
+model can then be used as the heart of an automatic differentiation
+reconstruction which retrieves the parameters that were used in the model.
+
+A main CDIModel class is defined in the base.py file, and models for
+various CDI geometries can be defined as subclasses of this base model.
+The subclasses of the main CDIModel class are required to implement a set of
+functions defined in the base.py file. Example implementations of
+these functions can be found in the code for the SimplePtycho class.
+
+Finally, it is recommended to read through the tutorial section on
+defining a new ptychography model before attempting to do so.
+
+"""
+
+# I don't believe that __all__ really needed, but it's nice to define it
+# to be explicit that import * is safe
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI']
+
+from CDTools.models.base import CDIModel
+from CDTools.models.simple_ptycho import SimplePtycho
+from CDTools.models.fancy_ptycho import FancyPtycho
+from CDTools.models.polarized_fancy_ptycho import PolarizedFancyPtycho
+from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
+from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
+from CDTools.models.rpi import RPI
+from CDTools.models.multimode_rpi import MultimodeRPI
+from CDTools.models.time_resolved_ptycho_calibration import TimeResolvedPtychoCalibration
+from CDTools.models.time_resolved_rpi import TimeResolvedRPI
+
+# Still needs to be updated for the new complex numbers
+#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
+

src/CDTools/models/base.py

@@ -0,0 +1,655 @@
+"""This module contains the base CDIModel class for CDI Models.
+
+The subclasses of the main CDIModel class are required to define their
+own implementations of the following functions:
+
+Loading and Saving
+------------------
+from_dataset
+    Creates a CDIModel from an appropriate CDataset
+simulate_to_dataset
+    Creates a CDataset from the simulation defined in the model
+save_results
+    Saves out a dictionary with the recovered parameters
+
+
+Simulation
+----------
+interaction
+    Simulates exit waves from experimental parameters
+forward_propagator
+    The propagator from the experiment plane to the detector plane
+backward_propagator
+    Optional, the propagator from the detector plane to the experiment plane
+measurement
+    Simulates the detector readout from a detector plane wavefront
+loss
+    the loss function to report and use for automatic differentiation
+
+"""
+
+import torch as t
+from torch.utils import data as torchdata
+from matplotlib import pyplot as plt
+from matplotlib.widgets import Slider
+from matplotlib import ticker
+import numpy as np
+import threading
+import queue
+import time
+from .complex_adam import MyAdam
+from .complex_lbfgs import MyLBFGS
+
+__all__ = ['CDIModel']
+
+
+class CDIModel(t.nn.Module):
+    """This base model defines all the functions that must be exposed for a valid CDIModel subclass
+
+    Most of the functions only raise a NotImplementedError at this level and
+    must be explicitly defined by any subclass - these are noted explocitly
+    in the module-level intro. The work of defining the various subclasses
+    boils down to creating an appropriate implementation for this set of
+    functions.
+    """
+
+    def __init__(self):
+        super(CDIModel,self).__init__()
+
+        self.loss_train = []
+        self.iteration_count = 0
+
+    def from_dataset(self, dataset):
+        raise NotImplementedError()
+
+
+    def interaction(self, *args):
+        raise NotImplementedError()
+
+
+    def forward_propagator(self, exit_wave):
+        raise NotImplementedError()
+
+
+    def backward_propagator(self, detector_wave):
+        raise NotImplementedError()
+
+
+    def measurement(self, detector_wave):
+        raise NotImplementedError()
+
+
+    def forward(self, *args):
+        """The complete forward model
+
+        This model relies on composing the interaction, forward propagator,
+        and measurement functions which are required to be defined by all
+        subclasses. It therefore should not be redefined by the subclasses.
+
+        The arguments to this function, for any given subclass, will be
+        the same as the arguments to the interaction function.
+        """
+        return self.measurement(self.forward_propagator(self.interaction(*args)))
+
+    def loss(self, sim_data, real_data):
+        raise NotImplementedError()
+
+
+    def to(self, *args, **kwargs):
+        super(CDIModel,self).to(*args,**kwargs)
+
+
+    def simulate_to_dataset(self, args_list):
+        raise NotImplementedError()
+
+    def save_results(self):
+        raise NotImplementedError()
+
+    def AD_optimize(self, iterations, data_loader,  optimizer,\
+                    scheduler=None, regularization_factor=None, thread=True,
+                    calculation_width=10):
+        """Runs a round of reconstruction using the provided optimizer
+
+        This is the basic automatic differentiation reconstruction tool
+        which all the other, algorithm-specific tools, use.
+
+        Like all the other optimization routines, it is defined as a
+        generator function which yields the average loss each epoch.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run
+        dataset : CDataset
+            The dataset to reconstruct against
+        optimizer : torch.optim.Optimizer
+            The optimizer to run the reconstruction with
+        scheduler : torch.optim.lr_scheduler._LRScheduler
+            Optional, a learning rate scheduler to use
+        regularization_factor : float or list(float)
+            Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
+        calculation_width : int
+            Default 10, how many translations to pass through at once for each round of gradient accumulation
+        """
+        # First, calculate the normalization
+        normalization = 0
+        for inputs, patterns in data_loader:
+            normalization += t.sum(patterns).cpu().numpy()
+
+        def run_iteration(stop_event=None):
+            loss = 0
+            N = 0
+            t0 = time.time()
+            for inputs, patterns in data_loader:
+                N += 1
+                def closure():
+                    optimizer.zero_grad()
+
+                    input_chunks = [[inp[i:i + calculation_width]
+                                     for inp in inputs]
+                                    for i in range(0, len(inputs[0]),
+                                                   calculation_width)]
+                    pattern_chunks = [patterns[i:i + calculation_width]
+                                      for i in range(0, len(inputs[0]),
+                                                     calculation_width)]
+
+                    total_loss = 0
+                    for inp, pats in zip(input_chunks, pattern_chunks):
+                        # This is just used to allow graceful exit when
+                        # threading
+                        if stop_event is not None and stop_event.is_set():
+                            exit()
+
+                        sim_patterns = self.forward(*inp)
+
+                        if hasattr(self, 'mask'):
+                            loss = self.loss(pats,sim_patterns, mask=self.mask)
+                        else:
+                            loss = self.loss(pats,sim_patterns)
+
+                        loss.backward()
+
+                        total_loss += loss.detach()
+
+                    if regularization_factor is not None \
+                       and hasattr(self, 'regularizer'):
+                        loss = self.regularizer(regularization_factor)
+                        loss.backward()
+                    return total_loss
+
+                loss += optimizer.step(closure).detach().cpu().numpy()
+
+            loss /= normalization
+            if scheduler is not None:
+                scheduler.step(loss)
+
+            self.loss_train.append(loss)
+            self.latest_iteration_time = time.time() - t0
+            return loss
+
+        if thread:
+            result_queue = queue.Queue()
+            stop_event = threading.Event()
+            def target():
+                try:
+                    result_queue.put(run_iteration(stop_event))
+                except Exception as e:
+                    # If something bad happens, put the exception into the
+                    # result queue
+                    result_queue.put(e)
+
+        for it in range(iterations):
+            if thread:
+                calc = threading.Thread(target=target, name='calculator', daemon=True)
+                try:
+                    calc.start()
+                    while calc.is_alive():
+                        if hasattr(self, 'figs'):
+                            self.figs[0].canvas.start_event_loop(0.01)
+                        else:
+                            calc.join()
+
+                except KeyboardInterrupt as e:
+                    stop_event.set()
+                    print('\nAsking execution thread to stop cleanly - please be patient.')
+                    calc.join()
+                    raise e
+
+                res = result_queue.get()
+
+                # If something went wrong in the thead, we'll get an exception
+                if isinstance(res, Exception):
+                    raise res
+
+                yield res
+
+            else:
+                yield run_iteration()
+
+
+    def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005,
+                      schedule=False, amsgrad=False, subset=None,
+                      regularization_factor=None, thread=True,
+                      calculation_width=10):
+        """Runs a round of reconstruction using the Adam optimizer
+
+        This is generally accepted to be the most robust algorithm for use
+        with ptychography. Like all the other optimization routines,
+        it is defined as a generator function, which yields the average
+        loss each epoch.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run
+        dataset : CDataset
+            The dataset to reconstruct against
+        batch_size : int
+            Optional, the size of the minibatches to use
+        lr : float
+            Optional, The learning rate (alpha) to use
+        schedule : float
+            Optional, whether to use the ReduceLROnPlateau scheduler
+        subset : list(int) or int
+            Optional, a pattern index or list of pattern indices to use
+        regularization_factor : float or list(float)
+            Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
+        calculation_width : int
+            Default 1, how many translations to pass through at once for each round of gradient accumulation
+
+        """
+
+        if subset is not None:
+            # if just one pattern, turn into a list for convenience
+            if type(subset) == type(1):
+                subset = [subset]
+            dataset = torchdata.Subset(dataset, subset)
+
+        # Make a dataloader
+        data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
+                                           shuffle=True)
+
+        # Define the optimizer
+        #optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
+        optimizer = MyAdam(self.parameters(), lr = lr, amsgrad=amsgrad)
+
+
+        # Define the scheduler
+        if schedule:
+            scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9)
+        else:
+            scheduler = None
+
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                scheduler=scheduler,
+                                regularization_factor=regularization_factor,
+                                thread=thread,
+                                calculation_width=calculation_width)
+
+
+    def LBFGS_optimize(self, iterations, dataset,
+                       lr=0.1,history_size=2, subset=None,
+                       regularization_factor=None, thread=True,
+                       calculation_width=10, line_search_fn=None):
+        """Runs a round of reconstruction using the L-BFGS optimizer
+
+        This algorithm is often less stable that Adam, however in certain
+        situations or geometries it can be shockingly efficient. Like all
+        the other optimization routines, it is defined as a generator
+        function which yields the average loss each epoch.
+
+        Note: There is no batch size, because it is a usually a bad idea to use
+        LBFGS on anything but all the data at onece
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run
+        dataset : CDataset
+            The dataset to reconstruct against
+        lr : float
+            Optional, the learning rate to use
+        history_size : int
+            Optional, the length of the history to use.
+        subset : list(int) or int
+            Optional, a pattern index or list of pattern indices to ues
+        regularization_factor : float or list(float)
+            Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
+
+        """
+        if subset is not None:
+            # if just one pattern, turn into a list for convenience
+            if type(subset) == type(1):
+                subset = [subset]
+            dataset = torchdata.Subset(dataset, subset)
+
+        # Make a dataloader. This basically does nothing but load all the
+        # data at once
+        data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
+
+
+        # Define the optimizer
+        optimizer = t.optim.LBFGS(self.parameters(),
+                                  lr = lr, history_size=history_size,
+                                  line_search_fn=line_search_fn)
+        #optimizer = MyLBFGS(self.parameters(),
+        #                    lr = lr, history_size=history_size)
+
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                regularization_factor=regularization_factor,
+                                thread=thread,
+                                calculation_width=calculation_width)
+
+
+    def SGD_optimize(self, iterations, dataset, batch_size=None,
+                     lr=0.01, momentum=0, dampening=0, weight_decay=0,
+                     nesterov=False, subset=None, regularization_factor=None,
+                     thread=True, calculation_width=10):
+        """Runs a round of reconstruction using the SGDoptimizer
+
+        This algorithm is often less stable that Adam, but it is simpler
+        and is the basic workhorse of gradience descent.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run
+        dataset : CDataset
+            The dataset to reconstruct against
+        batch_size : int
+            Optional, the size of the minibatches to use
+        lr : float
+            Optional, the learning rate to use
+        momentum : float
+            Optional, the length of the history to use.
+        subset : list(int) or int
+            Optional, a pattern index or list of pattern indices to use
+        regularization_factor : float or list(float)
+            Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
+        calculation_width : int
+            Default 1, how many translations to pass through at once for each round of gradient accumulation
+
+        """
+
+        if subset is not None:
+            # if just one pattern, turn into a list for convenience
+            if type(subset) == type(1):
+                subset = [subset]
+            dataset = torchdata.Subset(dataset, subset)
+
+        # Make a dataloader
+        if batch_size is not None:
+            data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
+                                               shuffle=True)
+        else:
+            data_loader = torchdata.DataLoader(dataset)
+
+
+        # Define the optimizer
+        optimizer = t.optim.SGD(self.parameters(),
+                                lr = lr, momentum=momentum,
+                                dampening=dampening,
+                                weight_decay=weight_decay,
+                                nesterov=nesterov)
+
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                regularization_factor=regularization_factor,
+                                thread=thread,
+                                calculation_width=calculation_width)
+
+
+    def report(self):
+        """Returns a string informing on the latest reconstruction iteration
+
+        Returns
+        -------
+        report : str
+            A string with basic info on the latest iteration
+        """
+        if hasattr(self, 'latest_iteration_time'):
+            return 'Iteration ' + str(len(self.loss_train)) + \
+                  ' completed in %0.2f s with loss ' %\
+                  self.latest_iteration_time + str(self.loss_train[-1])
+        else:
+            return 'No reconstruction iterations performed yet!'
+
+    # By default, the plot_list is empty
+    plot_list = []
+
+
+    def inspect(self, dataset=None, update=True):
+        """Plots all the plots defined in the model's plot_list attribute
+
+        If update is set to True, it will update any previously plotted set
+        of plots, if one exists, and then redraw them. Otherwise, it will
+        plot a new set, and any subsequent updates will update the new set
+
+        Optionally, a dataset can be passed, which then will plot any
+        registered plots which need to incorporate some information from
+        the dataset (such as geometry or a comparison with measured data).
+
+        Plots can be registered in any subclass by defining the plot_list
+        attribute. This should be a list of tuples in the following format:
+        ( 'Plot Title', function_to_generate_plot(self),
+        function_to_determine_whether_to_plot(self))
+
+        Where the third element in the tuple (a function that returns
+        True if the plot is relevant) is not required.
+
+        Parameters
+        ----------
+        dataset : CDataset
+            Optional, a dataset matched to the model type
+        update : bool
+            Default True, whether to update existing plots or plot new ones
+
+        """
+
+        #print('base models inspect: checking the object')
+        #a = self.obj.detach()
+        #def saveobj(a, filename):
+        #    a = np.abs(a)
+        #    plt.imshow(a)
+        #    plt.savefig(filename)
+        #f = ['base_a.png', 'base_b.png', 'base_c.png', 'base_d.png']
+        #comp = [a[i, j, :, :] for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1])]
+        #for i in range(4):
+        #    saveobj(comp[i], f[i])
+        
+        first_update = False
+        if update and hasattr(self, 'figs') and self.figs:
+            figs = self.figs
+        elif update:
+            figs = None
+            self.figs = []
+            first_update = True
+        else:
+            figs = None
+            self.figs = []
+
+        idx = 0
+        for plots in self.plot_list:
+            # If a conditional is included in the plot
+            try:
+                if len(plots) >=3 and not plots[2](self):
+                    continue
+            except TypeError as e:
+                if len(plots) >= 3 and not plots[2](self, dataset):
+                    continue
+
+            name = plots[0]
+            plotter = plots[1]
+
+            if figs is None:
+                fig = plt.figure()
+                self.figs.append(fig)
+            else:
+                fig = figs[idx]
+
+            try:
+                plotter(self,fig)
+                plt.title(name)
+
+            except TypeError as e:
+                if dataset is not None:
+                    try:
+                        plotter(self, fig, dataset)
+                        plt.title(name)
+
+                    except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
+                        pass
+
+            except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
+                pass
+
+            idx += 1
+
+            if update:
+                plt.draw()
+                fig.canvas.start_event_loop(0.001)
+
+        if first_update:
+            plt.pause(0.05 * len(self.figs))
+
+
+    def save_figures(self, prefix='', extension='.eps'):
+        """Saves all currently open inspection figures.
+
+        Note that this function is not very intelligent - so, for example,
+        if multiple probe modes are being reconstructed and the probe
+        plotting function allows one to scroll between different modes, it
+        will simply save whichever mode happens to be showing at the moment.
+        Therefore, this should not be treated as a good way of saving out
+        the full state of the reconstruction.
+
+        By default, the files will be named by the figure titles as defined
+        in the plot_list. Files can be saved with any extension suported by
+        matplotlib.pyplot.savefig.
+
+        Parameters
+        ----------
+        prefix : str
+            Optional, a string to prepend to the saved figure names
+        extention : strategy
+            Default is .eps, the file extension to save with.
+        """
+
+        if hasattr(self, 'figs') and self.figs:
+            figs = self.figs
+        else:
+            return # No figures to save
+
+        for fig in self.figs:
+            fig.savefig(prefix + fig.axes[0].get_title() + extension,
+                        bbox_inches = 'tight')
+
+
+    def compare(self, dataset, logarithmic=False):
+        """Opens a tool for comparing simulated and measured diffraction patterns
+
+        Parameters
+        ----------
+        dataset : CDataset
+            A dataset containing the simulated diffraction patterns to compare against
+        """
+
+        fig, axes = plt.subplots(1,3,figsize=(12,5.3))
+        fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
+        axslider = plt.axes([0.15,0.06,0.75,0.03])
+
+
+        def update_colorbar(im):
+            # If the update brought the colorbar out of whack
+            # (say, from clicking back in the navbar)
+            # Holy fuck this was annoying. Sorry future for how
+            # crappy this solution is.
+            #if not np.allclose(im.colorbar.ax.get_xlim(),
+            #                   (np.min(im.get_array()),
+            #                    np.max(im.get_array()))):
+            if hasattr(im, 'norecurse') and im.norecurse:
+                im.norecurse=False
+                return
+
+            im.norecurse=True
+            im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
+
+        def update(idx):
+            idx = int(idx) % len(dataset)
+            fig.pattern_idx = idx
+            updating = True if len(axes[0].images) >= 1 else False
+
+            inputs, output = dataset[idx]
+            sim_data = self.forward(*inputs).detach().cpu().numpy()
+            meas_data = output.detach().cpu().numpy()
+            if hasattr(self, 'mask') and self.mask is not None:
+                mask = self.mask.detach().cpu().numpy()
+            else:
+                mask = 1
+
+            if logarithmic:
+                sim_data =np.log(sim_data)/np.log(10)
+                meas_data = np.log(meas_data)/np.log(10)
+
+            if not updating:
+                axes[0].set_title('Simulated')
+                axes[1].set_title('Measured')
+                axes[2].set_title('Difference')
+
+                sim = axes[0].imshow(sim_data)
+                meas = axes[1].imshow(meas_data * mask)
+                diff = axes[2].imshow((sim_data-meas_data) * mask)
+
+                cb1 = plt.colorbar(sim, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
+                cb1.ax.tick_params(labelrotation=20)
+                cb1.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(sim))
+                cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
+                cb2.ax.tick_params(labelrotation=20)
+                cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
+                cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
+                cb3.ax.tick_params(labelrotation=20)
+                cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
+
+            else:
+
+                sim = axes[0].images[-1]
+                sim.set_data(sim_data)
+                update_colorbar(sim)
+
+                meas = axes[1].images[-1]
+                meas.set_data(meas_data * mask)
+                update_colorbar(meas)
+
+                diff = axes[2].images[-1]
+                diff.set_data((sim_data-meas_data) * mask)
+                update_colorbar(diff)
+
+
+        # This is dumb but the slider doesn't work unless a reference to it is
+        # kept somewhere...
+        self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d")
+        self.slider.on_changed(update)
+
+        def on_action(event):
+            if not hasattr(event, 'button'):
+                event.button = None
+            if not hasattr(event, 'key'):
+                event.key = None
+
+            if event.key == 'up' or event.button == 'up':
+                update(fig.pattern_idx - 1)
+            elif event.key == 'down' or event.button == 'down':
+                update(fig.pattern_idx + 1)
+            self.slider.set_val(fig.pattern_idx)
+            plt.draw()
+
+        fig.canvas.mpl_connect('key_press_event',on_action)
+        fig.canvas.mpl_connect('scroll_event',on_action)
+        update(0)

src/CDTools/models/bragg_2d_ptycho.py

@@ -0,0 +1,533 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools.propagators import generate_generalized_angular_spectrum_propagator as ggasp
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+__all__ = ['Bragg2DPtycho']
+
+#
+# Key ideas:
+# 1) To a first approximation, do the reconstruction on a parallelogram
+#    shaped grid on the sample which is conjugate to the detector coordinates
+# 2) Simulate a probe in those same coordinates, but propagate it back and
+#    forth (along with the translations), using the angular spectrum method
+# 3) Apply a correction to the simulated data to account for the tilt of the
+#    sample with respect to the detector
+# 4) Include a correction for the thickness of the sample
+#
+
+#
+# How to do this properly?
+# First thing to note is that the two corrections (probe propagation before
+# interaction and high-NA correction for the final diffraction measurement)
+# should be able to be turned on separately, since they show up in different
+# situations. In fact, I would like to focus on the first aspect initially
+# since I think that's the dominant issue we will contend with at CSX.
+#
+
+#
+# It also should be possible to choose an "auto" setting for the two
+# corrections, since the geometry information given should be enough to
+# decide if the correction is needed.
+# Probably the automatic check will have to be very conservative for the
+# probe propagation side since the model has no information about the
+# 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):
+
+    # Needed to do the real/complex split
+    #@property
+    #def obj(self):
+    #    return t.complex(self.obj_real, self.obj_imag)
+
+    #@property
+    #def probe(self):
+    #    return t.complex(self.probe_real, self.probe_imag)
+
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis, probe_guess, obj_guess,
+                 detector_slice=None,
+                 min_translation=t.tensor([0, 0], dtype=t.float32),
+                 median_propagation=t.tensor(0, dtype=t.float32),
+                 background=None, translation_offsets=None, mask=None,
+                 weights=None, translation_scale=1, saturation=None,
+                 probe_support=None, oversampling=1,
+                 propagate_probe=True, correct_tilt=True, lens=False):
+
+        # We need the detector geometry
+        # We need the probe basis (but in this case, we don't need the surface
+        # normal because it comes implied by the probe basis
+        # we do need the detector slice I suppose
+        # The min translation is also needed
+        # The median propagation should be needed as well
+        # translation_offsets can stay 2D for now
+        # propagate_probe and correct_tilt are important!
+
+        super(Bragg2DPtycho, 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.median_propagation = t.tensor(median_propagation)
+        
+        self.probe_basis = t.tensor(probe_basis)
+        self.detector_slice = copy(detector_slice)
+
+        # calculate the surface normal from the probe basis
+        surface_normal =  np.cross(np.array(probe_basis)[:,1],
+                           np.array(probe_basis)[:,0])
+        surface_normal /= np.linalg.norm(surface_normal)
+        self.surface_normal = t.tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 2:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
+        else:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess))        
+
+        # Not strictly necessary but otherwise it will return
+        # a probe with the stuff outside of the support unchanged after
+        # optimization
+        if probe_support is not None:
+            self.probe_support = t.tensor(probe_support, dtype=t.bool)
+            probe_guess = probe_guess * probe_support
+            # This seems dumb, but otherwise it winds up with a mixture
+            # of negative and positive zeros and it's super annoying when
+            # you look at the phase map
+            probe_guess[probe_guess == 0] = 0
+        else:
+            self.probe_support = t.ones(probe_guess[0].shape, dtype=t.bool)
+
+        #self.probe_real = t.nn.Parameter(probe_guess.real / self.probe_norm)
+        #self.probe_imag = t.nn.Parameter(probe_guess.imag / self.probe_norm)
+
+        #self.obj_real = t.nn.Parameter(obj_guess.real )
+        #self.obj_imag = t.nn.Parameter(obj_guess.imag)
+
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess)
+
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(
+                    self.probe[0][self.detector_slice].shape,
+                    dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+
+        self.background = t.nn.Parameter(background)
+
+        if weights is None:
+            self.weights = None
+        else:
+            # No incoherent + unstable here yet
+            self.weights = t.nn.Parameter(t.tensor(weights,
+                                                   dtype=t.float32))
+
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            t_o = t.tensor(translation_offsets, dtype=t.float32)
+            t_o = t_o / translation_scale
+            self.translation_offsets = t.nn.Parameter(t_o)
+
+        self.translation_scale = translation_scale
+
+        self.oversampling = oversampling
+
+        self.propagate_probe = propagate_probe
+        self.correct_tilt = correct_tilt
+        if correct_tilt:
+            # recall that here we always want the shape of the detector
+            # before it's cut down by the detector slice to match the
+            # physical detector region
+            probe_shape = self.probe[0].shape
+            
+            self.k_map, self.intensity_map = \
+                tools.propagators.generate_high_NA_k_intensity_map(
+                    self.probe_basis,
+                    self.detector_geometry['basis'] / oversampling,
+                    probe_shape,
+                    self.detector_geometry['distance'],
+                    self.wavelength,dtype=t.float32,
+                    lens=lens)
+        else:
+            self.k_map = None
+            self.intensity_map = None
+
+        self.prop_dir = t.tensor([0, 0, 1], dtype=t.float32)
+
+        # This propagator should be able to be multiplied by the propagation
+        # distance each time to get a propagator
+        self.universal_propagator = t.angle(ggasp(
+            self.probe.shape[1:],
+            self.probe_basis, self.wavelength,
+            t.tensor([0, 0, self.wavelength/(2*np.pi)], dtype=t.float32),
+            propagation_vector=self.prop_dir,
+            dtype=t.complex64,
+            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, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True, correct_tilt=True, lens=False, opt_for_fft=False):
+        
+        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 exit wave geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        ew_basis, ew_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center,
+                                                   padding=padding,
+                                                   opt_for_fft=opt_for_fft,
+                                                   oversampling=oversampling)
+        # now we grab the sample surface normal
+        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)
+
+        # and we use that to generate the probe basis
+
+        ew_normal = np.cross(np.array(ew_basis)[:,1],
+                           np.array(ew_basis)[:,0])
+        ew_normal /= np.linalg.norm(ew_normal)
+
+        # This is a bit of an odd way to do a projection but I think it's
+        # the most compact and reliable. We set up two matrix-vector equations
+        # to enforce the two conditions (on the plane normal to the surface
+        # normal and in the ew_normal direction from the original point
+
+        mat = np.vstack([np.eye(3) - np.outer(ew_normal,ew_normal),
+                         surface_normal])
+
+        # We know that this matrix multiplied by the final result must
+        # equal the input vector with a trailing 0, so we can do the
+        # projection with a pseudoinverse and removing the last column
+
+        projector = np.linalg.pinv(mat)[:, :3]
+        probe_basis = t.Tensor(np.dot(projector, ew_basis))
+
+        # Now we need a much better way to handle the translations here
+        # than translations_to_pixel
+        
+        # Next generate the object geometry from the probe geometry and
+        # the translations
+        p2s = tools.interactions.project_translations_to_sample
+        pix_translations, propagations = p2s(probe_basis, translations)
+
+
+        obj_size, min_translation = tools.initializers.calc_object_setup(ew_shape, pix_translations, padding=200)
+
+        median_propagation = t.median(propagations)
+
+        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
+        # Because the grid we defined on the sample is projected from the
+        # detector conjugate space, we can pretend that the grid is just in
+        # that space and use the standard initializations anyway
+        
+        if probe_size is None:
+            probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        else:
+            probe = tools.initializers.gaussian_probe(dataset, ew_basis, ew_shape, probe_size, propagation_distance=propagation_distance)
+
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(t.abs(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 = t.exp(1j*(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(probe[0].shape, dtype=t.bool)
+            xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
+            xs = xs - np.mean(xs)
+            ys = ys - np.mean(ys)
+            Rs = np.sqrt(xs**2 + ys**2)
+
+            probe_support[Rs < probe_support_radius] = 1
+            probe = probe * probe_support[None, :, :]
+
+        else:
+            probe_support = None
+
+        # Here we need to implement a simple condition to choose whether
+        # to propagate the probe or not
+        if not( propagate_probe is True or propagate_probe is False):
+            raise NotImplementedError('No auto option implemented yet')
+
+        if not(correct_tilt is True or correct_tilt is False):
+            raise NotImplementedError('No auto option implemented yet')
+
+            
+        return cls(wavelength, det_geo, probe_basis, probe, obj,
+                   detector_slice=det_slice,
+                   min_translation=min_translation,
+                   median_propagation =median_propagation,
+                   translation_offsets = translation_offsets,
+                   weights=weights, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   oversampling=oversampling,
+                   propagate_probe=propagate_probe,
+                   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
+        
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+
+        Ws = self.weights[index]
+        prs = Ws[...,None,None,None] * self.probe * self.probe_support[...,:,:]
+        # Now we need to propagate each of the probes
+
+        
+        for j in range(prs.shape[0]):
+            # I believe this -1 sign is in error, but I need a dataset with
+            # well understood geometry to figure it out
+            propagator = t.exp(
+                1j*(props[j]*(2*np.pi)/self.wavelength)
+                * self.universal_propagator)
+            prs[j] = tools.propagators.near_field(prs[j], propagator)
+
+        exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
+            prs, self.obj,pix_trans,
+            shift_probe=True, multiple_modes=True)
+
+        return exit_waves
+
+        
+    def forward_propagator(self, wavefields):
+        if self.correct_tilt:
+            return tools.propagators.high_NA_far_field(
+                wavefields,self.k_map,intensity_map=self.intensity_map)
+        else:
+            return tools.propagators.far_field(wavefields)
+
+
+    def backward_propagator(self, wavefields):
+        if self.correct_tilt:
+            assert NotImplementedError('Backward propagator not defined with tilt correction')
+        else:
+            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(Bragg2DPtycho, 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)
+
+        if self.k_map is not None:
+            self.k_map = self.k_map.to(*args,**kwargs)
+        if self.intensity_map is not None:
+            self.intensity_map = self.intensity_map.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.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)
+
+
+        
+    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.real.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
+
+
+    plot_list = [
+        ('Dominant Probe Amplitude',
+         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)),
+        ('Subdominant Probe Amplitude',
+         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),
+         lambda self: len(self.probe) >=2),
+        ('Object Amplitude', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig)),
+        ('Object Phase',
+         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',
+         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 = self.probe.detach().cpu().numpy()
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+        losses = np.array(self.loss_train)
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'obj':obj,
+                'background':background,
+                'weights':weights,
+                'losses':losses}

src/CDTools/models/complex_adam.py

@@ -0,0 +1,163 @@
+import math
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+from typing import List, Optional
+
+class MyAdam(Optimizer):
+    r"""Implements Adam algorithm.
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    The implementation of the L2 penalty follows changes proposed in
+    `Decoupled Weight Decay Regularization`_.
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(MyAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(MyAdam, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            max_exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group['betas']
+
+            for p in group['params']:
+                if p.grad is not None:
+                    params_with_grad.append(p)
+                    if p.grad.is_sparse:
+                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                    grads.append(p.grad)
+
+                    state = self.state[p]
+                    # Lazy state initialization
+                    if len(state) == 0:
+                        state['step'] = 0
+                        # Exponential moving average of gradient values
+                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                        # Exponential moving average of squared gradient values
+                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                        if group['amsgrad']:
+                            # Maintains max of all exp. moving avg. of sq. grad. values
+                            state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                    exp_avgs.append(state['exp_avg'])
+                    exp_avg_sqs.append(state['exp_avg_sq'])
+
+                    if group['amsgrad']:
+                        max_exp_avg_sqs.append(state['max_exp_avg_sq'])
+
+                    # update the steps for each param group update
+                    state['step'] += 1
+                    # record the step after step update
+                    state_steps.append(state['step'])
+
+            adam(params_with_grad,
+                   grads,
+                   exp_avgs,
+                   exp_avg_sqs,
+                   max_exp_avg_sqs,
+                   state_steps,
+                   amsgrad=group['amsgrad'],
+                   beta1=beta1,
+                   beta2=beta2,
+                   lr=group['lr'],
+                   weight_decay=group['weight_decay'],
+                   eps=group['eps'])
+        return loss
+
+
+    
+def adam(params: List[Tensor],
+         grads: List[Tensor],
+         exp_avgs: List[Tensor],
+         exp_avg_sqs: List[Tensor],
+         max_exp_avg_sqs: List[Tensor],
+         state_steps: List[int],
+         *,
+         amsgrad: bool,
+         beta1: float,
+         beta2: float,
+         lr: float,
+         weight_decay: float,
+         eps: float):
+    r"""Functional API that performs Adam algorithm computation.
+    See :class:`~torch.optim.Adam` for details.
+    """
+
+    for i, param in enumerate(params):
+
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step = state_steps[i]
+
+        bias_correction1 = 1 - beta1 ** step
+        bias_correction2 = 1 - beta2 ** step
+
+        if weight_decay != 0:
+            grad = grad.add(param, alpha=weight_decay)
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
+        if amsgrad:
+            # Maintains the maximum of all 2nd moment running avg. till now
+            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
+            # Use the max. for normalizing running avg. of gradient
+            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
+        else:
+            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
+
+        step_size = lr / bias_correction1
+
+        param.addcdiv_(exp_avg, denom, value=-step_size)

src/CDTools/models/complex_lbfgs.py

@@ -0,0 +1,485 @@
+import torch
+from functools import reduce
+from torch.optim.optimizer import Optimizer
+
+
+def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
+    # ported from https://github.com/torch/optim/blob/master/polyinterp.lua
+    # Compute bounds of interpolation area
+    if bounds is not None:
+        xmin_bound, xmax_bound = bounds
+    else:
+        xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
+
+    # Code for most common case: cubic interpolation of 2 points
+    #   w/ function and derivative values for both
+    # Solution in this case (where x2 is the farthest point):
+    #   d1 = g1 + g2 - 3*(f1-f2)/(x1-x2);
+    #   d2 = sqrt(d1^2 - g1*g2);
+    #   min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2));
+    #   t_new = min(max(min_pos,xmin_bound),xmax_bound);
+    d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
+    d2_square = d1**2 - g1 * g2
+    if d2_square >= 0:
+        d2 = d2_square.sqrt()
+        if x1 <= x2:
+            min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
+        else:
+            min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
+        return min(max(min_pos, xmin_bound), xmax_bound)
+    else:
+        return (xmin_bound + xmax_bound) / 2.
+
+
+def _strong_wolfe(obj_func,
+                  x,
+                  t,
+                  d,
+                  f,
+                  g,
+                  gtd,
+                  c1=1e-4,
+                  c2=0.9,
+                  tolerance_change=1e-9,
+                  max_ls=25):
+    # ported from https://github.com/torch/optim/blob/master/lswolfe.lua
+    d_norm = d.abs().max()
+    g = g.clone(memory_format=torch.contiguous_format)
+    # evaluate objective and gradient using initial step
+    f_new, g_new = obj_func(x, t, d)
+    ls_func_evals = 1
+    gtd_new = g_new.dot(d)
+
+    # bracket an interval containing a point satisfying the Wolfe criteria
+    t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd
+    done = False
+    ls_iter = 0
+    while ls_iter < max_ls:
+        # check conditions
+        if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev):
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
+            bracket_gtd = [gtd_prev, gtd_new]
+            break
+
+        if abs(gtd_new) <= -c2 * gtd:
+            bracket = [t]
+            bracket_f = [f_new]
+            bracket_g = [g_new]
+            done = True
+            break
+
+        if gtd_new >= 0:
+            bracket = [t_prev, t]
+            bracket_f = [f_prev, f_new]
+            bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
+            bracket_gtd = [gtd_prev, gtd_new]
+            break
+
+        # interpolate
+        min_step = t + 0.01 * (t - t_prev)
+        max_step = t * 10
+        tmp = t
+        t = _cubic_interpolate(
+            t_prev,
+            f_prev,
+            gtd_prev,
+            t,
+            f_new,
+            gtd_new,
+            bounds=(min_step, max_step))
+
+        # next step
+        t_prev = tmp
+        f_prev = f_new
+        g_prev = g_new.clone(memory_format=torch.contiguous_format)
+        gtd_prev = gtd_new
+        f_new, g_new = obj_func(x, t, d)
+        ls_func_evals += 1
+        gtd_new = g_new.dot(d)
+        ls_iter += 1
+
+    # reached max number of iterations?
+    if ls_iter == max_ls:
+        bracket = [0, t]
+        bracket_f = [f, f_new]
+        bracket_g = [g, g_new]
+
+    # zoom phase: we now have a point satisfying the criteria, or
+    # a bracket around it. We refine the bracket until we find the
+    # exact point satisfying the criteria
+    insuf_progress = False
+    # find high and low points in bracket
+    low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
+    while not done and ls_iter < max_ls:
+        # line-search bracket is so small
+        if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
+            break
+
+        # compute new trial value
+        t = _cubic_interpolate(bracket[0], bracket_f[0], bracket_gtd[0],
+                               bracket[1], bracket_f[1], bracket_gtd[1])
+
+        # test that we are making sufficient progress:
+        # in case `t` is so close to boundary, we mark that we are making
+        # insufficient progress, and if
+        #   + we have made insufficient progress in the last step, or
+        #   + `t` is at one of the boundary,
+        # we will move `t` to a position which is `0.1 * len(bracket)`
+        # away from the nearest boundary point.
+        eps = 0.1 * (max(bracket) - min(bracket))
+        if min(max(bracket) - t, t - min(bracket)) < eps:
+            # interpolation close to boundary
+            if insuf_progress or t >= max(bracket) or t <= min(bracket):
+                # evaluate at 0.1 away from boundary
+                if abs(t - max(bracket)) < abs(t - min(bracket)):
+                    t = max(bracket) - eps
+                else:
+                    t = min(bracket) + eps
+                insuf_progress = False
+            else:
+                insuf_progress = True
+        else:
+            insuf_progress = False
+
+        # Evaluate new point
+        f_new, g_new = obj_func(x, t, d)
+        ls_func_evals += 1
+        gtd_new = g_new.dot(d)
+        ls_iter += 1
+
+        if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]:
+            # Armijo condition not satisfied or not lower than lowest point
+            bracket[high_pos] = t
+            bracket_f[high_pos] = f_new
+            bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format)
+            bracket_gtd[high_pos] = gtd_new
+            low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
+        else:
+            if abs(gtd_new) <= -c2 * gtd:
+                # Wolfe conditions satisfied
+                done = True
+            elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
+                # old high becomes new low
+                bracket[high_pos] = bracket[low_pos]
+                bracket_f[high_pos] = bracket_f[low_pos]
+                bracket_g[high_pos] = bracket_g[low_pos]
+                bracket_gtd[high_pos] = bracket_gtd[low_pos]
+
+            # new point becomes new low
+            bracket[low_pos] = t
+            bracket_f[low_pos] = f_new
+            bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format)
+            bracket_gtd[low_pos] = gtd_new
+
+    # return stuff
+    t = bracket[low_pos]
+    f_new = bracket_f[low_pos]
+    g_new = bracket_g[low_pos]
+    return f_new, g_new, t, ls_func_evals
+
+
+class MyLBFGS(Optimizer):
+    """Implements L-BFGS algorithm, heavily inspired by `minFunc
+    <https://www.cs.ubc.ca/~schmidtm/Software/minFunc.html>`_.
+    .. warning::
+        This optimizer doesn't support per-parameter options and parameter
+        groups (there can be only one).
+    .. warning::
+        Right now all parameters have to be on a single device. This will be
+        improved in the future.
+    .. note::
+        This is a very memory intensive optimizer (it requires additional
+        ``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory
+        try reducing the history size, or use a different algorithm.
+    Args:
+        lr (float): learning rate (default: 1)
+        max_iter (int): maximal number of iterations per optimization step
+            (default: 20)
+        max_eval (int): maximal number of function evaluations per optimization
+            step (default: max_iter * 1.25).
+        tolerance_grad (float): termination tolerance on first order optimality
+            (default: 1e-5).
+        tolerance_change (float): termination tolerance on function
+            value/parameter changes (default: 1e-9).
+        history_size (int): update history size (default: 100).
+        line_search_fn (str): either 'strong_wolfe' or None (default: None).
+    """
+
+    def __init__(self,
+                 params,
+                 lr=1,
+                 max_iter=20,
+                 max_eval=None,
+                 tolerance_grad=1e-7,
+                 tolerance_change=1e-9,
+                 history_size=100,
+                 line_search_fn=None):
+        if max_eval is None:
+            max_eval = max_iter * 5 // 4
+        defaults = dict(
+            lr=lr,
+            max_iter=max_iter,
+            max_eval=max_eval,
+            tolerance_grad=tolerance_grad,
+            tolerance_change=tolerance_change,
+            history_size=history_size,
+            line_search_fn=line_search_fn)
+        super(MyLBFGS, self).__init__(params, defaults)
+
+        if len(self.param_groups) != 1:
+            raise ValueError("LBFGS doesn't support per-parameter options "
+                             "(parameter groups)")
+
+        self._params = self.param_groups[0]['params']
+        self._numel_cache = None
+
+    def _numel(self):
+        if self._numel_cache is None:
+            self._numel_cache = reduce(lambda total, p: total + p.numel(), self._params, 0)
+        return self._numel_cache
+
+    def _gather_flat_grad(self):
+        views = []
+        for p in self._params:
+            if p.grad is None:
+                view = p.new(p.numel()).zero_()
+            elif p.grad.is_sparse:
+                view = p.grad.to_dense().view(-1)
+            else:
+                view = p.grad.view(-1)
+            views.append(view)
+        return torch.cat(views, 0)
+
+    def _add_grad(self, step_size, update):
+        offset = 0
+        for p in self._params:
+            numel = p.numel()
+            if ((update.dtype==torch.complex64 or update.dtype==torch.complex128)
+                and (p.dtype==torch.float32 or p.dtype==torch.float64)):
+                p.add_(update[offset:offset + numel].real.view_as(p), alpha=step_size)
+            else:
+                p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)                
+            # view as to avoid deprecated pointwise semantics
+            #try:
+            #    print(p.dtype)
+            #    print(update.dtype)
+            #    p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
+            #    print('Worked fine')
+            #    print(update[offset:offset + numel].view_as(p))
+            #except:
+            #    print('Failed')
+            #    print(update[offset:offset + numel].view_as(p))
+            #    exit()
+            offset += numel
+        assert offset == self._numel()
+
+    def _clone_param(self):
+        return [p.clone(memory_format=torch.contiguous_format) for p in self._params]
+
+    def _set_param(self, params_data):
+        for p, pdata in zip(self._params, params_data):
+            p.copy_(pdata)
+
+    def _directional_evaluate(self, closure, x, t, d):
+        self._add_grad(t, d)
+        loss = float(closure())
+        flat_grad = self._gather_flat_grad()
+        self._set_param(x)
+        return loss, flat_grad
+
+    @torch.no_grad()
+    def step(self, closure):
+        """Performs a single optimization step.
+        Args:
+            closure (callable): A closure that reevaluates the model
+                and returns the loss.
+        """
+        assert len(self.param_groups) == 1
+
+        # Make sure the closure is always called with grad enabled
+        closure = torch.enable_grad()(closure)
+
+        group = self.param_groups[0]
+        lr = group['lr']
+        max_iter = group['max_iter']
+        max_eval = group['max_eval']
+        tolerance_grad = group['tolerance_grad']
+        tolerance_change = group['tolerance_change']
+        line_search_fn = group['line_search_fn']
+        history_size = group['history_size']
+
+        # NOTE: LBFGS has only global state, but we register it as state for
+        # the first param, because this helps with casting in load_state_dict
+        state = self.state[self._params[0]]
+        state.setdefault('func_evals', 0)
+        state.setdefault('n_iter', 0)
+
+        # evaluate initial f(x) and df/dx
+        orig_loss = closure()
+        loss = float(orig_loss)
+        current_evals = 1
+        state['func_evals'] += 1
+
+        flat_grad = self._gather_flat_grad()
+        opt_cond = flat_grad.abs().max() <= tolerance_grad
+        
+        # optimal condition
+        if opt_cond:
+            return orig_loss
+
+        # tensors cached in state (for tracing)
+        d = state.get('d')
+        t = state.get('t')
+        old_dirs = state.get('old_dirs')
+        old_stps = state.get('old_stps')
+        ro = state.get('ro')
+        H_diag = state.get('H_diag')
+        prev_flat_grad = state.get('prev_flat_grad')
+        prev_loss = state.get('prev_loss')
+
+        n_iter = 0
+        # optimize for a max of max_iter iterations
+        while n_iter < max_iter:
+            # keep track of nb of iterations
+            n_iter += 1
+            state['n_iter'] += 1
+
+            ############################################################
+            # compute gradient descent direction
+            ############################################################
+            if state['n_iter'] == 1:
+                d = flat_grad.neg()
+                old_dirs = []
+                old_stps = []
+                ro = []
+                H_diag = 1
+            else:
+                # do lbfgs update (update memory)
+                y = flat_grad.sub(prev_flat_grad)
+                s = d.mul(t)
+                ys = y.dot(s)  # y*s
+                if ys.abs() > 1e-10:
+                    # updating memory
+                    if len(old_dirs) == history_size:
+                        # shift history by one (limited-memory)
+                        old_dirs.pop(0)
+                        old_stps.pop(0)
+                        ro.pop(0)
+
+                    # store new direction/step
+                    old_dirs.append(y)
+                    old_stps.append(s)
+                    ro.append(1. / ys)
+
+                    # update scale of initial Hessian approximation
+                    H_diag = ys / y.dot(y)  # (y*y)
+
+                # compute the approximate (L-BFGS) inverse Hessian
+                # multiplied by the gradient
+                num_old = len(old_dirs)
+
+                if 'al' not in state:
+                    state['al'] = [None] * history_size
+                al = state['al']
+
+                # iteration in L-BFGS loop collapsed to use just one buffer
+                q = flat_grad.neg()
+                for i in range(num_old - 1, -1, -1):
+                    al[i] = old_stps[i].dot(q) * ro[i]
+                    q.add_(old_dirs[i], alpha=-al[i])
+
+                # multiply by initial Hessian
+                # r/d is the final direction
+                d = r = torch.mul(q, H_diag)
+                for i in range(num_old):
+                    be_i = old_dirs[i].dot(r) * ro[i]
+                    r.add_(old_stps[i], alpha=al[i] - be_i)
+
+            if prev_flat_grad is None:
+                prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format)
+            else:
+                prev_flat_grad.copy_(flat_grad)
+            prev_loss = loss
+
+            ############################################################
+            # compute step length
+            ############################################################
+            # reset initial guess for step size
+            if state['n_iter'] == 1:
+                t = min(1., 1. / flat_grad.abs().sum()) * lr
+            else:
+                t = lr
+
+            # directional derivative
+            gtd = flat_grad.dot(d)  # g * d
+            # Why did this used to be gtd > -tol_change?
+            
+            # directional derivative is below tolerance
+            if gtd.abs() < tolerance_change:
+                break
+
+            # optional line search: user function
+            ls_func_evals = 0
+            if line_search_fn is not None:
+                # perform line search, using user function
+                if line_search_fn != "strong_wolfe":
+                    raise RuntimeError("only 'strong_wolfe' is supported")
+                else:
+                    x_init = self._clone_param()
+
+                    def obj_func(x, t, d):
+                        return self._directional_evaluate(closure, x, t, d)
+
+                    loss, flat_grad, t, ls_func_evals = _strong_wolfe(
+                        obj_func, x_init, t, d, loss, flat_grad, gtd)
+                self._add_grad(t, d)
+                opt_cond = flat_grad.abs().max() <= tolerance_grad
+            else:
+                # no line search, simply move with fixed-step
+                self._add_grad(t, d)
+                if n_iter != max_iter:
+                    # re-evaluate function only if not in last iteration
+                    # the reason we do this: in a stochastic setting,
+                    # no use to re-evaluate that function here
+                    with torch.enable_grad():
+                        loss = float(closure())
+                    flat_grad = self._gather_flat_grad()
+                    opt_cond = flat_grad.abs().max() <= tolerance_grad
+                    ls_func_evals = 1
+
+            # update func eval
+            current_evals += ls_func_evals
+            state['func_evals'] += ls_func_evals
+
+            ############################################################
+            # check conditions
+            ############################################################
+            if n_iter == max_iter:
+                break
+
+            if current_evals >= max_eval:
+                break
+
+            # optimal condition
+            if opt_cond:
+                break
+
+            # lack of progress
+            if d.mul(t).abs().max() <= tolerance_change:
+                break
+
+            if abs(loss - prev_loss) < tolerance_change:
+                break
+
+        state['d'] = d
+        state['t'] = t
+        state['old_dirs'] = old_dirs
+        state['old_stps'] = old_stps
+        state['ro'] = ro
+        state['H_diag'] = H_diag
+        state['prev_flat_grad'] = prev_flat_grad
+        state['prev_loss'] = prev_loss
+
+        return orig_loss

src/CDTools/models/fancy_ptycho.py

@@ -0,0 +1,647 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools import analysis
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from scipy import linalg as sla
+from copy import copy
+
+__all__ = ['FancyPtycho']
+
+
+class FancyPtycho(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess,
+                 obj_guess,
+                 detector_slice=None,
+                 surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
+                 min_translation=t.tensor([0, 0], dtype=t.float32),
+                 background=None,
+                 translation_offsets=None,
+                 mask=None,
+                 weights=None,
+                 translation_scale=1,
+                 saturation=None,
+                 probe_support=None,
+                 oversampling=1,
+                 fourier_probe=False,
+                 loss='amplitude mse',
+                 units='um',
+                 simulate_probe_translation=False
+                 ):
+
+        super(FancyPtycho, self).__init__()
+        self.wavelength = t.tensor(wavelength)
+        self.detector_geometry = copy(detector_geometry)
+        det_geo = self.detector_geometry
+        if 'distance' in det_geo:
+            det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
+        if 'basis' in det_geo:
+            det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
+        if 'corner' in det_geo and det_geo['corner'] is not None:
+            det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
+
+        self.min_translation = t.tensor(min_translation)
+
+        self.probe_basis = t.tensor(probe_basis)
+        self.detector_slice = copy(detector_slice)
+        self.surface_normal = t.tensor(surface_normal)
+
+        self.saturation = saturation
+        self.units = units
+        self.fourier_probe = fourier_probe
+
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+        
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 2:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
+        else:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess))        
+        
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess)
+
+        if background is None:
+            if detector_slice is not None:
+                dummy_det = t.empty([s//oversampling
+                                     for s in self.probe[0].shape])
+                shape = dummy_det[self.detector_slice].shape
+                #shape = self.probe[0][self.detector_slice].shape
+            else:
+                shape = [s//oversampling for s in self.probe[0]]
+            background = 1e-6 * t.ones(shape, dtype=t.float32)
+            
+        self.background = t.nn.Parameter(background)
+
+        if weights is None:
+            self.weights = None
+        else:
+            # We now need to distinguish between real-valued per-image
+            # weights and complex-valued per-mode weight matrices
+            if len(weights.shape) == 1:
+                # This is if it's just a list of numbers
+                self.weights = t.nn.Parameter(t.tensor(weights,
+                                                       dtype=t.float32))
+            else:
+                # Now this is a matrix of weights, so it needs to be complex
+                self.weights = t.nn.Parameter(t.tensor(weights,
+                                                       dtype=t.complex64))
+
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            t_o = t.tensor(translation_offsets, dtype=t.float32)
+            t_o = t_o / translation_scale
+            self.translation_offsets = t.nn.Parameter(t_o)
+
+        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], dtype=t.bool)
+
+        self.oversampling = oversampling
+
+        self.simulate_probe_translation = simulate_probe_translation
+        if simulate_probe_translation:
+            Is = t.arange(self.probe.shape[-2], dtype=t.float32)
+            Js = t.arange(self.probe.shape[-1], dtype=t.float32)
+            Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
+            self.I_phase = 2 * np.pi* Is
+            self.J_phase = 2 * np.pi* Js
+            
+        # Here we set the appropriate loss function
+        if (loss.lower().strip() == 'amplitude mse'
+                or loss.lower().strip() == 'amplitude_mse'):
+            self.loss = tools.losses.amplitude_mse
+        elif (loss.lower().strip() == 'poisson nll'
+                or loss.lower().strip() == 'poisson_nll'):
+            self.loss = tools.losses.poisson_nll
+        else:
+            raise KeyError('Specified loss function not supported')
+
+
+    @classmethod
+    def from_dataset(cls,
+                     dataset,
+                     probe_size=None,
+                     randomize_ang=0,
+                     padding=0,
+                     n_modes=1,
+                     dm_rank=None,
+                     translation_scale=1,
+                     saturation=None,
+                     probe_support_radius=None,
+                     propagation_distance=None,
+                     scattering_mode=None,
+                     oversampling=1,
+                     auto_center=False,
+                     opt_for_fft=False,
+                     fourier_probe=False,
+                     loss='amplitude mse',
+                     units='um',
+                     simulate_probe_translation=False
+                     ):
+
+        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')
+
+        # We include the *extras to make this work even with datasets, like
+        # polarization dependent datasets, that might toss out extra inputs
+        (indices, translations, *extras), 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=opt_for_fft,
+                                                  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(surface_normal)
+
+        # 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(t.abs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
+
+        # For a Fourier space probe
+        if fourier_probe:
+            probe = tools.propagators.far_field(probe)
+
+        probe = t.stack([probe, ] + probe_stack)
+
+        obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
+
+        if dm_rank is not None and dm_rank != 0:
+            if dm_rank > n_modes:
+                raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
+            elif dm_rank == -1:
+                # dm_rank == -1 is defined to mean full-rank
+                dm_rank = n_modes
+
+            Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64)
+            # Start with as close to the identity matrix as possible,
+            # cutting of when we hit the specified maximum rank
+            for i in range(0, dm_rank):
+                Ws[:, i, i] = 1
+        else:
+            # dm_rank == None or dm_rank = 0 triggers a special case where
+            # a standard incoherent multi-mode model is used. This is the
+            # default, because it is so common.
+            # In this case, we define a set of weights which only has one index
+            Ws = t.ones(len(dataset))
+
+        if hasattr(dataset, 'intensities') and dataset.intensities is not None:
+            Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
+                   / t.mean(dataset.intensities))
+
+        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(probe[0].shape, dtype=t.bool)
+            xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
+            xs = xs - np.mean(xs)
+            ys = ys - np.mean(ys)
+            Rs = np.sqrt(xs**2 + ys**2)
+
+            probe_support[Rs < probe_support_radius] = 1
+            probe = probe * probe_support[None, :, :]
+
+        else:
+            probe_support = None
+
+        return cls(wavelength, det_geo, probe_basis, probe, obj,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets=translation_offsets,
+                   weights=Ws, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   fourier_probe=fourier_probe,
+                   oversampling=oversampling,
+                   loss=loss, units=units,
+                   simulate_probe_translation=simulate_probe_translation)
+
+
+    def interaction(self, index, translations, *args):
+
+        # The *args is included so that this can work even when given, say,
+        # a polarized ptycho dataset that might spit out more inputs.
+
+        # Step 1 is to convert the translations for each position into a
+        # value in pixels
+        pix_trans = tools.interactions.translations_to_pixel(
+            self.probe_basis,
+            translations,
+            surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+        # We then add on any recovered translation offset, if they exist
+        if self.translation_offsets is not None:
+            pix_trans += (self.translation_scale *
+                          self.translation_offsets[index])
+
+        # This restricts the basis probes to stay within the probe support
+        basis_prs = self.probe * self.probe_support[..., :, :]
+
+        # For a Fourier-space probe, we take an IFT
+        if self.fourier_probe:
+            basis_prs = tools.propagators.inverse_far_field(basis_prs)
+
+        # Now we construct the probes for each shot from the basis probes
+        if self.weights is not None:
+            Ws = self.weights[index]
+        else:
+            try:
+                Ws = t.ones(len(index)) # I'm positive this introduced a bug
+            except:
+                Ws = 1
+
+        if self.weights is None or len(self.weights[0].shape) == 0:
+            # If a purely stable coherent illumination is defined
+            prs = Ws[..., None, None, None] * basis_prs
+        else:
+            # If a frame-by-frame weight matrix is defined
+            # This takes the dot product of all the weight matrices with
+            # the probes. The output has dimensions of translation, then
+            # coherent mode index, then x,y, and then complex index
+            # Maybe this can be done with a matmul now?
+            prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3)
+
+        if self.simulate_probe_translation:
+            det_pix_trans = tools.interactions.translations_to_pixel(
+                    self.detector_geometry['basis'],
+                    translations,
+                    surface_normal=self.surface_normal)
+            
+            probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
+                                     self.I_phase[None,...] +
+                                     det_pix_trans[:,1,None,None] *
+                                     self.J_phase[None,...]))
+            prs = prs * probe_masks[...,None,:,:]
+
+
+        # Now we actually do the interaction, using the sinc subpixel
+        # translation model as per usual
+        exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
+            prs, self.obj, pix_trans,
+            shift_probe=True, multiple_modes=True)
+        return 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)
+
+
+    # Note: No "loss" function is defined here, because it is added
+    # dynamically during object creation in __init__
+
+    def to(self, *args, **kwargs):
+        super(FancyPtycho, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args, **kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if 'distance' in det_geo:
+            det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
+        if 'basis' in det_geo:
+            det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
+        if 'corner' in det_geo and det_geo['corner'] is not None:
+            det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
+
+        if self.mask is not None:
+            self.mask = self.mask.to(*args, **kwargs)
+
+        if self.simulate_probe_translation:
+            self.I_phase = self.I_phase.to(*args, **kwargs)
+            self.J_phase = self.J_phase.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.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+
+    def sim_to_dataset(self, args_list, calculation_width=None):
+        # 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
+
+        data = []
+        len(indices)
+        if calculation_width is None:
+            calculation_width = len(indices)
+        index_chunks = [indices[i:i + calculation_width]
+                        for i in range(0, len(indices),
+                                       calculation_width)]
+        translation_chunks = [translations[i:i + calculation_width]
+                              for i in range(0, len(indices),
+                                             calculation_width)]
+        
+            
+        # Then we simulate the results
+        data = [self.forward(idx, trans).detach()
+                for idx, trans in zip(index_chunks, translation_chunks)]
+
+        data = t.cat(data, dim=0)
+        # 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=t.float32, device=self.probe.device)
+        if (hasattr(self, 'translation_offsets') and
+            self.translation_offsets is not None):
+            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
+        else:
+            return translations
+
+
+    def get_rhos(self):
+        # If this is the general unified mode model
+        if self.weights.dim() >= 2:
+            Ws = self.weights.detach().cpu().numpy()
+            rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj())
+            return rhos_out
+        # This is the purely incoherent case
+        else:
+            return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
+                            dtype=np.complex64)
+
+    def tidy_probes(self, normalization=1, normalize=False):
+        """Tidies up the probes
+        
+        What we want to do here is use all the information on all the probes
+        to calculate a natural basis for the experiment, and update all the
+        density matrices to operate in that updated basis
+        """
+
+        # First we treat the purely incoherent case
+
+        # I don't love this pattern of using an if statement with a return
+        # to catch this case, but because it's so much simpler than the
+        # unified mode case I think it's appropriate
+        if self.weights.dim() == 1:
+            probe = self.probe.detach().cpu().numpy()
+            ortho_probes = analysis.orthogonalize_probes(probe)
+            self.probe.data = t.as_tensor(
+                ortho_probes,
+                device=self.probe.device,
+                dtype=self.probe.dtype)
+            return
+
+        # This is for the unified mode case
+
+        # Note to future: We could probably do this more cleanly with an
+        # SVD directly on the Ws matrix, instead of an eigendecomposition
+        # of the rho matrix.
+
+        rhos = self.get_rhos()
+        overall_rho = np.mean(rhos, axis=0)
+        probe = self.probe.detach().cpu().numpy()
+        ortho_probes, A = analysis.orthogonalize_probes(
+            probe, density_matrix=overall_rho,
+            keep_transform=True, normalize=normalize)
+        Aconj = A.conj()
+        Atrans = np.transpose(A)
+        new_rhos = np.matmul(Atrans, np.matmul(rhos, Aconj))
+
+        new_rhos /= normalization
+        ortho_probes *= np.sqrt(normalization)
+
+        dm_rank = self.weights.shape[1]
+
+        new_Ws = []
+        for rho in new_rhos:
+            # These are returned from smallest to largest - we want to keep
+            # the largest ones
+            w, v = sla.eigh(rho)
+            w = w[::-1][:dm_rank]
+            v = v[:, ::-1][:, :dm_rank]
+            # For situations where the rank of the density matrix is not
+            # full in reality, but we keep more modes around than needed,
+            # some ws can go negative due to numerical error! This is
+            # extremely rare, but comon enough to cause crashes occasionally
+            # when there are thousands of individual matrices to transform
+            # every time this is called.
+            w = np.maximum(w, 0)
+
+            new_Ws.append(np.dot(np.diag(np.sqrt(w)), v.transpose()))
+
+        new_Ws = np.array(new_Ws)
+
+        self.weights.data = t.as_tensor(
+            new_Ws, dtype=self.weights.dtype, device=self.weights.device)
+
+        self.probe.data = t.as_tensor(
+            ortho_probes, device=self.probe.device, dtype=self.probe.dtype)
+
+
+    def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs):
+        def get_probes(idx):
+            basis_prs = self.probe * self.probe_support[..., :, :]
+            prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
+                        axis=-4)
+            ortho_probes = analysis.orthogonalize_probes(prs)
+
+            if mode.lower() == 'amplitude':
+                return np.abs(ortho_probes.detach().cpu().numpy())
+            if mode.lower() == 'root_sum_intensity':
+                return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,
+                              axis=0)
+            if mode.lower() == 'phase':
+                return np.angle(ortho_probes.detach().cpu().numpy())
+
+        probe_matrix = np.zeros([self.probe.shape[0]]*2,
+                                dtype=np.complex64)
+        np_probes = self.probe.detach().cpu().numpy()
+        for i in range(probe_matrix.shape[0]):
+            for j in range(probe_matrix.shape[0]):
+                probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
+
+        weights = self.weights.detach().cpu().numpy()
+
+        probe_intensities = np.sum(np.tensordot(weights, probe_matrix, axes=1)
+                                   * weights.conj(), axis=2)
+
+        # Imaginary part is already essentially zero up to rounding error
+        probe_intensities = np.real(probe_intensities)
+
+        values = np.sum(probe_intensities, axis=1)
+        if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
+            cmap = 'viridis'
+        else:
+            cmap = 'twilight'
+
+        p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity', cmap=cmap, **kwargs),
+
+        
+    plot_list = [
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('Basis Probe Fourier Space Amplitudes',
+         lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Basis Probe Fourier Space Phases',
+         lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Basis Probe Real Space Amplitudes',
+         lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Basis Probe Real Space Phases',
+         lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
+         lambda self: len(self.weights.shape) >= 2),
+        ('% Power in Top Mode (only accurate after tidy_probes)',
+         lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
+         lambda self: len(self.weights.shape) >= 2),
+        ('Object Amplitude',
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
+#    def plot_errors(self, dataset):
+        
+    
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        probe = self.probe.detach().cpu().numpy()
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+
+        return {'basis': basis, 'translation': translations,
+                'probe': probe, 'obj': obj,
+                'background': background,
+                'weights': weights}

src/CDTools/models/multimode_rpi.py

@@ -0,0 +1,359 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools.interactions import RPI_interaction
+from CDTools.tools import initializers
+from scipy.ndimage.morphology import binary_dilation
+import numpy as np
+from copy import copy
+
+__all__ = ['MultimodeRPI']
+
+
+
+__all__ = ['RPI']
+
+class MultimodeRPI(CDIModel):
+
+    @property
+    def obj(self):
+        return t.complex(self.obj_real, self.obj_imag)
+
+    @property
+    def weights(self):
+        ws = t.complex(self.weights_real, self.weights_imag) 
+        return ws / 10# / self.obj_real.size().numel()
+    
+    def __init__(self, wavelength, detector_geometry, probe_basis,
+                 probe, obj_guess, detector_slice=None,
+                 background=None, mask=None, saturation=None,
+                 obj_support=None, oversampling=1, weight_matrix=False):
+
+        super(MultimodeRPI, 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.probe_basis = t.tensor(probe_basis)
+
+        scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
+                                 probe.shape[-2]/obj_guess.shape[-2]])
+        self.obj_basis = self.probe_basis * scale_factor
+        self.detector_slice = detector_slice
+
+        # Maybe something to include in a bit
+        # self.surface_normal = t.tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+            
+        self.probe = t.tensor(probe, dtype=t.complex64)
+
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+        
+        self.obj_real = t.nn.Parameter(obj_guess.real)
+        self.obj_imag = t.nn.Parameter(obj_guess.imag)
+        
+        self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel())
+        self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0]))
+
+        if not weight_matrix:
+            self.weights_real.requires_grad=False
+            self.weights_imag.requires_grad=False
+        
+        # Wait for LBFGS to be updated for complex-valued parameters
+        # 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,
+                    dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+                
+        self.background = t.tensor(background, dtype=t.float32)
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support[None, ...]
+        else:
+            self.obj_support = t.ones_like(self.obj[0, ...])
+
+        self.oversampling = oversampling
+
+
+    @classmethod
+    def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0):
+        raise NotImplementedError()
+        
+        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')
+        # We only need the patterns here, not the inputs associated with them.
+        _, 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=opt_for_fft,
+                                                   oversampling=oversampling)
+
+        if not isinstance(probe,t.Tensor):
+            probe = t.as_tensor(probe)
+        
+        # Potentially need all of this orientation stuff later
+        
+        #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(surface_normal)
+
+
+        if background is None and hasattr(dataset, 'background') \
+           and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        elif background is not None:
+            background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
+
+        det_geo = dataset.detector_geometry
+
+        # If no mask is given, but one exists in the dataset, load it.
+        if mask is None and hasattr(dataset, 'mask') \
+           and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+
+        # Now we initialize the object
+        if obj_size is None:
+            # This is a standard size for a well-matched probe and detector
+            obj_size = (np.array(probe_shape) // 2).astype(int)
+
+        if initialization.lower().strip() == 'random':
+            # I think something to do with the fact that the object is defined
+            # on a coarser grid needs to be accounted for here that is not
+            # accounted for yet
+            scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
+            obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
+        elif initialization.lower().strip() == 'spectral':
+            if background is not None:
+                obj_guess = initializers.RPI_spectral_init(
+                    patterns[0], probe, obj_size, mask=mask,
+                    background=background**2, n_modes=n_modes)
+            else:
+                obj_guess = initializers.RPI_spectral_init(
+                    patterns[0], probe, obj_size, mask=mask,
+                    n_modes=n_modes)
+                
+        else:
+            raise KeyError('Initialization "' + str(initialization) + \
+                           '" invalid - use "spectral" or "random"')
+
+        probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
+        probe_fft = tools.propagators.far_field(probe_intensity)
+        pad0l = (probe.shape[-2] - obj_size[-2])//2
+        pad0r = probe.shape[-2] - obj_size[-2] - pad0l
+        pad1l = (probe.shape[-1] - obj_size[-1])//2
+        pad1r = probe.shape[-1] - obj_size[-1] - pad1l
+        probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
+        probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
+
+        obj_support = probe_lr > t.max(probe_lr) * probe_threshold
+        obj_support = t.as_tensor(binary_dilation(obj_support))
+
+        return cls(wavelength, det_geo, probe_basis,
+                   probe, obj_guess, detector_slice=det_slice,
+                   background=background, mask=mask, saturation=saturation,
+                   obj_support=obj_support, oversampling=oversampling,
+                   weight_matrix=weight_matrix)
+
+
+    def random_init(self, pattern):
+        scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
+        self.obj.data = scale * t.exp(
+            2j * np.pi * t.rand(self.obj.shape)).to(
+                dtype=self.obj.dtype, device=self.obj.device)
+        
+    def spectral_init(self, pattern):
+        if self.background is not None:
+            self.obj.data = initializers.RPI_spectral_init(
+                pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
+                background=self.background**2, n_modes=self.obj.shape[0]).to(
+                    dtype=self.obj.dtype, device=self.obj.device)
+        else:
+            self.obj.data = initializers.RPI_spectral_init(
+                pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
+                n_modes=self.obj.shape[0]).to(
+                    dtype=self.obj.dtype, device=self.obj.device)
+    
+    # Needs work
+    def interaction(self, index, *args):
+        # including *args allows this to work with all sorts of datasets
+        # that might include other information in with the index in their
+        # "input" parameters (such as translations for a ptychography dataset).
+        # This makes it seamless to use such a dataset even though those
+        # extra arguments will not be used.
+
+        
+        all_exit_waves = []
+
+        # Mix the probes with the weight matrix
+        prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3)
+        
+        for i in range(self.probe.shape[0]):
+            pr = prs[i]
+            # Here we have a 3D probe (one single mode)
+            # and a 4D object (multiple modes mixing incoherently)
+            exit_waves = RPI_interaction(pr,
+                                         self.obj_support * self.obj[i])
+            all_exit_waves.append(exit_waves.unsqueeze(0))
+
+        # This creates a bunch of modes generated from all possible combos
+        # of the probe and object modes all strung out along the first index
+
+        output = t.cat(all_exit_waves)
+
+        # If we have multiple indexes input, we unsqueeze and repeat the stack
+        # of wavefields enough times to simulate each requested index. This
+        # seems silly, but it enables (for example) one to do a reconstruction
+        # from a set of diffraction patterns that are all known to be from the
+        # same object.
+        try:
+            # will fail if index has no length, for example when index
+            # is just an int. In this case, we just do nothing instead
+            output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
+        except TypeError:
+            pass
+        return output
+
+
+    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):
+        # Here I'm taking advantage of an undocumented feature in the
+        # incoherent_sum measurement function where it will work with
+        # a 4D wavefield array as well as a 5D array.
+        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)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
+
+    def regularizer(self, factors):
+        return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+            + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
+        
+    def to(self, *args, **kwargs):
+        super(MultimodeRPI, 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.probe = self.probe.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.obj_basis = self.obj_basis.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.background = self.background.to(*args, **kwargs)
+        
+        # Maybe include in a bit
+        #self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+    def sim_to_dataset(self, args_list):
+        raise NotImplementedError('No sim to dataset yet, sorry!')
+
+    plot_list = [
+        ('Root Sum Squared Amplitude of all Probes',
+         lambda self, fig: p.plot_amplitude(
+             np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
+             fig=fig, basis=self.probe_basis)),
+        ('Object Amplitudes', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
+                                            basis=self.obj_basis)),
+        ('Object Phases',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig,
+                                        basis=self.obj_basis))
+    ]
+
+
+    def save_results(self, dataset=None, full_obj=False):
+        # dataset is set as a kwarg here because it isn't needed, but the
+        # common pattern is to pass a dataset. This makes it okay if one
+        # continues to use that standard pattern
+        probe_basis = self.probe_basis.detach().cpu().numpy()
+        obj_basis = self.obj_basis.detach().cpu().numpy()
+        probe = self.probe.detach().cpu().numpy()
+        # Provide the option to save out the subdominant objects or
+        # just the dominant one
+        if full_obj:
+            obj = self.obj.detach().cpu().numpy()
+        else:
+            obj = self.obj[0].detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        
+        return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
+                'probe': probe,'obj': obj,
+                'background': background}
+

src/CDTools/models/multislice_2d_ptycho.py

@@ -0,0 +1,637 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import analysis, image_processing
+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
+from functools import reduce
+
+__all__ = ['Multislice2DPtycho']
+
+
+
+class Multislice2DPtycho(CDIModel):
+
+    @property
+    def probe(self):
+        return t.complex(self.probe_real, self.probe_imag)
+
+    @property
+    def obj(self):
+        return t.complex(self.obj_real, self.obj_imag)
+
+    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,
+                 oversampling=1,
+                 bandlimit=None,
+                 subpixel=True,
+                 exponentiate_obj=True,
+                 fourier_probe=False,
+                 prevent_aliasing=True,
+                 phase_only=False,
+                 units='um',
+                 ):
+
+        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 = copy(detector_slice)
+        self.surface_normal = t.tensor(surface_normal)
+
+        self.saturation = saturation
+        self.subpixel = subpixel
+        self.exponentiate_obj = exponentiate_obj
+        self.fourier_probe = fourier_probe
+        self.units = units
+        self.phase_only = phase_only
+        self.prevent_aliasing = prevent_aliasing
+
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+
+        # 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(t.abs(probe_guess[0]))
+        else:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess))
+
+        pg = probe_guess / self.probe_norm
+        self.probe_real = t.nn.Parameter(pg.real)
+        self.probe_imag = t.nn.Parameter(pg.imag)
+
+        self.obj_real = t.nn.Parameter(obj_guess.real)
+        self.obj_imag = t.nn.Parameter(obj_guess.imag)
+
+        #self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
+        #                            / self.probe_norm)
+        #self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
+
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(
+                    self.probe[0][self.detector_slice].shape, dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+
+        self.background = t.nn.Parameter(background)
+
+        if weights is None:
+            self.weights = None
+        else:
+            # We now need to distinguish between real-valued per-image
+            # weights and complex-valued per-mode weight matrices
+            if len(weights.shape) == 1:
+                # This is if it's just a list of numbers
+                self.weights = t.nn.Parameter(t.tensor(weights,
+                                                       dtype=t.float32))
+            else:
+                # Now this is a matrix of weights, so it needs to be complex
+                self.weights = t.nn.Parameter(t.tensor(weights,
+                                                       dtype=t.complex64))
+
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            t_o = t.tensor(translation_offsets, dtype=t.float32)
+            t_o = t_o / translation_scale
+            self.translation_offsets = t.nn.Parameter(t_o)
+
+        self.translation_scale = translation_scale
+
+        if probe_support is not None:
+            self.probe_support = t.tensor(probe_support, dtype=t.bool)
+        else:
+            self.probe_support = None
+
+        self.oversampling = oversampling
+
+        spacing = np.linalg.norm(self.probe_basis, axis=0)
+        shape = np.array(self.probe.shape[1:])
+        if prevent_aliasing:
+            shape *= 2
+            spacing /= 2
+
+        self.bandlimit = bandlimit
+
+        self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, self.bandlimit)
+
+
+    @classmethod
+    def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None):
+
+        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=100)
+
+        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
+        probe = tools.initializers.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling)
+        #probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(t.abs(probe))
+        if n_modes >=2:
+            probe_stack = list(0.01*tools.initializers.generate_subdominant_modes(probe,n_modes-1,circular=False))
+            #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)
+        else:
+            probe = t.unsqueeze(probe,0)
+
+        # For a Fourier space probe
+        if fourier_probe:
+            probe = tools.propagators.far_field(probe)
+
+        # Consider a different start
+        if exponentiate_obj:
+            obj = t.zeros(obj_size, dtype=t.complex64)
+        else:
+            obj = t.exp(1j*t.zeros(obj_size))
+        # If we will use a separate object per slice
+        if not replicate_slice:
+            obj = t.stack([obj]*nz)
+            
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
+
+        if dm_rank is not None and dm_rank != 0:
+            if dm_rank > n_modes:
+                raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
+            elif dm_rank == -1:
+                # dm_rank == -1 is defined to mean full-rank
+                dm_rank = n_modes
+            Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
+            # Start with as close to the identity matrix as possible,
+            # cutting of when we hit the specified maximum rank
+            for i in range(0,dm_rank):
+                Ws[:,i,i,0] = 1
+        else:
+            # dm_rank == None or dm_rank = 0 triggers a special case where
+            # a standard incoherent multi-mode model is used. This is the
+            # default, because it is so common.
+            # In this case, we define a set of weights which only has one index
+            Ws = 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(probe[0].shape, dtype=t.bool)
+            xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
+            xs = xs - np.mean(xs)
+            ys = ys - np.mean(ys)
+            Rs = np.sqrt(xs**2 + ys**2)
+        
+            probe_support[Rs<probe_support_radius] = 1
+            probe = probe * probe_support[None,:,:]
+
+        else:
+            probe_support = None
+ 
+        probe = probe
+        
+        return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   probe_support=probe_support,
+                   min_translation=min_translation,
+                   translation_offsets = translation_offsets,
+                   weights=Ws, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   oversampling=oversampling,
+                   bandlimit=bandlimit,
+                   subpixel=subpixel,
+                   exponentiate_obj=exponentiate_obj,
+                   units=units, fourier_probe=fourier_probe,
+                   phase_only=phase_only,
+                   prevent_aliasing=prevent_aliasing)
+                   
+    
+    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]
+        
+            # This restricts the basis probes to stay within the probe support
+        if self.probe_support is not None:
+            basis_prs = self.probe * self.probe_support[...,:,:]
+        else:
+            basis_prs = self.probe 
+
+        # For a Fourier-space probe, we take an IFT
+        if self.fourier_probe:
+            basis_prs = tools.propagators.inverse_far_field(basis_prs)
+
+        if self.prevent_aliasing:
+            #pix_trans = pix_trans * 2
+            basis_prs = image_processing.fourier_upsample(basis_prs)
+
+        # Now we construct the probes for each shot from the basis probes
+        Ws = self.weights[index]
+
+        if len(self.weights[0].shape) == 0:
+            # If a purely stable coherent illumination is defined
+            # No cmult because Ws is real in this case
+            prs = Ws[...,None,None,None] * basis_prs
+        else:
+            # If a frame-by-frame weight matrix is defined
+            # This takes the dot product of all the weight matrices with
+            # the probes. The output has dimensions of translation, then
+            # coherent mode index, then x,y, and then complex index
+            prs = t.sum(Ws[...,None,None] * basis_prs, axis=-3)
+            
+
+        if self.exponentiate_obj:
+            if self.phase_only:
+                obj = t.exp(1j*self.obj.real)
+            else:
+                obj = t.exp(1j*self.obj)
+        else:
+            obj = self.obj
+
+        #if self.prevent_aliasing:
+        #    obj = image_processing.fourier_upsample(obj)
+
+        exit_waves = self.probe_norm * prs
+        for i in range(self.nz):
+            # If only one object slice
+            if self.obj.dim() == 2:
+                if i == 0 and self.subpixel:
+                    # We only need to apply the subpixel shift to the first
+                    # slice, because it shifts the probe
+                    exit_waves = tools.interactions.ptycho_2D_sinc(
+                        exit_waves, obj, pix_trans,
+                        shift_probe=True, multiple_modes=True)
+                else:
+                    exit_waves = tools.interactions.ptycho_2D_round(
+                        exit_waves, obj, pix_trans,
+                        multiple_modes=True,upsample_obj=self.prevent_aliasing)
+                                        
+                    
+            elif self.obj.dim() == 3:
+                # If separate slices
+                if i == 0 and self.subpixel:
+                    exit_waves = tools.interactions.ptycho_2D_sinc(
+                        exit_waves, obj[i], pix_trans,
+                        shift_probe=True, multiple_modes=True)
+                else:
+                    exit_waves = tools.interactions.ptycho_2D_round(
+                        exit_waves, obj[i], pix_trans,
+                        multiple_modes=True, upsample_obj=self.prevent_aliasing)
+
+            #if self.iteration_count >= 1:
+            #    plt.imshow(t.abs(tools.propagators.far_field(
+            #        exit_waves[0,0].detach()).cpu()))
+            #    plt.colorbar()
+            #    plt.show()
+            if i < self.nz-1: #on all but the last iteration
+                exit_waves = tools.propagators.near_field(
+                    exit_waves,self.as_prop)
+            #plt.imshow(t.abs(exit_waves[0,0].detach().cpu()))
+            #plt.show()
+
+        return exit_waves
+    
+        
+    def forward_propagator(self, wavefields):
+        if self.prevent_aliasing:
+            left = [self.probe.shape[-2]//2,self.probe.shape[-1]//2]
+            right = [self.probe.shape[-2]//2+self.probe.shape[-2],
+                     self.probe.shape[-1]//2+self.probe.shape[-1]]
+            
+            return tools.propagators.far_field(wavefields)[...,left[0]:right[0],
+                                                           left[1]:right[1]]
+        else:
+            return tools.propagators.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)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5)
+
+    
+    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)
+        if self.probe_support is not None:
+            self.probe_support = self.probe_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.real.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
+
+    
+    def get_rhos(self):
+        # If this is the general unified mode model
+        if self.weights.dim() >= 2:
+            Ws = self.weights.detach().cpu().numpy()
+            rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
+            return rhos_out
+        # This is the purely incoherent case
+        else:
+            return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
+                            dtype=np.complex64)
+
+        
+    def tidy_probes(self, normalization=1, normalize=False):
+        """Tidies up the probes
+        
+        What we want to do here is use all the information on all the probes
+        to calculate a natural basis for the experiment, and update all the
+        density matrices to operate in that updated basis
+        
+        """
+        
+        # First we treat the purely incoherent case
+        
+        # I don't love this pattern of using an if statement with a return
+        # to catch this case, but because it's so much simpler than the
+        # unified mode case I think it's appropriate
+        if self.weights.dim() == 1:
+            probe = self.probe.detach().cpu().numpy()
+            ortho_probes = analysis.orthogonalize_probes(probe)
+            self.probe.data = t.as_tensor(ortho_probes,
+                        device=self.probe.device,dtype=self.probe.dtype)
+            return
+
+        # This is for the unified mode case
+        
+        # Note to future: We could probably do this more cleanly with an
+        # SVD directly on the Ws matrix, instead of an eigendecomposition
+        # of the rho matrix.
+        
+        rhos = self.get_rhos()
+        overall_rho = np.mean(rhos,axis=0)
+        probe = self.probe.detach().cpu().numpy()
+        ortho_probes, A = analysis.orthogonalize_probes(probe,
+                                        density_matrix=overall_rho,
+                                        keep_transform=True,
+                                        normalize=normalize)
+        Aconj = A.conj()
+        Atrans = np.transpose(A)
+        new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
+
+        new_rhos /= normalization
+        ortho_probes *= np.sqrt(normalization)
+
+        dm_rank = self.weights.shape[1]
+        
+        new_Ws = []
+        for rho in new_rhos:
+            # These are returned from smallest to largest - we want to keep
+            # the largest ones
+            w,v = sla.eigh(rho)
+            w = w[::-1][:dm_rank]
+            v = v[:,::-1][:,:dm_rank]
+            # For situations where the rank of the density matrix is not
+            # full in reality, but we keep more modes around than needed,
+            # some ws can go negative due to numerical error! This is
+            # extremely rare, but comon enough to cause crashes occasionally
+            # when there are thousands of individual matrices to transform
+            # every time this is called.
+            w = np.maximum(w,0)
+            
+            new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
+            
+        new_Ws = np.array(new_Ws)
+
+        self.weights.data = t.as_tensor(new_Ws,
+            dtype=self.weights.dtype,device=self.weights.device)
+        
+        self.probe.data = t.as_tensor(ortho_probes,
+            device=self.probe.device,dtype=self.probe.dtype)
+        
+
+    # Needs to be updated to allow for plotting to an existing figure
+    plot_list = [ 
+        ('Probe Fourier Space Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Probe Fourier Space Phase',
+         lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Probe Real Space Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Probe Real Space Phase',
+         lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()),axis=0), fig=fig),
+         lambda self: len(self.weights.shape) >=2),
+        ('% Power in Top Mode (only accurate after tidy_probes)',
+         lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig,units=self.units),
+         lambda self: len(self.weights.shape) >=2),
+        ('Slice by Slice Real Part of T', 
+         lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
+         lambda self: self.exponentiate_obj),
+        ('Slice by Slice Imaginary Part of T',
+         lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
+         lambda self: self.exponentiate_obj),
+        ('Integrated Real Part of T', 
+         lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
+         lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
+        ('Integrated Imaginary Part of T',
+         lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
+         lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
+        ('Slice by Slice Amplitude of Object Function', 
+         lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
+         lambda self: not self.exponentiate_obj),
+        ('Slice by Slice Phase of Object Function',
+         lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
+         lambda self: not self.exponentiate_obj),
+        ('Amplitude of Stacked Object Function',
+         lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
+         lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3),
+        ('Phase of Stacked Object Function',
+         lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
+         lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+        ('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()
+        if self.fourier_probe:
+            probe = tools.propagators.inverse_far_field(self.probe)
+        else:
+            probe = self.probe
+            
+        probe = probe.detach().cpu().numpy()
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        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}

src/CDTools/models/polarized_fancy_ptycho.py

@@ -0,0 +1,556 @@
+import torch as t
+from CDTools.models import CDIModel, FancyPtycho
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import plotting as p
+# from CDTools.tools import polarized_plotting as pp
+from CDTools.tools import analysis
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from scipy import linalg as sla
+from copy import copy
+from CDTools.tools import polarization
+
+__all__ = ['PolarizedFancyPtycho']
+
+class PolarizedFancyPtycho(FancyPtycho):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_guess, polarizer, analyzer,
+                 detector_slice=None,
+                 surface_normal=np.array([0.,0.,1.]),
+                 min_translation = t.Tensor([0,0]),
+                 background = None, translation_offsets=None,
+                 polarizer_offsets=None, analyzer_offsets=None,
+                 polarizer_scale=1, analyzer_scale=1,  mask=None,
+                 weights = None, translation_scale = 1, saturation=None,
+                 probe_support = None, oversampling=1,
+                 loss='amplitude mse',units='um'):
+
+        super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_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 = weights, translation_scale = 1, saturation=None,
+                 probe_support = None, oversampling=1,
+                 loss='amplitude mse',units='um')
+
+        if polarizer_offsets is None:
+            self.polarizer_offsets = None
+        else:
+            self.polarizer_offsets = t.nn.Parameter(t.tensor(polarizer_offsets).to(dtype=t.float32)) / polarizer_scale
+
+        if analyzer_offsets is None:
+            self.analyzer_offsets = None
+        else:
+            self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale
+        
+        self.polarizer = polarizer
+        self.analyzer = analyzer
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        if probe_guess.dim() > 4:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
+        else:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess))
+
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+
+    @classmethod
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True):
+
+        # When using this method, remember to pass through the inputs
+        model = FancyPtycho.from_dataset(
+            dataset,
+            probe_size=probe_size,
+            randomize_ang=randomize_ang,
+            padding=padding,
+            n_modes=n_modes,
+            dm_rank=dm_rank,
+            translation_scale=translation_scale,
+            saturation=saturation,
+            probe_support_radius=probe_support_radius,
+            propagation_distance=propagation_distance,
+            scattering_mode=scattering_mode,
+            oversampling=oversampling,
+            auto_center=auto_center,
+            opt_for_fft=opt_for_fft,
+            loss=loss,
+            units=units)
+
+
+        # Mutate the class to its subclass
+        model.__class__ = cls
+
+        if left_polarized:
+            x = 1j
+        else:
+            x = -1j
+
+        probe = model.probe.detach()
+        probe = t.cat((probe, probe * x), dim=-3)
+        probe_max = t.max(t.abs(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)
+        #print('probe', type(probe), probe.shape)
+        model.probe.data = probe
+        #print(model.probe.shape)
+        # obj = t.stack((model.obj.data, model.obj.data), dim=-3)
+        # model.obj.data = t.stack((obj.data, obj.data), dim=-4)
+        # obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+        obj = model.obj.detach()
+        # Abe - Probably something identity matrix-like would be a better
+        # initialization (e.g. ((obj,0*obj),(0*obj,obj))
+        obj = t.stack((obj, obj), dim=-3)
+        obj = t.stack((obj, obj), dim=-4)
+        #print('object', type(obj), obj.shape)
+        model.obj.data = obj
+        #print('polarized fancy ptycho from datset obj')
+        a = obj.detach()
+        #plt.imshow(np.real(a[0, 0, :, :]))
+        #plt.figure()
+        #plt.imshow(np.real(a[0, 1, :, :]))
+        #plt.show()
+
+        # tensor vs tensor.data
+        return model
+
+    polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
+
+    @classmethod
+    def from_dataset2(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'):
+
+        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')
+
+        # We include the *extras to make this work even with datasets, like
+        # polarization dependent datasets, that might toss out extra inputs
+        (indices, translations, polarizer, analyzer), 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
+
+        if left_polarized:
+            x = 1j
+        else:
+            x = -1j
+
+
+        # 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=opt_for_fft,
+                                                  oversampling=oversampling)
+
+        probe_shape = t.stack((2, probe_shape), dim=-3)
+        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(surface_normal)
+
+        # 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(t.abs(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)
+        probe_x, probe_y = probe, probe * x
+        probe = t.stact((probe_x, probe_y), dim=-3)
+
+        a = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+        b = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+        c = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+        d = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+
+        ab = t.stack((a, b), dim=-3)
+        cd = t.stack((c, d), dim=-3)
+        obj = t.stack((ab, cd), dim=-4)
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
+
+        if dm_rank is not None and dm_rank != 0:
+            if dm_rank > n_modes:
+                raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
+            elif dm_rank == -1:
+                # dm_rank == -1 is defined to mean full-rank
+                dm_rank = n_modes
+
+            Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64)
+            # Start with as close to the identity matrix as possible,
+            # cutting of when we hit the specified maximum rank
+            for i in range(0, dm_rank):
+                Ws[:, i, i] = 1
+        else:
+            # dm_rank == None or dm_rank = 0 triggers a special case where
+            # a standard incoherent multi-mode model is used. This is the
+            # default, because it is so common.
+            # In this case, we define a set of weights which only has one index
+            Ws = 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(probe[0].shape, dtype=t.bool)
+            xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
+            xs = xs - np.mean(xs)
+            ys = ys - np.mean(ys)
+            Rs = np.sqrt(xs**2 + ys**2)
+
+            probe_support[Rs < probe_support_radius] = 1
+            probe = probe * probe_support[None, :, :]
+
+        else:
+            probe_support = None
+
+        return cls(wavelength, det_geo, probe_basis, probe, obj,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets=translation_offsets,
+                   weights=Ws, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   oversampling=oversampling,
+                   loss=loss, units=units)
+
+
+
+    def interaction(self, index, translations, polarizer, analyzer, test=False):
+
+        # Step 1 is to convert the translations for each position into a
+        # value in pixels
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations,
+                                                             surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+        # We then add on any recovered translation offset, if they exist
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+
+
+        # This restricts the basis probes to stay within the probe support
+        basis_prs = self.probe * self.probe_support[...,:,:] # This makes no sense
+        # self.probe is an Nx2xXxY stach of probes
+
+        # Now we construct the probes for each shot from the basis probes
+        Ws = self.weights[index]
+        if len(self.weights[0].shape) == 0:
+            # If a purely stable coherent illumination is defined
+            # Ws is a tensor of length M, M is the number of frames to be processed
+            prs = Ws[...,None,None,None,None] * basis_prs
+        else:
+            raise NotImplementedError('Unstable Modes not Implemented for polarized light')
+
+        pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
+        exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
+            pol_probes, self.obj, pix_trans,
+            shift_probe=True, multiple_modes=True, polarized=True)
+        # We're losing some efficiency here, because we only need to keep
+        # around the scalar wavefield after analyzing the waves.
+        # But I think it's not a huge issue - Abe
+        analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
+        return analyzed_exit_waves
+
+
+    def vectorial_wavefields(wavefields, func, *args, **kwargs):
+        wavefields_x = wavefields[..., 0, :, :, :]
+        wavefields_y = wavefields[..., 1, :, :, :]
+        out_x = func(wavefields_x, *args, **kwargs)
+        out_y = func(wavefields_y, *args, **kwargs)
+        out = t.stack((out_x, out_y), dim=-4)
+        return out[..., None, :, :]
+
+    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):
+        wavefields_x = wavefields[..., 0, :, :]
+        wavefields_y = wavefields[..., 1, :, :]
+        out_x = tools.measurements.quadratic_background(wavefields_x,
+                            self.background,
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+        # now, set bckgr to 0 since t shouldn't be calculated twice
+        out_y = tools.measurements.quadratic_background(wavefields_y,
+                            0,
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+
+        return out_x + out_y
+
+
+    # Note: No "loss" function is defined here, because it is added
+    # dynamically during object creation in __init__
+
+    def to(self, *args, **kwargs):
+        super(PolarizedFancyPtycho, self).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=t.float32,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
+
+
+
+    def get_rhos(self):
+        # If this is the general unified mode model
+        if self.weights.dim() >= 2:
+            Ws = self.weights.detach().cpu().numpy()
+            rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
+            return rhos_out
+        # This is the purely incoherent case
+        else:
+            return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
+                            dtype=np.complex64)
+
+    def tidy_probes(self, normalization=1, normalize=False):
+        """Tidies up the probes
+
+        What we want to do here is use all the information on all the probes
+        to calculate a natural basis for the experiment, and update all the
+        density matrices to operate in that updated basis
+
+        """
+
+        # First we treat the purely incoherent case
+
+        # I don't love this pattern of using an if statement with a return
+        # to catch this case, but because it's so much simpler than the
+        # unified mode case I think it's appropriate
+        if self.weights.dim() == 1:
+            probe = self.probe.detach().cpu().numpy()
+            ortho_probes = analysis.orthogonalize_probes(probe)
+            self.probe.data = t.as_tensor(ortho_probes,
+                    device=self.probe.device,dtype=self.probe.dtype)
+            return
+
+        # This is for the unified mode case
+
+        # Note to future: We could probably do this more cleanly with an
+        # SVD directly on the Ws matrix, instead of an eigendecomposition
+        # of the rho matrix.
+
+        rhos = self.get_rhos()
+        overall_rho = np.mean(rhos,axis=0)
+        probe = self.probe.detach().cpu().numpy()
+        ortho_probes, A = analysis.orthogonalize_probes(probe,
+                                        density_matrix=overall_rho,
+                                        keep_transform=True,
+                                        normalize=normalize)
+        Aconj = A.conj()
+        Atrans = np.transpose(A)
+        new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
+
+        new_rhos /= normalization
+        ortho_probes *= np.sqrt(normalization)
+
+        dm_rank = self.weights.shape[1]
+
+        new_Ws = []
+        for rho in new_rhos:
+            # These are returned from smallest to largest - we want to keep
+            # the largest ones
+            w,v = sla.eigh(rho)
+            w = w[::-1][:dm_rank]
+            v = v[:,::-1][:,:dm_rank]
+            # For situations where the rank of the density matrix is not
+            # full in reality, but we keep more modes around than needed,
+            # some ws can go negative due to numerical error! This is
+            # extremely rare, but comon enough to cause crashes occasionally
+            # when there are thousands of individual matrices to transform
+            # every time this is called.
+            w = np.maximum(w,0)
+
+            new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
+
+        new_Ws = np.array(new_Ws)
+
+        self.weights.data = t.as_tensor(new_Ws,
+            dtype=self.weights.dtype,device=self.weights.device)
+
+        self.probe.data = t.as_tensor(ortho_probes,
+            device=self.probe.device,dtype=self.probe.dtype)
+
+
+    def plot_wavefront_variation(self, dataset,fig=None,mode='amplitude',**kwargs):
+        def get_probes(idx):
+            basis_prs = self.probe * self.probe_support[...,:,:]
+            prs = t.sum(self.weights[idx,:,:,None,None] * basis_prs, axis=-4)
+            ortho_probes = analysis.orthogonalize_probes(prs)
+
+            if mode.lower() == 'amplitude':
+                return np.abs(ortho_probes.detach().cpu().numpy())
+            if mode.lower() == 'root_sum_intensity':
+                return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,axis=0)
+            if mode.lower() == 'phase':
+                return np.angle(ortho_probes.detach().cpu().numpy())
+
+        probe_matrix = np.zeros([self.probe.shape[0]]*2,
+                                dtype=np.complex64)
+        np_probes = self.probe.detach().cpu().numpy()
+        for i in range(probe_matrix.shape[0]):
+            for j in range(probe_matrix.shape[0]):
+                probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
+
+
+        weights = self.weights.detach().cpu().numpy()
+
+        probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
+                                   weights.conj(),axis=2)
+
+        # Imaginary part is already essentially zero up to rounding error
+        probe_intensities = np.real(probe_intensities)
+
+        values = np.sum(probe_intensities,axis=1)
+        if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
+            cmap = 'viridis'
+        else:
+            cmap = 'twilight'
+
+        p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity',cmap=cmap,**kwargs),
+
+
+    plot_list = [
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('',
+         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'),
+         lambda self: len(self.weights.shape) >= 2),
+        ('Basis Probe Amplitudes (scroll to view modes)',
+         lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Basis Probe Phases (scroll to view modes)',
+         lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
+         lambda self: len(self.weights.shape) >= 2),
+        ('% Power in Top Mode (only accurate after tidy_probes)',
+         lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
+         lambda self: len(self.weights.shape) >= 2),
+        ('Object Amplitude',
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+        ('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 = self.probe.detach().cpu().numpy()
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'obj':obj,
+                'background':background,
+                'weights':weights}

src/CDTools/models/rpi.py

@@ -0,0 +1,437 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools.interactions import RPI_interaction
+from CDTools.tools import initializers
+from scipy.ndimage.morphology import binary_dilation
+import numpy as np
+from copy import copy
+import time
+
+__all__ = ['RPI']
+
+#
+# This model has to work a bit differently from a ptychography model
+# because a typical RPI dataset will have lots of images, each of which
+# can be reconstructed on it's own. I can see a few ways to approach this:
+#
+# 1) Enfore 1 image per dataset as a restriction for internal consistency
+# of the "model" idea
+#
+# 2) Override some of the base functions of the CDIModel to accept optional
+# parameters that make it work on larger datasets
+#
+# I think #1 is basically untenable, because it would require creating
+# a new dataset and model for each reconstruction - which means instantiating
+# tons of stuff and storing all sorts of excess information for each frame,
+# when that could easily be reused for other frames. As a result, I think I
+# will try the following pattern:
+#
+# 1) A model will contain a single object "guess" at all times
+# 2) Constructing a model from a data will automatically instantiate the object guess from the first diffraction pattern in the dataset
+# 3) All the optimization functions will get an additional argument for the index of the diffraction pattern to reconstruct. This could either be handled by editing the CDIModel base class to pass through some kwargs, or by overriding all the optimization functions explicitly.
+# 4) A few convenience functions can be written to re-initialize the object array from any image / index in the dataset
+# 5) A new function can be written to reconstruct the entire dataset by running through each pattern one at a time.
+#
+# The advantage of this approach is that I can start by writing a class that
+# will only reconstruct the first diffraction pattern from a dataset and then
+# extend it.
+#
+# Final note: It is worth seeing whether it is possible to include the
+# probe propagation explicitly as a parameter which can be reconstructed
+# via gradient descent.
+#
+#
+
+__all__ = ['RPI']
+
+class RPI(CDIModel):
+
+    @property
+    def obj(self):
+        return t.complex(self.obj_real, self.obj_imag)
+
+    @property
+    def weights(self):
+        ws = t.complex(self.weights_real, self.weights_imag) 
+        return ws / 10# / self.obj_real.size().numel()
+
+
+    
+    def __init__(self, wavelength, detector_geometry, probe_basis,
+                 probe, obj_guess, detector_slice=None,
+                 background=None, mask=None, saturation=None,
+                 obj_support=None, oversampling=1, weight_matrix=False):
+
+        super(RPI, 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.probe_basis = t.tensor(probe_basis)
+
+        scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
+                                 probe.shape[-2]/obj_guess.shape[-2]])
+        self.obj_basis = self.probe_basis * scale_factor
+        self.detector_slice = detector_slice
+
+        # Maybe something to include in a bit
+        # self.surface_normal = t.tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+            
+        self.probe = t.tensor(probe, dtype=t.complex64)
+        
+        if obj_guess.dim() == 2:
+            obj_guess = obj_guess[None, :, :]
+
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+        
+        self.obj_real = t.nn.Parameter(obj_guess.real)
+        self.obj_imag = t.nn.Parameter(obj_guess.imag)
+        
+        self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel())
+        self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0]))
+
+        if not weight_matrix:
+            self.weights_real.requires_grad=False
+            self.weights_imag.requires_grad=False
+        
+        # Wait for LBFGS to be updated for complex-valued parameters
+        # 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,
+                    dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+                
+        self.background = t.tensor(background, dtype=t.float32)
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support[None, ...]
+        else:
+            self.obj_support = t.ones_like(self.obj[0, ...])
+
+        self.oversampling = oversampling
+
+
+    @classmethod
+    def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0):
+        
+        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')
+        # We only need the patterns here, not the inputs associated with them.
+        _, 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=opt_for_fft,
+                                                   oversampling=oversampling)
+
+        if not isinstance(probe,t.Tensor):
+            probe = t.as_tensor(probe)
+        
+        # Potentially need all of this orientation stuff later
+        
+        #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(surface_normal)
+
+
+        if background is None and hasattr(dataset, 'background') \
+           and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        elif background is not None:
+            background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
+
+        det_geo = dataset.detector_geometry
+
+        # If no mask is given, but one exists in the dataset, load it.
+        if mask is None and hasattr(dataset, 'mask') \
+           and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+
+        # Now we initialize the object
+        if obj_size is None:
+            # This is a standard size for a well-matched probe and detector
+            obj_size = (np.array(probe_shape) // 2).astype(int)
+
+        if initialization.lower().strip() == 'random':
+            # I think something to do with the fact that the object is defined
+            # on a coarser grid needs to be accounted for here that is not
+            # accounted for yet
+            scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
+            obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
+        elif initialization.lower().strip() == 'spectral':
+            if background is not None:
+                obj_guess = initializers.RPI_spectral_init(
+                    patterns[0], probe, obj_size, mask=mask,
+                    background=background**2, n_modes=n_modes)
+            else:
+                obj_guess = initializers.RPI_spectral_init(
+                    patterns[0], probe, obj_size, mask=mask,
+                    n_modes=n_modes)
+                
+        else:
+            raise KeyError('Initialization "' + str(initialization) + \
+                           '" invalid - use "spectral" or "random"')
+
+        probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
+        probe_fft = tools.propagators.far_field(probe_intensity)
+        pad0l = (probe.shape[-2] - obj_size[-2])//2
+        pad0r = probe.shape[-2] - obj_size[-2] - pad0l
+        pad1l = (probe.shape[-1] - obj_size[-1])//2
+        pad1r = probe.shape[-1] - obj_size[-1] - pad1l
+        probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
+        probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
+
+        obj_support = probe_lr > t.max(probe_lr) * probe_threshold
+        obj_support = t.as_tensor(binary_dilation(obj_support))
+
+        return cls(wavelength, det_geo, probe_basis,
+                   probe, obj_guess, detector_slice=det_slice,
+                   background=background, mask=mask, saturation=saturation,
+                   obj_support=obj_support, oversampling=oversampling,
+                   weight_matrix=weight_matrix)
+
+    @classmethod
+    def from_calibration(cls, calibration, obj_size=None, n_modes=1, saturation=None):
+        
+        wavelength = calibration['wavelength']
+        probe_basis = t.as_tensor(calibration['basis'])
+        probe = t.as_tensor(calibration['probe'])
+        if 'background' in calibration:
+            background = t.as_tensor(calibration['background'])
+        else:
+            background = t.zeros_like(probe.real)
+        if 'mask' in calibration:
+            mask = t.as_tensor(calibration['mask'])
+        else:
+            mask = t.ones_like(probe.real, dtype=t.bool)
+
+        # Now we initialize the object
+        if obj_size is None:
+            # This is a standard size for a well-matched probe and detector
+            obj_size = (np.array(probe.shape[-2:]) // 2).astype(int)
+
+        obj_guess = t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
+
+        det_geo = {'distance': 1,
+                   'basis': wavelength / probe_basis} # Is this even right?
+                
+        return cls(wavelength, det_geo, probe_basis,
+                   probe, obj_guess,
+                   background=background, mask=mask)
+
+    
+    def random_init(self, pattern):
+        scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
+        self.obj.data = scale * t.exp(
+            2j * np.pi * t.rand(self.obj.shape)).to(
+                dtype=self.obj.dtype, device=self.obj.device)
+        
+    def spectral_init(self, pattern):
+        if self.background is not None:
+            self.obj.data = initializers.RPI_spectral_init(
+                pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
+                background=self.background**2, n_modes=self.obj.shape[0]).to(
+                    dtype=self.obj.dtype, device=self.obj.device)
+        else:
+            self.obj.data = initializers.RPI_spectral_init(
+                pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
+                n_modes=self.obj.shape[0]).to(
+                    dtype=self.obj.dtype, device=self.obj.device)
+    
+    # Needs work
+    def interaction(self, index, *args):
+        # including *args allows this to work with all sorts of datasets
+        # that might include other information in with the index in their
+        # "input" parameters (such as translations for a ptychography dataset).
+        # This makes it seamless to use such a dataset even though those
+        # extra arguments will not be used.
+
+        
+        all_exit_waves = []
+
+        # Mix the probes with the weight matrix
+        prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3)
+        
+        for i in range(self.probe.shape[0]):
+            pr = prs[i]
+            # Here we have a 3D probe (one single mode)
+            # and a 4D object (multiple modes mixing incoherently)
+            exit_waves = RPI_interaction(pr,
+                                         self.obj_support * self.obj)
+                
+            all_exit_waves.append(exit_waves)
+
+        # This creates a bunch of modes generated from all possible combos
+        # of the probe and object modes all strung out along the first index
+        output = t.cat(all_exit_waves)
+        
+        # If we have multiple indexes input, we unsqueeze and repeat the stack
+        # of wavefields enough times to simulate each requested index. This
+        # seems silly, but it enables (for example) one to do a reconstruction
+        # from a set of diffraction patterns that are all known to be from the
+        # same object.
+        try:
+            # will fail if index has no length, for example when index
+            # is just an int. In this case, we just do nothing instead
+            output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
+        except TypeError:
+            pass
+        return output
+
+
+    def forward_propagator(self, wavefields):
+        p = tools.propagators.far_field(wavefields)
+        return p
+
+
+    def backward_propagator(self, wavefields):
+        return tools.propagators.inverse_far_field(wavefields)
+
+    
+    def measurement(self, wavefields):
+        # Here I'm taking advantage of an undocumented feature in the
+        # incoherent_sum measurement function where it will work with
+        # a 4D wavefield array as well as a 5D array.
+        m = tools.measurements.quadratic_background(wavefields,
+                            self.background,
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+        return m
+    
+    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 regularizer(self, factors):
+        if self.obj.shape[0] == 1:
+            return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2)
+        else:
+            return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+                + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
+        
+    def to(self, *args, **kwargs):
+        super(RPI, 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.probe = self.probe.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.obj_basis = self.obj_basis.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.background = self.background.to(*args, **kwargs)
+        
+        # Maybe include in a bit
+        #self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+    def sim_to_dataset(self, args_list):
+        raise NotImplementedError('No sim to dataset yet, sorry!')
+
+    plot_list = [
+        ('Root Sum Squared Amplitude of all Probes',
+         lambda self, fig: p.plot_amplitude(
+             np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
+             fig=fig, basis=self.probe_basis)),
+        ('Dominant Object Amplitude', 
+         lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig,
+                                            basis=self.obj_basis)),
+        ('Dominant Object Phase',
+         lambda self, fig: p.plot_phase(self.obj[0], fig=fig,
+                                        basis=self.obj_basis)),
+        ('Subdominant Object Amplitude',
+         lambda self, fig: p.plot_amplitude(self.obj[1], fig=fig,
+                                            basis=self.obj_basis),
+         lambda self: len(self.obj) >=2),
+        ('Subdominant Object Phase',
+         lambda self, fig: p.plot_phase(self.obj[1], fig=fig,
+                                        basis=self.obj_basis),
+         lambda self: len(self.obj) >=2)
+    ]
+
+
+    def save_results(self, dataset=None, full_obj=False):
+        # dataset is set as a kwarg here because it isn't needed, but the
+        # common pattern is to pass a dataset. This makes it okay if one
+        # continues to use that standard pattern
+        probe_basis = self.probe_basis.detach().cpu().numpy()
+        obj_basis = self.obj_basis.detach().cpu().numpy()
+        probe = self.probe.detach().cpu().numpy()
+        # Provide the option to save out the subdominant objects or
+        # just the dominant one
+        if full_obj:
+            obj = self.obj.detach().cpu().numpy()
+        else:
+            obj = self.obj[0].detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        
+        return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
+                'probe': probe,'obj': obj,
+                'background': background}
+

src/CDTools/models/s_matrix_ptycho.py

@@ -0,0 +1,430 @@
+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
+
+__all__ = ['SMatrixPtycho']
+
+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,
+                 probe_planes = 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
+        # Remember that for S-matrix we have several probes for different
+        # planes
+        if probe_guess.dim() > 4:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0,0].to(t.float32)))
+        else:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].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[-3],
+                        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)
+
+        # This maps indices to probe planes to be used. If none, it defaults
+        # to always being plane 0
+        if probe_planes is None:
+            self.probe_planes = None
+        else:
+            self.probe_planes = t.LongTensor(probe_planes)
+
+        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):
+
+        datasets = [dataset]
+        propagation_distances = [propagation_distance]
+
+        # We only return the 0th element because in the general case, the
+        # constructor needs to return a stacked datset in addition to
+        # a model, but for the case of one dataset we only need to return
+        # the model.
+        return cls.from_datasets(datasets, probe_convergence_radius,
+                                 locality_radius=locality_radius,
+                                 probe_size=probe_size,
+                                 randomize_ang=randomize_ang,
+                                 padding=padding,
+                                 n_modes=n_modes,
+                                 translation_scale=translation_scale,
+                                 saturation=saturation,
+                                 propagation_distances=propagation_distances,
+                                 scattering_mode=scattering_mode,
+                                 oversampling=oversampling)[0]
+    
+    # This is for the multi-focal-plane case, where each dataset will correspond
+    # to a different focal plane. The guess propagation distance for each
+    # dataset can be set individually but otherwise the probes are
+    # reconstructed entirely separately. All datasets are assumed to have
+    # the same basic parameters (wavelength, detector geometry, etc) and share
+    # the same origin in the x-y plane.
+    @classmethod
+    def from_datasets(cls, datasets, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distances=None, scattering_mode=None, oversampling=1):
+        
+        wavelength = datasets[0].wavelength
+        det_basis = datasets[0].detector_geometry['basis']
+        det_shape = datasets[0][0][1].shape
+        distance = datasets[0].detector_geometry['distance']
+
+        # 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,
+                                                padding=padding,
+                                                opt_for_fft=False,
+                                                oversampling=oversampling)
+        
+        if propagation_distances is None:
+            propagation_distances = [None] * len(datasets)
+        
+        # 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))
+
+        # always do this on the cpu
+        probe_planes = []
+        translations = []
+        patterns = []
+        probes = []
+        for i, dataset in enumerate(datasets):
+            get_as_args = dataset.get_as_args
+            
+            dataset.get_as(device='cpu')
+            (indices, tx), pats = dataset[:]
+            dataset.get_as(*get_as_args[0],**get_as_args[1])
+            translations.append(tx)
+            probe_planes.extend([i]*tx.shape[0])
+            patterns.append(pats)
+
+            # Finally, initialize the probe and  object using this information
+            if locality_radius != 0:
+                probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], 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_distances[i], oversampling=oversampling)
+
+            # 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)
+            probes.append(probe)
+
+        translations = t.cat(translations)
+        patterns = t.cat(patterns)
+        probes = t.stack(probes)
+
+        
+        if hasattr(datasets[0], 'sample_info') and \
+           datasets[0].sample_info is not None and \
+           'orientation' in datasets[0].sample_info:
+            surface_normal = datasets[0].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(datasets[0].background)
+        else:
+            background = None
+
+
+        
+        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((translations.shape[0],2)) - 0.5)
+
+        weights = t.ones(translations.shape[0])
+        
+        if hasattr(datasets[0], 'mask') and datasets[0].mask is not None:
+            mask = datasets[0].mask.to(t.bool)
+        else:
+            mask = None
+
+        
+        probe_support = t.zeros_like(probes[0,0].to(dtype=t.float32))
+
+        xs, ys = np.mgrid[:probes.shape[-3],:probes.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
+        probes = probes * probe_support[None,None,:,:]
+
+        
+        model = cls(wavelength, det_geo, probe_basis, probes, probe_support,
+                    s_matrix,
+                    detector_slice=det_slice,
+                    surface_normal=surface_normal,
+                    min_translation=min_translation,
+                    translation_offsets = translation_offsets,
+                    probe_planes = probe_planes,
+                    weights=weights, mask=mask, background=background,
+                    translation_scale=translation_scale,
+                    saturation=saturation,
+                    oversampling=oversampling)
+
+        # Now we need to produce a concatenated dataset to be used to
+        # train the model
+        dataset = Ptycho2DDataset(translations, patterns)
+        return model, dataset
+                   
+    
+    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]
+        if self.probe_planes is not None:
+            probes_set = self.probe[self.probe_planes[index]]
+        else:
+            probes_set = self.probe[[0]*translations.shape[0]]
+
+
+        all_exit_waves = []
+        for i in range(probes_set.shape[1]):
+            exit_waves = []
+            
+            for j in range(probes_set.shape[0]):
+                pr = tools.propagators.inverse_far_field(probes_set[j,i] * self.probe_fourier_support)
+                exit_wave = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(pr, self.s_matrix, pix_trans[j], shift_probe=True)
+                exit_waves.append(exit_wave)
+            
+            exit_waves = t.stack(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)
+        if self.probe_planes is not None:
+            self.probe_planes = self.probe_planes.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 = [
+        ('First Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[0,0], fig=fig, basis=self.probe_basis)),
+        ('First Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[0,0], fig=fig, basis=self.probe_basis)),
+        ('Second Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[1,0], fig=fig, basis=self.probe_basis),
+         lambda self: self.probe.shape[0] >=2),
+        ('Second Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[1,0], fig=fig, basis=self.probe_basis),
+         lambda self: self.probe.shape[0] >=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()
+        wavelength = self.wavelength.cpu().numpy()
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'s_matrix':s_matrix,
+                'background':background,
+                'weights':weights, 'wavelength':wavelength}

src/CDTools/models/simple_ptycho.py

@@ -0,0 +1,274 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import plotting as p
+from copy import copy
+from torch.utils import data as torchdata
+from datetime import datetime
+import numpy as np
+
+__all__ = ['SimplePtycho']
+
+class SimplePtycho(CDIModel):
+    """A simple ptychography model for exploring ideas and extensions
+
+
+
+    """
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis, detector_slice,
+                 probe_guess, obj_guess, min_translation = [0,0],
+                 surface_normal=np.array([0.,0.,1.]), mask=None):
+
+        super(SimplePtycho,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 = copy(detector_slice)
+
+        self.surface_normal = t.tensor(surface_normal)
+
+        if mask is None:
+            self.mask = None
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        self.probe_norm = t.max(t.abs(probe_guess))
+
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess)
+
+
+    @classmethod
+    def from_dataset(cls, dataset):
+        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])
+
+        center = tools.image_processing.centroid(t.sum(patterns,dim=0))
+
+        # Then, generate the probe geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center)
+
+        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.])
+
+        # 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)
+
+        # 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).to(dtype=t.complex64)
+        det_geo = dataset.detector_geometry
+
+
+        if hasattr(dataset, 'mask') and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+        else:
+            mask = None
+
+        return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask, surface_normal=surface_normal)
+
+
+    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
+
+        return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
+                                                  self.obj,
+                                                  pix_trans)
+
+
+    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.intensity(wavefields,
+                                            detector_slice=self.detector_slice)
+
+
+    def loss(self, real_data, sim_data, mask=None):
+        return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
+
+
+    def to(self, *args, **kwargs):
+        super(SimplePtycho, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args,**kwargs)
+        # 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.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)
+
+
+
+    plot_list = [
+        ('Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
+        ('Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
+        ('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))
+    ]
+
+
+
+    def save_results(self):
+        probe = self.probe.detach().cpu().numpy()
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        return {'probe':probe,'obj':obj}
+
+
+    def ePIE(self, iterations, dataset, beta = 1.0):
+        """Runs an ePIE reconstruction as described in `Maiden et al. (2017) <https://www.osapublishing.org/optica/abstract.cfm?uri=optica-4-7-736>`_.
+        Optional parameters are:
+
+        :arg ``iterations``: Controls the number of iterations run, defaults to 1.
+        :arg ``beta``: Algorithmic parameter described in Maiden's implementation of rPIE. Defaults to 0.15.
+        :arg ``probe``: Initial probe wavefunction.
+        :arg ``object``: Initial object wavefunction.
+        """
+        probe_shape = self.probe.shape
+
+        if self.mask is not None:
+            mask = self.mask[...,None]
+        else:
+            mask=None
+
+        def probe_update(exit_wave, exit_wave_corrected, probe, object, translation):
+            new_probe = probe + tools.cmath.cmult(beta * tools.cmath.cconj(object[translation])/(self.probe_norm*t.max(tools.cmath.cabssq(object))), exit_wave_corrected-exit_wave)
+            return new_probe
+
+        def object_update(exit_wave, exit_wave_corrected, probe, object, translation):
+            new_object = object.clone()
+            new_object[translation] = object[translation] + tools.cmath.cmult(beta * tools.cmath.cconj(probe)/(self.probe_norm*t.max(tools.cmath.cabssq(probe))), exit_wave_corrected-exit_wave)
+            return new_object
+
+        with t.no_grad():
+            data_loader = torchdata.DataLoader(dataset, shuffle=True)
+
+            for it in range(iterations):
+                loss = []
+                for (i, [translations]), [patterns] in data_loader:
+                    probe = self.probe.data.clone()
+                    object = self.obj.data.clone()
+
+                    exit_wave = self.interaction(i, translations).clone()
+                    # Apply modulus constraint
+                    exit_wave_corrected = exit_wave.clone()
+                    exit_wave_corrected = self.forward_propagator(exit_wave_corrected.clone())
+                    exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(exit_wave_corrected.clone()[self.detector_slice], patterns, mask = mask)
+                    exit_wave_corrected = self.backward_propagator(exit_wave_corrected.clone())
+
+                    # Calculate the section of the object wavefunction to be modified
+                    pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                                         translations)
+                    pix_trans -= self.min_translation
+
+                    pix_trans = t.round(pix_trans).to(dtype=t.int32).detach().cpu().numpy()
+
+                    object_slice = np.s_[pix_trans[0]:
+                                      pix_trans[0]+probe_shape[0],
+                                      pix_trans[1]:
+                                      pix_trans[1]+probe_shape[1]]
+
+                    # Apply probe and object updates
+                    self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
+                    self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
+
+                    # Calculate loss
+                    loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
+
+                yield t.mean(t.tensor(loss)).cpu().numpy()

src/CDTools/models/time_resolved_ptycho_calibration.py

@@ -0,0 +1,504 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools import analysis
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from scipy import linalg as sla
+from copy import copy
+
+#
+# Basic points:
+# Just one probe mode, no need to overcomplicate things.
+# Weights is only a list of numbers, no matrices or anything like that
+# Mandatory probe support in Fourier space
+#
+# When loading from a dataset, we need information on the zone plate geometry
+#
+
+__all__ = ['TimeResolvedPtychoCalibration']
+
+
+class TimeResolvedPtychoCalibration(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_guess,
+                 fourier_times, probe_fourier_support,
+                 times, time_dependence, frame_delays,
+                 detector_slice=None,
+                 surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
+                 min_translation=t.tensor([0, 0], dtype=t.float32),
+                 background=None, translation_offsets=None, mask=None,
+                 weights=None, translation_scale=1, saturation=None,
+                 oversampling=1,
+                 loss='amplitude mse', units='um',
+                 simulate_probe_translation=False):
+
+        super(TimeResolvedPtychoCalibration, self).__init__()
+        self.wavelength = t.tensor(wavelength)
+        self.detector_geometry = copy(detector_geometry)
+        det_geo = self.detector_geometry
+        if 'distance' in det_geo:
+            det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
+        if 'basis' in det_geo:
+            det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
+        if 'corner' in det_geo:
+            det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
+
+        self.min_translation = t.tensor(min_translation)
+
+        self.probe_basis = t.tensor(probe_basis)
+        self.detector_slice = copy(detector_slice)
+        self.surface_normal = t.tensor(surface_normal)
+
+        self.saturation = saturation
+        self.units = units
+
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+
+        self.probe_norm = 1 * t.max(t.abs(probe_guess))        
+
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess)
+
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(
+                    self.probe[self.detector_slice].shape,
+                    dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+
+        self.background = t.nn.Parameter(background)
+
+        if weights is None:
+            self.weights = None
+        else:
+            self.weights = t.nn.Parameter(t.tensor(weights,
+                                                   dtype=t.float32))
+
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            t_o = t.tensor(translation_offsets, dtype=t.float32)
+            t_o = t_o / translation_scale
+            self.translation_offsets = t.nn.Parameter(t_o)
+
+        self.translation_scale = translation_scale
+
+        self.probe_fourier_support = probe_fourier_support
+        self.fourier_times = fourier_times
+        self.times = times
+        self.time_dependence = t.nn.Parameter(time_dependence)
+        self.frame_delays = frame_delays
+
+        self.oversampling = oversampling
+
+        self.simulate_probe_translation = simulate_probe_translation
+        if simulate_probe_translation:
+            Is = t.arange(self.probe.shape[-2], dtype=t.float32)
+            Js = t.arange(self.probe.shape[-1], dtype=t.float32)
+            Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
+            self.I_phase = 2 * np.pi* Is
+            self.J_phase = 2 * np.pi* Js
+            
+        # Here we set the appropriate loss function
+        if (loss.lower().strip() == 'amplitude mse'
+                or loss.lower().strip() == 'amplitude_mse'):
+            self.loss = tools.losses.amplitude_mse
+        elif (loss.lower().strip() == 'poisson nll'
+                or loss.lower().strip() == 'poisson_nll'):
+            self.loss = tools.losses.poisson_nll
+        else:
+            raise KeyError('Specified loss function not supported')
+
+
+    @classmethod
+    def from_dataset(cls, dataset, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', simulate_probe_translation=False):
+
+        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')
+
+        # We include the *extras to make this work even with datasets, like
+        # polarization dependent datasets, that might toss out extra inputs
+        (indices, translations, *extras), 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=opt_for_fft,
+                                                  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(surface_normal)
+
+        # 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
+        probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+
+        probe = tools.propagators.far_field(probe)
+        
+        obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
+        # we define a set of weights which only has one index
+        Ws = t.ones(len(dataset))
+
+        if hasattr(dataset, 'mask') and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+        else:
+            mask = None
+
+        # What do we know about the zp?
+        # delta_r, N, and beamstop ratio. It's probably best, though,
+        # to just read diameter, beamstop_diameter and focal_length directly
+        # because that is the most general even if the optic isn't truly
+        # a zone plate.
+        zp_distance = zp_geometry['focal_length']
+        probe_size = probe_basis * t.as_tensor(probe_shape, dtype=t.float32)
+        pinv_basis = t.tensor(np.linalg.pinv(probe_size).transpose()).to(t.float32)
+        zone_plate_basis = pinv_basis * wavelength * zp_distance
+        zone_plate_steps = t.sum(zone_plate_basis,axis=1)
+        # This may very well mix up x & y and fail on non-square detectors
+        probe_fourier_support = t.zeros(probe.shape, dtype=t.bool)
+        xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
+        xs = zone_plate_steps[1] * (xs - np.mean(xs))
+        ys = zone_plate_steps[0] * (ys - np.mean(ys))
+        Rs = np.sqrt(xs**2 + ys**2)
+
+        distances = np.sqrt(zp_distance**2 + xs**2 + ys**2)
+        times = (distances - t.min(distances)) / 2.99792e8
+
+        # This sets the support of the probe and also restricts the
+        # timing matrix so it only considers times that are actually
+        # in the window defined by the probe fourier support
+        probe_fourier_support[Rs < zp_geometry['diameter']/2] = 1
+        times[Rs > zp_geometry['diameter']/2] = 0
+        times[Rs > zp_geometry['diameter']/2] = t.max(times)        
+        
+        probe_fourier_support[Rs < zp_geometry['beamstop_diameter']/2] = 0
+        times[Rs < zp_geometry['beamstop_diameter']/2] = t.max(times)
+        times[Rs < zp_geometry['beamstop_diameter']/2] = t.min(times)
+
+        fourier_times = times - t.min(times)
+        probe = probe * probe_fourier_support
+
+        # This is now the time axis for the probe's envelope
+        times = t.linspace(0, time_window, n_times+1)
+        time_dependence = t.ones(n_times, dtype=t.complex64)
+        frame_delays = t.linspace(0, t.max(fourier_times) + time_window, n_frames+2)
+        frame_delays -= time_window
+        frame_delays = frame_delays[1:-1]
+        
+        return cls(wavelength, det_geo, probe_basis, probe, obj,
+                   fourier_times, probe_fourier_support,
+                   times, time_dependence, frame_delays,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets=translation_offsets,
+                   weights=Ws, mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   oversampling=oversampling,
+                   loss=loss, units=units,
+                   simulate_probe_translation=simulate_probe_translation)
+
+
+    def interaction(self, index, translations, *args):
+        # The *args is included so that this can work even when given, say,
+        # a polarized ptycho dataset that might spit out more inputs.
+
+        # Step 1 is to convert the translations for each position into a
+        # value in pixels
+        pix_trans = tools.interactions.translations_to_pixel(
+            self.probe_basis,
+            translations,
+            surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+        # We then add on any recovered translation offset, if they exist
+        if self.translation_offsets is not None:
+            pix_trans += (self.translation_scale *
+                          self.translation_offsets[index])
+
+        probes = self.get_probes(space='real')        
+        Ws = self.weights[index]
+
+        # This might not work well
+        prs = Ws[...,None,None,None] * probes
+        #prs = t.sum(Ws[..., None, None, None] * probes, axis=-3)
+        #print(prs.shape)
+
+        if self.simulate_probe_translation:
+            det_pix_trans = tools.interactions.translations_to_pixel(
+                    self.detector_geometry['basis'],
+                    translations,
+                    surface_normal=self.surface_normal)
+            
+            probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
+                                     self.I_phase[None,...] +
+                                     det_pix_trans[:,1,None,None] *
+                                     self.J_phase[None,...]))
+            prs = prs * probe_masks[...,None,:,:]
+
+
+        # Now we actually do the interaction, using the sinc subpixel
+        # translation model as per usual
+        exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
+            prs, self.obj, pix_trans,
+            shift_probe=True, multiple_modes=True)
+        return 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)
+
+
+    # Note: No "loss" function is defined here, because it is added
+    # dynamically during object creation in __init__
+
+    def to(self, *args, **kwargs):
+        super(TimeResolvedPtychoCalibration, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args, **kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if 'distance' in det_geo:
+            det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
+        if 'basis' in det_geo:
+            det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
+        if 'corner' in det_geo:
+            det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
+
+        if self.mask is not None:
+            self.mask = self.mask.to(*args, **kwargs)
+
+        if self.simulate_probe_translation:
+            self.I_phase = self.I_phase.to(*args, **kwargs)
+            self.J_phase = self.J_phase.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.fourier_times = self.fourier_times.to(*args, **kwargs)
+        self.times = self.times.to(*args, **kwargs)
+        self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+
+    def sim_to_dataset(self, args_list, calculation_width=None):
+        # 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
+
+        data = []
+        len(indices)
+        if calculation_width is None:
+            calculation_width = len(indices)
+        index_chunks = [indices[i:i + calculation_width]
+                        for i in range(0, len(indices),
+                                       calculation_width)]
+        translation_chunks = [translations[i:i + calculation_width]
+                              for i in range(0, len(indices),
+                                             calculation_width)]
+        
+            
+        # Then we simulate the results
+        data = [self.forward(idx, trans).detach()
+                for idx, trans in zip(index_chunks, translation_chunks)]
+
+        data = t.cat(data, dim=0)
+        # 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=t.float32, device=self.probe.device)
+        if (hasattr(self, 'translation_offsets') and
+            self.translation_offsets is not None):
+            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
+        else:
+            return translations
+
+
+    def get_probes(self, space='real'):
+        # This is the part where I need to create the probe modes using the
+        # time-dependent stuff
+
+        # This restricts the basis probes to stay within the probe support
+        optic_mask = self.probe * self.probe_fourier_support
+
+        probes = t.zeros([len(self.frame_delays)] + list(self.probe.shape),
+                         dtype=self.probe.dtype, device=self.probe.device)
+        for i, delay in enumerate(self.frame_delays):
+            indices = t.bucketize(self.fourier_times, self.times + delay)
+            clamped_indices = t.clamp(indices-1, max=len(self.time_dependence)-1)
+            illumination = t.take(self.time_dependence, clamped_indices)
+            illumination[indices==0] = 0
+            illumination[indices==len(self.time_dependence)+1] = 0
+            probes[i] = illumination * optic_mask
+
+        if space.lower()=='real':
+            return tools.propagators.inverse_far_field(probes)
+        elif space.lower()=='fourier' or space.lower()=='reciprocal':
+            return probes
+
+        
+    plot_list = [
+        ('Probe Amplitudes (scroll to view modes)',
+         lambda self, fig: p.plot_amplitude(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Probe Phases (scroll to view modes)',
+         lambda self, fig: p.plot_phase(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Fourier Probe Amplitudes (scroll to view modes)',
+         lambda self, fig: p.plot_amplitude(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Fourier Probe Phases (scroll to view modes)',
+         lambda self, fig: p.plot_phase(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Optic Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Optic Phase',
+         lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
+         lambda self: len(self.weights.shape) >= 2),
+        ('% Power in Top Mode (only accurate after tidy_probes)',
+         lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
+         lambda self: len(self.weights.shape) >= 2),
+        ('Object Amplitude',
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)),
+        ('Time Structure',
+         lambda self, fig: (plt.figure(fig.number) and (plt.clf() or True) and plt.plot(self.time_dependence.real.detach().cpu().numpy()) and plt.plot(self.time_dependence.imag.detach().cpu().numpy())))
+    ]
+
+#    def plot_errors(self, dataset):
+        
+    
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        optic = self.probe.detach().cpu().numpy()
+        optic = optic * self.probe_norm.detach().cpu().numpy()
+        time_dependence = self.time_dependence.detach().cpu().numpy()
+        times = self.times.detach().cpu().numpy()
+        fourier_times = self.fourier_times.detach().cpu().numpy()
+        probes = self.get_probes().detach().cpu().numpy()
+        obj = self.obj.detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+
+        return {'basis': basis, 'translation': translations,
+                'probes': probes, 'optic': optic,
+                'times': times, 'time_dependence': time_dependence,
+                'fourier_times': fourier_times,
+                'obj': obj,
+                'background': background,
+                'weights': weights,
+                }

src/CDTools/models/time_resolved_rpi.py

@@ -0,0 +1,302 @@
+import torch as t
+from CDTools.models import CDIModel
+from CDTools import tools
+from CDTools.tools import plotting as p
+from CDTools.tools.interactions import RPI_interaction
+from CDTools.tools import initializers
+from scipy.ndimage.morphology import binary_dilation
+import numpy as np
+from copy import copy
+
+__all__ = ['MultimodeRPI']
+
+
+
+__all__ = ['RPI']
+
+class TimeResolvedRPI(CDIModel):
+
+    @property
+    def obj(self):
+        return t.complex(self.obj_real, self.obj_imag)
+
+    
+    def __init__(self, wavelength, detector_geometry, probe_basis,
+                 probe, obj_guess, framerate, detector_slice=None,
+                 background=None, mask=None, saturation=None,
+                 obj_support=None, oversampling=1):
+
+        super(TimeResolvedRPI, self).__init__()
+
+        self.wavelength = t.tensor(wavelength)
+        self.framerate = framerate
+        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.probe_basis = t.tensor(probe_basis)
+
+        scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
+                                 probe.shape[-2]/obj_guess.shape[-2]])
+        self.obj_basis = self.probe_basis * scale_factor
+        self.detector_slice = detector_slice
+
+        # Maybe something to include in a bit
+        # self.surface_normal = t.tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.tensor(mask, dtype=t.bool)
+
+            
+        self.probe = t.tensor(probe, dtype=t.complex64)
+
+        obj_guess = t.tensor(obj_guess, dtype=t.complex64)
+        
+        self.obj_real = t.nn.Parameter(obj_guess.real)
+        self.obj_imag = t.nn.Parameter(obj_guess.imag)
+        
+        # Wait for LBFGS to be updated for complex-valued parameters
+        # 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,
+                    dtype=t.float32)
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape,
+                                           dtype=t.float32)
+                
+        self.background = t.tensor(background, dtype=t.float32)
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support[None, ...]
+        else:
+            self.obj_support = t.ones_like(self.obj[0, ...])
+
+        self.oversampling = oversampling
+
+
+    @classmethod
+    def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, probe_threshold=0):
+        
+        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')
+        # We only need the patterns here, not the inputs associated with them.
+        _, 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=opt_for_fft,
+                                                   oversampling=oversampling)
+
+        if not isinstance(probe,t.Tensor):
+            probe = t.as_tensor(probe)
+        
+        if background is None and hasattr(dataset, 'background') \
+           and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        elif background is not None:
+            background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
+
+        det_geo = dataset.detector_geometry
+
+        # If no mask is given, but one exists in the dataset, load it.
+        if mask is None and hasattr(dataset, 'mask') \
+           and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+
+        # Now we initialize the object
+        if obj_size is None:
+            # This is a standard size for a well-matched probe and detector
+            obj_size = list((np.array(probe_shape) // 2).astype(int))
+
+        # I think something to do with the fact that the object is defined
+        # on a coarser grid needs to be accounted for here that is not
+        # accounted for yet
+        scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
+        n_modes = (probe.shape[0] - 1) // framerate + 1
+
+        obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
+        
+
+        probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
+        probe_fft = tools.propagators.far_field(probe_intensity)
+        pad0l = (probe.shape[-2] - obj_size[-2])//2
+        pad0r = probe.shape[-2] - obj_size[-2] - pad0l
+        pad1l = (probe.shape[-1] - obj_size[-1])//2
+        pad1r = probe.shape[-1] - obj_size[-1] - pad1l
+        probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
+        probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
+
+        obj_support = probe_lr > t.max(probe_lr) * probe_threshold
+        obj_support = t.as_tensor(binary_dilation(obj_support))
+
+        return cls(wavelength, det_geo, probe_basis,
+                   probe, obj_guess, framerate, detector_slice=det_slice,
+                   background=background, mask=mask, saturation=saturation,
+                   obj_support=obj_support, oversampling=oversampling)
+
+
+    def random_init(self, pattern):
+        scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
+        self.obj.data = scale * t.exp(
+            2j * np.pi * t.rand(self.obj.shape)).to(
+                dtype=self.obj.dtype, device=self.obj.device)
+        
+    
+    # Needs work
+    def interaction(self, index, *args):
+        # including *args allows this to work with all sorts of datasets
+        # that might include other information in with the index in their
+        # "input" parameters (such as translations for a ptychography dataset).
+        # This makes it seamless to use such a dataset even though those
+        # extra arguments will not be used.
+
+        
+        all_exit_waves = []
+
+        # Mix the probes with the weight matrix
+        prs = self.probe
+        
+        for i in range(self.probe.shape[0]):
+            obj_frame = i // self.framerate
+            pr = prs[i]
+            exit_waves = RPI_interaction(pr,
+                                         self.obj_support * self.obj[obj_frame])
+            all_exit_waves.append(exit_waves.unsqueeze(0))
+
+        # This creates a bunch of modes generated from all possible combos
+        # of the probe and object modes all strung out along the first index
+
+        output = t.cat(all_exit_waves)
+
+        # If we have multiple indexes input, we unsqueeze and repeat the stack
+        # of wavefields enough times to simulate each requested index. This
+        # seems silly, but it enables (for example) one to do a reconstruction
+        # from a set of diffraction patterns that are all known to be from the
+        # same object.
+        try:
+            # will fail if index has no length, for example when index
+            # is just an int. In this case, we just do nothing instead
+            output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
+        except TypeError:
+            pass
+        return output
+
+
+    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):
+        # Here I'm taking advantage of an undocumented feature in the
+        # incoherent_sum measurement function where it will work with
+        # a 4D wavefield array as well as a 5D array.
+        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)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
+
+    def regularizer(self, factors):
+        return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+            + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
+        
+    def to(self, *args, **kwargs):
+        super(TimeResolvedRPI, 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.probe = self.probe.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.obj_basis = self.obj_basis.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.background = self.background.to(*args, **kwargs)
+        
+        # Maybe include in a bit
+        #self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+    def sim_to_dataset(self, args_list):
+        raise NotImplementedError('No sim to dataset yet, sorry!')
+
+    plot_list = [
+        ('Probe Amplitudes',
+         lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
+        ('Object Amplitudes', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
+                                            basis=self.obj_basis)),
+        ('Object Phases',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig,
+                                        basis=self.obj_basis))
+    ]
+
+
+    def save_results(self, dataset=None, full_obj=False):
+        # dataset is set as a kwarg here because it isn't needed, but the
+        # common pattern is to pass a dataset. This makes it okay if one
+        # continues to use that standard pattern
+        probe_basis = self.probe_basis.detach().cpu().numpy()
+        obj_basis = self.obj_basis.detach().cpu().numpy()
+        probe = self.probe.detach().cpu().numpy()
+        # Provide the option to save out the subdominant objects or
+        # just the dominant one
+        if full_obj:
+            obj = self.obj.detach().cpu().numpy()
+        else:
+            obj = self.obj[0].detach().cpu().numpy()
+        background = self.background.detach().cpu().numpy()**2
+        
+        return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
+                'probe': probe,'obj': obj,
+                'background': background}
+

src/CDTools/tools/__init__.py

@@ -0,0 +1,27 @@
+""" This module contains various packages of tools supporting common needs
+
+It at times feels a bit silly to keep all the tools siloed under a particular
+subpackage such as "projectors" or "measurements", especially given that
+"flat is better than nested", but I find something comforting about the
+organization and the ability to only pull in the particular set of tools
+that one needs for a specific application.
+
+The submodules are all structured as modules with their own __init__ files,
+which use an import * statement to import from a file defining the various
+functions. This is done to prevent leakage of imported packages into the
+namespace of CDTools. I know, we're all consenting adults, but I just hate
+having numpy and torch defined under CDTools.tools.cmath, you know?
+
+"""
+
+from CDTools.tools import losses
+from CDTools.tools import data
+from CDTools.tools import image_processing
+from CDTools.tools import initializers
+from CDTools.tools import plotting
+from CDTools.tools import interactions
+from CDTools.tools import propagators
+from CDTools.tools import measurements
+from CDTools.tools import analysis
+from CDTools.tools import atoms
+from CDTools.tools import polarization

src/CDTools/tools/analysis/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.analysis.analysis import *

src/CDTools/tools/analysis/analysis.py

@@ -0,0 +1,960 @@
+"""Contains basic functions for analyzing the results of reconstructions
+
+The functions in this module are designed to work either with pytorch tensors
+or numpy arrays, so they can be used either directly after reconstructions
+on the attributes of the models themselves, or after-the-fact once the
+data has been stored in numpy arrays.
+"""
+
+import torch as t
+import numpy as np
+from CDTools.tools import image_processing as ip
+from scipy import fftpack
+from scipy import linalg as sla
+from scipy import special
+
+__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
+           'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
+           'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
+           'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error']
+
+
+def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
+    """Orthogonalizes a set of incoherently mixing probes
+
+    The strategy is to define a reduced orthogonal basis that spans
+    all of the retrieved probes, and then build the density matrix
+    defined by the probes in that basis. After diagonalization, the
+    eigenvectors can be recast into the original basis and returned
+
+    By default, it assumes that the set of probes are defined just as
+    standard incoherently mixing probe modes, and orthogonalizes them.
+    However, if a density matrix is explicitly given, it will instead
+    consider the problem of extracting the eigenbasis of the matrix
+    probes * denstity_matrix * probes^dagger, where probes is the
+    column matrix of the given probe functions. This latter problem arises
+    in the generalization of the probe mixing model, and reduces to the
+    simpler case when the density matrix is equal to the identity matrix
+    
+    If the parameter "keep_transform" is set, the function will additionally
+    return the matrix A such that A * ortho_probes^dagger = probes^dagger
+
+    If the parameter "normalize" is False (as is the default), the variation
+    in intensities in the probe modes will be kept in the probe modes, as is
+    natural for a purely incoherent model. If it is set to "True", the
+    returned probe modes will all be normalized instead.
+
+    Parameters
+    ----------
+    probes : array
+        An l x n x m complex array representing  a stack of probes
+    density_matrix : np.array
+        An optional l x l density matrix further elaborating on the state
+    keep_transform : bool
+        Default False, whether to return the map from probes to ortho_probes
+    normalize : bool
+        Default False, whether to normalize the probe modes
+
+    Returns
+    -------
+    ortho_probes: array
+        An l x n x m complex array representing a stack of probes
+    """
+    
+    try:
+        probes = probes.detach().cpu().numpy()
+        send_to_torch = True
+    except:
+        send_to_torch = False
+
+    # We can do the orthogonalization with an SVD, so first we have to
+    # reshape the final two dimensions (the image shape) into a single
+    # vectorized dimension. This matrix is probes^dagger, hence the
+    # conjugation
+    probes_mat = probes.reshape(probes.shape[0],
+                                probes.shape[1]*probes.shape[2])
+
+    if density_matrix is None:
+        density_matrix = np.eye(probes.shape[0])
+    
+        
+    # next we want to extract the eigendecomposition of the density matrix
+    # itself
+    w,v = sla.eigh(density_matrix)
+    w = w[::-1]
+    v = v[:,::-1]
+
+    # We do this just to avoid total failure when the density
+    # matrix is not positive definite.
+    # In most cases (such as when rho is generated directly from some other
+    # matrix A such that rho=A A^dagger), w should never have any negative
+    # entries.
+    w = np.maximum(w,0)
+
+    B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose())
+
+    #u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    u,s,vh = sla.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    
+
+    if normalize:
+        ortho_probes = vh.reshape(probes.shape[0],
+                                  probes.shape[1],
+                                  probes.shape[2])
+
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,np.dot(u,np.diag(s)))
+        #A_dagger = np.dot(np.linalg.pinv(np.diag(s)),
+        #                  np.dot(np.transpose(u).conj(),B_dagger))
+    else:
+        ortho_probes = np.dot(np.diag(s),vh).reshape(probes.shape[0],
+                                                     probes.shape[1],
+                                                     probes.shape[2])
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,u)
+        #A_dagger = np.dot(np.transpose(u).conj(),B_dagger)
+    
+    if send_to_torch:
+        ortho_probes = t.as_tensor(np.stack(ortho_probes))
+        A = t.as_tensor(A)
+
+    if keep_transform:
+        return ortho_probes, A#_dagger
+    else:
+        return ortho_probes 
+
+    
+def standardize(probe, obj, obj_slice=None, correct_ramp=False):
+    """Standardizes a probe and object to prepare them for comparison
+
+    There are a number of ambiguities in the definition of a ptychographic
+    reconstruction. This function makes an explicit choice for each ambiguity
+    to allow comparisons between independent reconstructions without confusing
+    these ambiguities for real differences between the reconstructions.
+
+    The ambiguities and standardizations are:
+
+    1) a. Probe and object can be scaled inversely to one another
+       b. So we set the probe intensity to an average per-pixel value of 1
+    2) a. The probe and object can aquire equal and opposite phase ramps
+       b. So we set the centroid of the FFT of the probe to zero frequency
+    3) a. The probe and object can each acquire an arbitrary overall phase
+       b. So we set the phase of the sum of all values of both the probe and object to 0
+
+    When dealing with the properties of the object, a slice is used by
+    default as the edges of the object often are dominated by unphysical
+    noise. The default slice is from 3/8 to 5/8 of the way across. If the
+    probe is actually a stack of incoherently mixing probes, then the
+    dominant probe mode (assumed to be the first in the list) is used, but
+    all the probes are updated with the same factors.
+
+    Parameters
+    ----------
+    probe : array
+        A complex array storing a retrieved probe or stack of incoherently mixed probes
+    obj : array
+        A complex array storing a retrieved object
+    obj_slice : slice
+        Optional, a slice to take from the object for calculating normalizations
+    correct_ramp : bool
+        Default False, whether to correct for the relative phase ramps
+
+    Returns
+    -------
+    standardized_probe : array
+        The standardized probe
+    standardized_obj : array
+        The standardized object
+
+    """
+    # First, we normalize the probe intensity to a fixed value.
+    probe_np = False
+    if isinstance(probe, np.ndarray):
+        probe = t.as_tensor(probe, dtype=t.complex64)
+        probe_np = True
+    obj_np = False
+    if isinstance(obj, np.ndarray):
+        obj = t.as_tensor(obj,dtype=t.complex64)
+        obj_np = True
+
+    # If this is a single probe and not a stack of probes
+    if len(probe.shape) == 2:
+        probe = probe[None,...]
+        single_probe = True
+    else:
+        single_probe = False
+
+    normalization = t.sqrt(t.mean(t.abs(probe[0])**2))
+    probe = probe / normalization
+    obj = obj * normalization
+
+    # Default slice of the object to use for alignment, etc.
+    if obj_slice is None:
+        obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
+                          (obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
+
+
+    if correct_ramp:
+        # Need to check if this is actually working and, if not, why not
+        center_freq = ip.centroid(t.abs(t.fft.fftshift(t.fft.fft2(probe[0]),
+                                                       dim=(-1,-2)))**2)
+        center_freq -= t.div(t.tensor(probe[0].shape,dtype=t.float32),2,rounding_mode='floor')
+        center_freq /= t.as_tensor(probe[0].shape,dtype=t.float32)
+
+        Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
+        probe_phase_ramp = t.exp(2j * np.pi *
+                                 (center_freq[0] * t.tensor(Is).to(t.float32) +
+                                  center_freq[1] * t.tensor(Js).to(t.float32)))
+        probe = probe *  t.conj(probe_phase_ramp)
+        Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]]
+        obj_phase_ramp = t.exp(2j*np.pi *
+                               (center_freq[0] * t.tensor(Is).to(t.float32) +
+                                center_freq[1] * t.tensor(Js).to(t.float32)))
+        obj = obj * obj_phase_ramp
+
+    # Then, we set them to consistent absolute phases
+
+    obj_angle = t.angle(t.sum(obj[obj_slice]))
+    obj = obj *  t.exp(-1j*obj_angle)
+
+    for i in range(probe.shape[0]):
+        probe_angle = t.angle(t.sum(probe[i]))
+        probe[i] = probe[i] * t.exp(-1j*probe_angle)
+
+    if single_probe:
+        probe = probe[0]
+
+    if probe_np:
+        probe = probe.detach().cpu().numpy()
+    if obj_np:
+        obj = obj.detach().cpu().numpy()
+
+    return probe, obj
+
+
+
+def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, correct_ramp=False):
+    """Takes a collection of reconstructions and outputs a single synthesized probe and object
+
+    The function first standardizes the sets of probes and objects using the
+    standardize function, passing through the relevant options. Then it
+    calculates the closest overlap of subsequent frames to subpixel
+    precision and uses a sinc interpolation to shift all the probes and objects
+    to a common frame. Then the images are summed.
+
+    Parameters
+    ----------
+    probes : list(array)
+        A list of probes or stacks of probe modes
+    objects : list(array)
+        A list of objects
+    use_probe : bool
+        Default False, whether to use the probe or object for alignment
+    obj_slice : slice
+        Optional, A slice of the object to use for alignment and normalization
+    correct_ramp : bool
+        Default False, whether to correct for a relative phase ramp in the probe and object
+
+    Returns
+    -------
+    synth_probe : array
+        The synthesized probe
+    synth_obj : array
+        The synthesized object
+    obj_stack : list(array)
+        A list of standardized objects, for further processing
+    """
+
+    # This should be cleaned up so it accepts anything array_like
+    probe_np = False
+    if isinstance(probes[0], np.ndarray):
+        probes = [t.as_tensor(probe,dtype=t.complex64) for probe in probes]
+        probe_np = True
+    obj_np = False
+    if isinstance(objects[0], np.ndarray):
+        objects = [t.as_tensor(obj,dtype=t.complex64) for obj in objects]
+        obj_np = True
+
+    obj_shape = np.min(np.array([obj.shape for obj in objects]),axis=0)
+    objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects]
+
+    if obj_slice is None:
+        obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5,
+                          (objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5]
+
+
+
+    synth_probe, synth_obj = standardize(probes[0].clone(), objects[0].clone(), obj_slice=obj_slice,correct_ramp=correct_ramp)
+
+    obj_stack = [synth_obj]
+
+    for i, (probe, obj) in enumerate(zip(probes[1:],objects[1:])):
+        probe, obj = standardize(probe.clone(), obj.clone(), obj_slice=obj_slice,correct_ramp=correct_ramp)
+        if use_probe:
+            shift = ip.find_shift(synth_probe[0],probe[0], resolution=50)
+        else:
+            shift = ip.find_shift(synth_obj[obj_slice],obj[obj_slice], resolution=50)
+
+
+        obj = ip.sinc_subpixel_shift(obj,np.array(shift))
+
+        if len(probe.shape) == 3:
+            probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift))
+                             for p in probe],dim=0)
+        else:
+            probe = ip.sinc_subpixel_shift(probe,tuple(shift))
+
+        synth_probe = synth_probe + probe
+        synth_obj = synth_obj + obj
+        obj_stack.append(obj)
+
+
+    # If there only was one image
+    try:
+        i
+    except:
+        i = -1
+
+    if probe_np:
+        synth_probe = synth_probe.numpy()
+    if obj_np:
+        synth_obj = synth_obj.numpy()
+        obj_stack = [obj.numpy() for obj in obj_stack]
+
+    return synth_probe/(i+2), synth_obj/(i+2), obj_stack
+
+
+
+def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
+    """Calculates a PRTF between each the individual objects and a synthesized one
+
+    The consistency PRTF at any given spatial frequency is defined as the ratio
+    between the intensity of any given reconstruction and the intensity
+    of a synthesized or averaged reconstruction at that spatial frequency.
+    Typically, the PRTF is averaged over spatial frequencies with the same
+    magnitude.
+
+    Parameters
+    ----------
+    synth_obj : array
+        The synthesized object in the numerator of the PRTF
+    objects : list(array)
+        A list of objects or diffraction patterns for the denomenator of the PRTF
+    basis : array
+        The basis for the reconstruction array to allow output in physical unit
+    obj_slice : slice
+        Optional, a slice of the objects to use for calculating the PRTF
+    nbins : int
+        Optional, number of bins to use in the histogram. Defaults to a sensible value
+
+    Returns
+    -------
+    freqs : array
+        The frequencies for the PRTF
+    PRTF : array
+        The values of the PRTF
+    """
+
+    obj_np = False
+    if isinstance(objects[0], np.ndarray):
+        objects = [t.as_tensor(obj, dtype=t.complex64) for obj in objects]
+        obj_np = True
+    if isinstance(synth_obj, np.ndarray):
+        synth_obj = t.as_tensor(synth_obj, dtype=t.complex64)
+
+    if isinstance(basis, t.Tensor):
+        basis = basis.detach().cpu().numpy()
+
+    if obj_slice is None:
+        obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5,
+                          (objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5]
+
+    if nbins is None:
+        nbins = np.max(synth_obj[obj_slice].shape) // 4
+
+    synth_fft = (t.abs(t.fft.fftshift(t.fft.fft2(synth_obj[obj_slice]), dim=(-1,-2)))**2).numpy()
+
+
+    di = np.linalg.norm(basis[:,0])
+    dj = np.linalg.norm(basis[:,1])
+
+    i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di))
+    j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj))
+
+    Js,Is = np.meshgrid(j_freqs,i_freqs)
+    Rs = np.sqrt(Is**2+Js**2)
+
+
+    synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft)
+
+    prtfs = []
+    for obj in objects:
+        obj = obj[obj_slice]
+        single_fft = (t.abs(t.fft.fftshift(t.fft.fft2(obj),
+                                           dim=(-1,-2)))**2).numpy()
+        single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
+
+        prtfs.append(synth_ints/single_ints)
+
+
+    prtf = np.mean(prtfs,axis=0)
+
+    if not obj_np:
+        bins = t.Tensor(bins)
+        prtf = t.Tensor(prtf)
+
+    return bins[:-1], prtf
+
+
+
+def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
+    """Calculates a cross-correlation between two images with their autocorrelations deconvolved.
+
+    This is formally defined as the inverse Fourier transform of the normalized
+    product of the Fourier transforms of the two images. It results in a
+    kernel, whose characteristic size is related to the exactness of the
+    possible alignment between the two images, on top of a random background
+
+    Parameters
+    ----------
+    im1 : array
+        The first image, as a complex or real valued array
+    im2 : array
+        The first image, as a complex or real valued array
+    im_slice : slice
+        Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
+
+    Returns
+    -------
+    corr : array
+        The complex-valued deconvolved cross-correlation, in real space
+
+    """
+
+    im_np = False
+    if isinstance(im1, np.ndarray):
+        im1 = t.as_tensor(im1)
+        im_np = True
+    if isinstance(im2, np.ndarray):
+        im2 = t.as_tensor(im2)
+        im_np = True
+
+    if im_slice is None:
+        im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
+                          (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
+
+
+    cor_fft = t.fft.fft2(im1[im_slice]) * \
+        t.conj(t.fft.fft2(im2[im_slice]))
+
+    # Not sure if this is more or less stable than just the correlation
+    # maximum - requires some testing
+    cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
+
+    if im_np:
+        cor = cor.numpy()
+
+    return cor
+
+
+def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
+    """Calculates a Fourier ring correlation between two images
+
+    This function requires an input of a basis to allow for FRC calculations
+    to be related to physical units.
+
+    Like other analysis functions, this can take input in numpy or pytorch,
+    and will return output in the respective format.
+
+    Parameters
+    ----------
+    im1 : array
+        The first image, a complex or real valued array
+    im2 : array
+        The first image, a complex or real valued array
+    basis : array
+        The basis for the images, defined as is standard for datasets
+    im_slice : slice
+        Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
+    nbins : int
+        Number of bins to break the FRC up into
+    snr : float
+        The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
+    limit : str
+        Default is 'side'. What is the highest frequency to calculate the FRC to? If 'side', it chooses the side of the Fourier transform, if 'corner' it goes fully to the corner.
+
+    Returns
+    -------
+    freqs : array
+        The frequencies associated with each FRC value
+    FRC : array
+        The FRC values
+    threshold : array
+        The threshold curve for comparison
+
+    """
+
+    im_np = False
+    if isinstance(im1, np.ndarray):
+        im1 = t.as_tensor(im1)
+        im_np = True
+    if isinstance(im2, np.ndarray):
+        im2 = t.as_tensor(im2)
+        im_np = True
+
+    if isinstance(basis, np.ndarray):
+        basis = t.tensor(basis)
+
+    if im_slice is None:
+        im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
+                          (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
+
+    if nbins is None:
+        nbins = np.max(im1[im_slice].shape) // 8
+
+    f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
+    f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
+    cor_fft = f1 * t.conj(f2)
+
+    F1 = t.abs(f1)**2
+    F2 = t.abs(f2)**2
+
+
+    di = np.linalg.norm(basis[:,0])
+    dj = np.linalg.norm(basis[:,1])
+
+    i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
+    j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
+
+    Js,Is = np.meshgrid(j_freqs,i_freqs)
+    Rs = np.sqrt(Is**2+Js**2)
+
+    if limit.lower().strip() == 'side':
+        max_i = np.max(i_freqs)
+        max_j = np.max(j_freqs)
+        frc_range = [0, max(max_i,max_j)]
+    elif limit.lower().strip() == 'corner':
+        frc_range = [0, np.max(Rs)]
+    else:
+        raise ValueError('Invalid FRC limit: choose "side" or "corner"')
+
+    numerator, bins = np.histogram(Rs, bins=nbins, range=frc_range,
+                                   weights=cor_fft.numpy())
+    denominator_F1, bins = np.histogram(Rs, bins=nbins, range=frc_range,
+                                        weights=F1.detach().cpu().numpy())
+    denominator_F2, bins = np.histogram(Rs, bins=nbins, range=frc_range,
+                                        weights=F2.detach().cpu().numpy())
+    n_pix, bins = np.histogram(Rs, bins=nbins, range=frc_range)
+
+    
+    #n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window
+    
+    frc = np.abs(numerator) / np.sqrt(denominator_F1*denominator_F2)
+
+    # This moves from combined-image SNR to single-image SNR
+    snr /= 2
+
+    # NOTE: I should update this  to produce lots of different threshold curves
+    # sigma, 2sigma, 3sigma, traditional FRC 1-bit, my better one, n_pix, etc.
+    
+    threshold = (snr + (2 * np.sqrt(snr) + 1) / np.sqrt(n_pix)) / \
+        (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
+
+    my_threshold = np.sqrt(snr**2 + (2*snr**2 + 2*snr + 1)/n_pix) / \
+        np.sqrt(snr**2 + 2 * snr + 1 + 2*snr**2 / n_pix)
+
+    twosigma_threshold = 2/ np.sqrt(n_pix)
+    
+    if not im_np:
+        bins = t.tensor(bins)
+        frc = t.tensor(frc)
+        threshold = t.tensor(threshold)
+
+    return bins[:-1], frc, threshold
+
+
+def calc_vn_entropy(matrix):
+    """Calculates the Von Neumann entropy of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to calculate the entropy of
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The entropy or entropies of the arrays
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        eigs = [eig / np.sum(eig) for eig in eigs]
+        # And calculate the VN entropy!
+        entropies = [-np.sum(special.xlogy(eig,eig)) for eig in eigs]
+        return np.array(entropies)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
+        return entropy
+
+    
+def calc_top_mode_fraction(matrix):
+    """Calculates the fraction of total power in the top mode of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to work from
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The fraction of power in the top mode of each matrix
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        fractions = [np.max(eig) / np.sum(eig) for eig in eigs]
+        return np.array(fractions)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        fraction = np.max(eig) / np.sum(eig)
+        return fraction
+    
+
+def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
+                   dims=2):
+    """Calculates the root-mean-squared error between two complex wavefields
+
+    The formal definition of this function is:
+
+    output = norm * sqrt(mean(abs(field_1 - gamma * field_2)**2))
+    
+    Where norm is an optional normalization factor, and gamma is an
+    optional phase factor which is appropriate when the wavefields suffer
+    from a global phase degeneracy as is often the case in diffractive
+    imaging.
+
+    The normalization is defined as the square root of the total intensity
+    contained in field_1, which is appropriate when field_1 represents a
+    known ground truth:
+
+    norm = sqrt(mean(abs(field_1)**2))
+    
+    The phase offset is an analytic expression for the phase offset which
+    will minimize the RMS error between the two wavefields:
+
+    gamma = exp(1j * angle(sum(field_1 * conj(field_2))))
+    
+    This implementation is stable even in cases where field_1 and field_2
+    are completely orthogonal.
+
+    In the definitions above, the field_n are n-dimensional wavefields. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the default is 2 for a 2D wavefield.
+    
+    Parameters
+    ----------
+    field_1 : array
+        The first complex-valued field
+    field_2 : array
+        The second complex-valued field
+    align_phases : bool
+        Default is True, whether to account for a global phase offset
+    normalize : bool
+        Default is False, whether to normalize to the intensity of field_1
+    dims : (int or tuple of python:ints)
+        Default is 2, the number of final dimensions to reduce over.
+
+    
+    Returns
+    -------
+    rms_error : float or t.Tensor
+        The RMS error, or tensor of RMS errors, depending on the dim argument
+
+    """
+
+    sumdims = tuple(d - dims for d in range(dims))
+        
+    if align_phases:
+        # Keepdim allows us to broadcast the result correctly when we
+        # multiply by the fields
+        gamma = t.exp(1j * t.angle(t.sum(field_1 * t.conj(field_2), dim=sumdims,
+                                         keepdim=True)))
+    else:
+        gamma = 1
+
+    if normalize:
+        norm = 1 / t.mean(t.abs(field_1)**2, dim=sumdims)
+    else:
+        norm = 1
+
+    difference = field_1 - gamma * field_2
+    
+    return t.sqrt(norm * t.mean(t.abs(difference)**2, dim=sumdims))
+
+
+def calc_fidelity(fields_1, fields_2, dims=2):
+    """Calculates the fidelity between two density matrices
+
+    The fidelity is a comparison metric between two density matrices
+    (i.e. mutual coherence functions) that extends the idea of the
+    overlap to incoherent light. As a reminder, the overlap between two
+    fields is:
+
+    overlap = abs(sum(field_1 * field_2))**2
+    
+    Whereas the fidelity is defined as:
+    
+    fidelity = trace(sqrt(sqrt(dm_1) <dot> dm_2 <dot> sqrt(dm_1)))**2
+
+    where dm_n refers to the density matrix encoded by fields_n such
+    that dm_n = fields_n <dot> fields_<n>.conjtranspose(), sqrt
+    refers to the matrix square root, and <dot> is the matrix product.
+    
+    This is not a practical implementation, however, as it is not feasible
+    to explicitly construct the matrices dm_1 and dm_2 in memory. Therefore,
+    we take advantage of the alternate definition based directly on the
+    fields_<n> parameter:
+
+    fidelity = sum(svdvals(fields_1 <dot> fields_2.conjtranspose()))**2
+    
+    In the definitions above, the fields_n are regarded as collections of
+    wavefields, where each wavefield is by default 2-dimensional. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the fields_n arguments must always have at least one more dimension
+    than the dims argument. Any additional dimensions are treated as batch
+    dimensions.
+    
+    Parameters
+    ----------
+    fields_1 : array
+        The first set of complex-valued field modes
+    fields_2 : array
+        The second set of complex-valued field modes
+    dims : int
+        Default is 2, the number of final dimensions to reduce over.
+
+    
+    Returns
+    -------
+    fidelity : float or t.Tensor
+        The fidelity, or tensor of fidelities, depending on the dim argument
+
+    """
+
+    fields_1 = t.as_tensor(fields_1)
+    fields_2 = t.as_tensor(fields_2)
+    
+    mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj()
+    sumdims = tuple(d - dims for d in range(dims))
+    mat = t.sum(mult,dim=sumdims)
+
+    # Because I think this is the nuclear norm squared, I would like to swap
+    # Out the definition for this, but I need to test it before swapping.
+    # It also probably makes sense to implement sqrt_fidelity separately
+    # because that's more important
+    #return t.linalg.matrix_norm(mat, ord='nuc')**2
+    
+    # I think this is just the nuclear norm.
+    svdvals = t.linalg.svdvals(mat)
+    return t.sum(svdvals, dim=-1)**2
+
+
+def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
+    """Calculates a generalization of the root-mean-squared error between two complex wavefields
+
+    This function calculates an generalization of the RMS error which uses the
+    concept of fidelity to extend it to capture the error between
+    incoherent wavefields, defined as a mode decomposition. The extension has
+    several nice properties, in particular:
+
+    1) For coherent wavefields, it precisely matches the RMS error including
+       a correction for the global phase degeneracy (align_phases=True)
+    2) All mode decompositions of either field that correspond to the same
+       density matrix / mutual coherence function will produce the same 
+       output
+    3) The error will only be zero when comparing mode decompositions that
+       correspond to the same density matrix.
+    4) Due to (2), one need not worry about the ordering of the modes,
+       properly orthogonalizing the modes, and it is even possible to
+       compare mode decompositions with different numbers of modes.    
+    
+    The formal definition of this function is:
+
+    output = norm * sqrt(mean(abs(fields_1)**2)
+                         + mean(abs(fields_2)**2)
+                         - 2 * sqrt(fidelity(fields_1,fields_2)))
+    
+    Where norm is an optional normalization factor, and the fidelity is
+    defined based on the mean, rather than the sum, to match the convention
+    for the root *mean* squared error.
+
+    The normalization is defined as the square root of the total intensity
+    contained in fields_1, which is appropriate when fields_1 represents a
+    known ground truth:
+
+    norm = sqrt(mean(abs(fields_1)**2))
+
+    In the definitions above, the fields_n are regarded as collections of
+    wavefields, where each wavefield is by default 2-dimensional. The
+    dimensionality of the wavefields can be altered via the dims argument,
+    but the fields_n arguments must always have at least one more dimension
+    than the dims argument. Any additional dimensions are treated as batch
+    dimensions.
+    
+    Parameters
+    ----------
+    fields_1 : array
+        The first set of complex-valued field modes
+    fields_2 : array
+        The second set of complex-valued field modes
+    normalize : bool
+        Default is False, whether to normalize to the intensity of fields_1
+    dims : (int or tuple of python:ints)
+        Default is 2, the number of final dimensions to reduce over.
+
+    Returns
+    -------
+    rms_error : float or t.Tensor
+        The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
+
+    """
+    fields_1 = t.as_tensor(fields_1)
+    fields_2 = t.as_tensor(fields_2)
+
+    npix = t.prod(t.as_tensor(fields_1.shape[-dims:],dtype=t.int32))
+    
+    sumdims = tuple(d - dims - 1 for d in range(dims+1))
+    fields_1_intensity = t.sum(t.abs(fields_1)**2,dim=sumdims) / npix
+    fields_2_intensity = t.sum(t.abs(fields_2)**2,dim=sumdims) / npix
+    fidelity = calc_fidelity(fields_1, fields_2, dims=dims) / npix**2
+
+    result = fields_1_intensity + fields_2_intensity - 2 * t.sqrt(fidelity)
+    
+    if normalize:
+        result /= fields_1_intensity
+
+    return t.sqrt(result)
+    
+
+def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.):
+    """Calculates a Fourier ring correlation between two images
+
+    This function requires an input of a basis to allow for FRC calculations
+    to be related to physical units.
+
+    Like other analysis functions, this can take input in numpy or pytorch,
+    and will return output in the respective format.
+
+    Parameters
+    ----------
+    im1 : array
+        The first image, a complex or real valued array
+    im2 : array
+        The first image, a complex or real valued array
+    basis : array
+        The basis for the images, defined as is standard for datasets
+    im_slice : slice
+        Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
+    nbins : int
+        Number of bins to break the FRC up into
+    snr : float
+        The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
+
+    Returns
+    -------
+    freqs : array
+        The frequencies associated with each FRC value
+    FRC : array
+        The FRC values
+    threshold : array
+        The threshold curve for comparison
+
+    """
+
+    im_np = False
+    if isinstance(fields_1, np.ndarray):
+        fields_1 = t.as_tensor(fields_)
+        im_np = True
+    if isinstance(fields_2, np.ndarray):
+        fields_2 = t.as_tensor(fields_2)
+        im_np = True
+
+    if isinstance(basis, np.ndarray):
+        basis = t.tensor(basis)
+
+    if im_slice is None:
+        im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
+                          (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
+
+    if nbins is None:
+        nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4
+
+
+    f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
+    f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
+    cor_fft = f1 * t.conj(f2)
+        
+
+    F1 = t.abs(f1)**2
+    F2 = t.abs(f2)**2
+
+
+    di = np.linalg.norm(basis[:,0])
+    dj = np.linalg.norm(basis[:,1])
+
+    i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
+    j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
+
+    Js,Is = np.meshgrid(j_freqs,i_freqs)
+    Rs = np.sqrt(Is**2+Js**2)
+
+    # This line is used to get a set of bins that matches the logic
+    # used by np.histogram, so that this function will match the choices
+    # of bin edges that comes from the non-generalized version. This also
+    # gets us the count on the number of pixels per bin so we can calculate
+    # the threshold curve
+    n_pix, bins = np.histogram(Rs,bins=nbins)
+
+    frc = []
+    for i in range(len(bins)-1):
+        mask = t.logical_and(Rs<bins[i+1], Rs>=bins[i])
+        masked_f1 = f1 * mask[...,:,:]
+        masked_f2 = f2 * mask[...,:,:]
+        numerator = t.sqrt(calc_fidelity(masked_f1, masked_f2))
+        denominator_f1 = t.sqrt(calc_fidelity(masked_f1, masked_f1))
+        denominator_f2 = t.sqrt(calc_fidelity(masked_f2, masked_f2))
+        frc.append(numerator / t.sqrt((denominator_f1 * denominator_f2)))
+
+    frc = np.array(frc)
+
+    # This moves from combined-image SNR to single-image SNR
+    snr /= 2
+
+    threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
+        (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
+
+    if not im_np:
+        bins = t.tensor(bins)
+        frc = t.tensor(frc)
+        threshold = t.tensor(threshold)
+
+    return bins[:-1], frc, threshold
+
+

src/CDTools/tools/atoms/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.atoms.atoms import *

src/CDTools/tools/data/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.data.data import *

src/CDTools/tools/data/data.py

@@ -0,0 +1,781 @@
+"""Contains the base functions for loading and saving data from/to .cxi files
+
+These functions are used when constructing a new dataset class to pull
+specific desired information from a .cxi file. These functions should
+handle all the needed conversions between standard formats (for example,
+transposes of the basis arrays, shifting from object to probe motion, etc).
+"""
+
+import h5py
+import numpy as np
+import numbers
+import datetime
+import dateutil.parser
+import torch as t
+from contextlib import contextmanager
+
+__all__ = ['get_entry_info',
+           'get_sample_info',
+           'get_wavelength',
+           'get_detector_geometry',
+           'get_mask',
+           'get_dark',
+           'get_data',
+           'get_shot_to_shot_info',
+           'get_ptycho_translations',
+           'create_cxi',
+           'add_entry_info',
+           'add_sample_info',
+           'add_source',
+           'add_detector',
+           'add_mask',
+           'add_dark',
+           'add_data',
+           'add_shot_to_shot_info',
+           'add_ptycho_translations']
+
+
+#
+# Functions to inspect the basic attributes of a cxi file represented as an
+# h5 file object
+#
+
+
+def get_entry_info(cxi_file):
+    """Returns a dictionary with the basic metadata from the cxi file's entry_1 attribute
+
+    String type metadata is read out as a string, and datetime metadata
+    is converted to python datetime objects if the string is properly
+    formatted.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    entry_info : dict
+        A dictionary with basic metadata defined in the cxi file
+
+    """
+    e1 = cxi_file['entry_1']
+    metadata_attrs = ['title',
+                      'experiment_identifier',
+                      'experiment_description',
+                      'program_name',
+                      'start_time',
+                      'end_time']
+    metadata = {attr: str(e1[attr][()].decode()) for attr in metadata_attrs
+                if attr in e1}
+    
+
+    for attr in metadata_attrs:
+        if attr in e1:
+            try:
+                metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode()))
+            except ValueError:
+                metadata[attr] = str(e1[attr][()].decode())
+    return metadata
+
+
+def get_sample_info(cxi_file):
+    """Returns a dictionary with the basic metadata from the cxi file's entry_1/sample_1 attribute
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    sample_info : dict
+        A dictionary with basic metadata from the sample defined in the cxi file
+
+    """
+    if 'entry_1/sample_1' not in cxi_file:
+        return None
+    
+    s1 = cxi_file['entry_1/sample_1']
+    metadata_attrs = ['name','description','unit_cell_group']
+
+    metadata = {}
+    for attr in metadata_attrs:
+        # Somehow different ways of saving can lead to different ways to
+        # decode it here, so we try both
+        if attr in s1:
+            try:
+                metadata[attr] = str(s1[attr][()].decode())
+            except AttributeError as e:
+                metadata[attr] = str(np.array(s1[attr][:])[0].decode())
+
+    float_attrs = ['concentration',
+                   'mass',
+                   'temperature',
+                   'thickness',
+                   'unit_cell_volume']
+
+    for attr in float_attrs:
+        if attr in s1:
+            metadata[attr] = np.float32(s1[attr][()])
+
+    if 'unit_cell' in s1:
+        metadata['unit_cell'] = np.array(s1['unit_cell']).astype(np.float32)
+
+
+    if 'geometry_1/orientation' in s1:
+        orient = np.array(s1['geometry_1/orientation']).astype(np.float32)
+        xvec = orient[:3] / np.linalg.norm(orient[:3])
+        yvec = orient[3:] / np.linalg.norm(orient[3:])
+        metadata['orientation'] = np.array([xvec,yvec,
+                                            np.cross(xvec,yvec)])
+
+    if 'geometry_1/surface_normal' in s1:
+        snorm = np.array(s1['geometry_1/surface_normal']).astype(np.float32)
+        xvec = np.cross(np.array([0.,1.,0.]), snorm)
+        xvec /= np.linalg.norm(xvec)
+        yvec = np.cross(snorm, xvec)
+        yvec /= np.linalg.norm(yvec)
+        metadata['orientation'] = np.array([xvec, yvec, snorm])
+
+    # Check if the metadata is empty
+    if metadata == {}:
+        metadata = None
+
+    return metadata
+
+
+def get_wavelength(cxi_file):
+    """Returns the wavelength of the source defined in the cxi file object, in m
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    wavelength: np.float32
+        The wavelength of the source defined in the cxi file
+    """
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'source_1/wavelength' in i1:
+        wavelength = np.float32(i1['source_1/wavelength'])
+    elif 'source_1/energy' in i1:
+        energy = np.float32(i1['source_1/energy'])
+        wavelength = 1.9864459e-25 / energy
+    else:
+        raise KeyError('Neither Wavelength or Energy Defined in provided .cxi File')
+
+    return wavelength
+
+
+def get_detector_geometry(cxi_file):
+    """Returns a standardized description of the detector geometry defined in the cxi file object
+
+    It makes intelligent assumptions based on the definitions in the cxi
+    file definition. The standardized description of the geometry that it
+    outputs includes the sample to detector distance, the corner location
+    of the detector, and the basis vectors defining the detector. It can
+    only handle detectors defined as rectangular grids of pixels.
+
+    The distance and corner_location values are technically overdetermining
+    the detector location, but for many experiments (particularly
+    transmission experiments), the distance is needed and the exact
+    corner location is not. If the corner location is not reported in
+    the cxi file, no attempt will be made to calculate it.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    distance : np.float32
+        The sample to detector distance, in m
+    basis_vectors : np.array
+        The basis vectors for the detector
+    corner_location : np.array
+        The real-space location of the (0,0) pixel in the detector
+
+    """
+    i1 = cxi_file['entry_1/instrument_1']
+    d1 = i1['detector_1']
+
+    if 'detector_1/basis_vectors' in i1:
+        basis_vectors = np.array(d1['basis_vectors'])
+        if basis_vectors.shape == (2,3):
+            basis_vectors = basis_vectors.T
+    else:
+        # This whole thing just to account for all the ways people can
+        # implicitly define the x or y pixel size for a detector. I've
+        # seen too many of these in the wild, unfortunately...
+        try:
+            x_pixel_size = np.float32(d1['x_pixel_size'])
+        except:
+            x_pixel_size = None
+        try:
+            y_pixel_size = np.float32(d1['y_pixel_size'])
+        except:
+            y_pixel_size = None
+
+        if x_pixel_size is None and y_pixel_size is not None:
+            x_pixel_size = y_pixel_size
+        elif x_pixel_size is not None and y_pixel_size is None:
+            y_pixel_size = x_pixel_size
+        if x_pixel_size is None and y_pixel_size is None:
+            raise KeyError('Detector pixel size not defined in file.')
+        basis_vectors = np.array([[0,-y_pixel_size,0],
+                                  [-x_pixel_size,0,0]]).transpose()
+
+    try:
+        distance = np.float32(d1['distance'])
+    except:
+        distance = None
+    try:
+        corner_position = np.array(d1['corner_position'])
+    except:
+        corner_position = None
+
+    # Don't pretend to calculate corner position from distance if it's not
+    # defined, but do calculate distance from corner position if distance is
+    # not defined. If neither is defined, then raise an error.
+    if distance is None and corner_position is not None:
+        detector_normal = np.cross(basis_vectors[:,0],
+                                   basis_vectors[:,1])
+        detector_normal /= np.linalg.norm(detector_normal)
+        distance = np.linalg.norm(np.dot(corner_position, detector_normal))
+
+    if distance is None and corner_position is None:
+        raise KeyError('Neither sample to detector distance nor corner position is defined in file.')
+
+    return distance, basis_vectors, corner_position
+
+
+def get_mask(cxi_file):
+    """Returns the detector mask defined in the cxi file object
+
+    This function converts from the format specified in the cxi file
+    definition to a simple on/off mask, where a value of 1 defines a
+    good pixel (on) and a value of 0 defines a bad pixel (off).
+
+    If any bit is set in the mask at all, it will be defined as a bad
+    pixel, with the exception of pixels marked exactly as 0x00001000,
+    which is defined to mean that the pixel has signal above the
+    background. These pixels are treated as on pixels.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    mask : np.array
+        An array storing the mask from the cxi file
+    """
+
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1/mask' in i1:
+        mask = np.array(i1['detector_1/mask']).astype(np.uint32)
+        mask_on = np.equal(mask,np.uint32(0))
+        mask_has_signal = np.equal(mask,np.uint32(0x00001000))
+        return np.logical_or(mask_on,mask_has_signal).astype(bool)
+    else:
+        return None
+
+
+def get_dark(cxi_file):
+    """Returns an array with a dark image to use for initialization of a background model
+
+    This looks for a set of dark images at
+    entry_1/instrument_1/detector_1/data_dark. If the darks exist, it will
+    return the mean of the array along all axes but the last two. That is,
+    if the dark image is a single image, it will return that image. If it
+    is a stack of images, it will return the mean along the stack axis.
+
+    If the darks do not exist, it will return None.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    dark : np.array
+        An array storing the dark image
+    """
+
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1/data_dark' in i1:
+        darks = np.array(i1['detector_1/data_dark'])
+        dims = tuple(range(len(darks.shape) - 2))
+        darks = np.nanmean(darks,axis=dims)
+    else:
+        darks = None
+
+    return darks
+
+
+def get_data(cxi_file, cut_zeros = True):
+    """Returns an array with the full stack of detector data defined in the cxi file object
+
+    This function will make sure to check all the various places that it's
+    okay to store the data in, to ensure that it can find the data regardless
+    of whether the creator of the .cxi file has remembered to link the data
+    to all the required locations.
+
+    It will return the data array in whatever shape it's defined in.
+
+    It will also read out the axes attribute of the data into a list of strings.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+    cut_zeros : bool
+            Default True, whether to set all negative data to zero
+
+    Returns
+    -------
+    data : np.array
+        An array storing the data defined in the cxi file
+    axes : list(str)
+        A list of the axes defined in the axes attribute, if any
+    """
+
+    # Possible locations for the data
+    if 'entry_1/data_1/data' in cxi_file:
+        pull_from = 'entry_1/data_1/data'
+    elif 'entry_1/instrument_1/detector_1/data' in cxi_file:
+        pull_from = 'entry_1/instrument_1/detector_1/data'
+    else:
+        raise KeyError('Data is not defined within cxi file')
+
+    data = cxi_file[pull_from][:]
+    # Use maximum in-place to avoid allocating any more memory than is needed
+    if cut_zeros:
+        np.maximum(data,0,data)
+
+    if 'axes' in cxi_file[pull_from].attrs:
+        try:
+            axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
+        except AttributeError as e: # Weird string vs bytes thing, ehhh
+            axes = str(cxi_file[pull_from].attrs['axes']).split(':')
+
+        axes = [axis.strip().lower() for axis in axes]
+    else:
+        axes = None
+
+    return data, axes
+
+
+def get_shot_to_shot_info(cxi_file, field_name):
+    """Gets a specified dataset of shot-to-shot information from the cxi file
+
+    The data is assumed to be in the form of an array, with one dimension
+    being the number of patterns being stored in the dataset. This is
+    helpful for storing additional readback data on the shot-to-shot
+    level that may be important but doens't have a clearly defined
+    place to be stored in the .cxi file specification. Such data includes
+    shot-to-shot probe intensity measurements, polarizer positions, etc.
+
+    It will look for this data in 3 places (in the following order):
+
+    1) entry_1/data_1/<field_name>
+    2) entry_1/sample_1/geometry_1/<field_name>
+    3) entry_1/instrument_1/detector_1/<field_name>
+
+    This function is also used internally to read out the translations
+    associated with a ptychography experiment
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+    field_name : str
+        The name of the field to be read from
+
+    Returns
+    -------
+    data : np.array
+        An array storing the translations defined in the cxi file
+    """
+    if 'entry_1/data_1/' + field_name in cxi_file:
+        pull_from = 'entry_1/data_1/' + field_name
+    elif 'entry_1/sample_1/geometry_1/' + field_name in cxi_file:
+        pull_from = 'entry_1/sample_1/geometry_1/' + field_name
+    elif 'entry_1/instrument_1/detector_1/' in cxi_file:
+        pull_from = 'entry_1/instrument_1/detector_1/' + field_name
+    else:
+        raise KeyError('Data is not defined within cxi file')
+
+    return np.array(cxi_file[pull_from]).astype(np.float32)
+
+
+def get_ptycho_translations(cxi_file):
+    """Gets an array of x,y,z translations, if such an array has been defined in the file
+
+    It negates the translations, because the CXI file format is designed
+    to specify translations of the samples and the CDTools code specifies
+    translations of the optics.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    translations : np.array
+        An array storing the translations defined in the cxi file
+    """
+
+    translations = get_shot_to_shot_info(cxi_file, 'translation')
+    return -translations
+
+
+#
+# Now we move on to the helper functions for writing CXI files
+#
+
+
+def create_cxi(filename):
+    """Creates a new cxi file with a single entry group
+
+    Parameters
+    ----------
+    filename : str
+        The path at which to create the file
+    """
+    file_obj = h5py.File(filename,'w')
+    file_obj.create_dataset('cxi_version', data=160)
+    file_obj.create_dataset('number_of_entries',data=1)
+    e1f = file_obj.create_group('entry_1')
+    return file_obj
+
+
+def add_entry_info(cxi_file, metadata):
+    """Adds a dictionary of entry metadata to the entry_1 group of a cxi file object
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the info to
+    metadata : dict
+        A dictionary containing all the metadata to be stored
+    """
+    # Just the string and datetime types should be relevant but all are
+    # included in case the cxi spec becomes more permissive
+    for key, value in metadata.items():
+        if isinstance(value,(str,bytes)):
+            cxi_file['entry_1'][key] = np.string_(value)
+        elif isinstance(value, datetime.datetime):
+            cxi_file['entry_1'][key] = np.string_(value.isoformat())
+        elif isinstance(value, numbers.Number):
+            cxi_file['entry_1'][key] = value
+        elif isinstance(value, (np.ndarray,list,tuple)):
+            cxi_file['entry_1'].create_dataset(key, data=np.asarray(value))
+        elif isinstance(value, t.Tensor):
+            asnumpy = value.detach().cpu().numpy()
+            cxi_file['entry_1'].create_dataset(key, data=asnumpy)
+
+
+
+def add_sample_info(cxi_file, metadata):
+    """Adds a dictionary of entry metadata to the entry_1/sample_1 group of a cxi file object
+
+    This function will create the sample_1 attribute if it doesn't already exist
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the info to
+    metadata : dict
+        A dictionary containing all the metadata to be stored
+    """
+    if 'entry_1/sample_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('sample_1')
+    s1 = cxi_file['entry_1/sample_1']
+
+    if 'orientation' in metadata:
+        if 'geometry_1' not in s1:
+            s1.create_group('geometry_1')
+        # Only store the part of this matrix as defined in the CXI file spec
+        s1['geometry_1'].create_dataset('orientation',
+                                        data=metadata['orientation'].ravel()[:6])
+
+    for key, value in metadata.items():
+        if key == 'orientation':
+            continue # this is a special case
+        if isinstance(value,(str,bytes)):
+            s1[key] = np.string_(value)
+        elif isinstance(value, datetime.datetime):
+            s1[key] = np.string_(value.isoformat())
+        elif isinstance(value, numbers.Number):
+            s1[key] = value
+        elif isinstance(value, (np.ndarray,list,tuple)):
+            s1.create_dataset(key, data=np.asarray(value))
+        elif isinstance(value, t.Tensor):
+            asnumpy = value.detach().cpu().numpy()
+            s1.create_dataset(key, data=asnumpy)
+
+
+def add_source(cxi_file, wavelength):
+    """Adds the entry_1/source_1 group to a cxi file object
+
+    It stores the energy and wavelength attributes in the source_1 group,
+    given a wavelength to define them from.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the source to
+    wavelength : float
+        The wavelength of light
+    """
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'source_1' not in i1:
+        i1.create_group('source_1')
+    s1 = i1['source_1']
+    s1['wavelength'] = np.float32(wavelength)
+    s1['energy'] = np.float32(1.9864459e-25 / wavelength)
+
+
+
+def add_detector(cxi_file, distance, basis, corner=None):
+    """Adds the entry_1/instrument_1/detector_1 group to a cxi file object
+
+    It will define all the relevant parameters - distance, pixel size,
+    detector basis, and corner position (if relevant) based on the provided
+    information.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the detector to
+    distance : float
+        The sample to detector distance
+    basis : array
+        The detector basis
+    corner : array
+        Optional, the corner position of the detector
+
+    """
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    d1 = i1['detector_1']
+
+    d1['distance'] = np.float32(distance)
+
+    if isinstance(basis, t.Tensor):
+        basis = basis.detach().cpu().numpy()
+    if basis.shape == (2,3):
+        basis = basis.T
+    d1['x_pixel_size'] = np.linalg.norm(basis[:,1])
+    d1['y_pixel_size'] = np.linalg.norm(basis[:,0])
+    d1.create_dataset('basis_vectors', data=basis)
+
+    if corner is not None:
+        if isinstance(corner, t.Tensor):
+            corner = corner.detach().cpu().numpy()
+        d1.create_dataset('corner_position',data=corner)
+
+
+def add_mask(cxi_file, mask):
+    """Adds the specified mask to the cxi file
+
+    It places the mask into the mask dataset under
+    entry_1/instrument_1/detector_1. The internal mask is defined
+    simply as a 1 for an "on" pixel and a 0 for an "off" pixel, and
+    the saved mask is exactly the opposite. This is simpler than the
+    most general mask allowed by the cxi file format but it captures the
+    distinction between pixels to be used and pixels not to be used.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the mask to
+    mask : array
+        The mask to save out to the file
+    """
+
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    d1 = i1['detector_1']
+    if isinstance(mask, t.Tensor):
+        mask = mask.detach().cpu().numpy()
+
+    mask_to_save = np.zeros(mask.shape).astype(np.uint32)
+    mask_to_save[mask == 0] = 1
+    d1.create_dataset('mask',data=mask_to_save)
+
+
+def add_dark(cxi_file, dark):
+    """Adds the specified dark image to a cxi file
+
+    It places the dark image data into the data_dark dataset under
+    entry_1/instrument_1/detector_1.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the mask to
+    dark : array
+        The dark image(s) to save out to the file
+    """
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    d1 = i1['detector_1']
+    if isinstance(dark, t.Tensor):
+        dark = dark.detach().cpu().numpy()
+
+    d1.create_dataset('data_dark',data=dark)
+
+
+def add_data(cxi_file, data, axes=None, compression='gzip',
+             chunks=True):
+    """Adds the specified data to the cxi file
+
+    It will add the data unchanged to the file, placing it in two spots:
+
+    1) The entry_1/instrument_1/detector_1/data path
+    2) A softlink at entry_1/data_1/data
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the data to
+    data : array
+        The data to be saved
+    axes : list(str)
+        Optional, a list of axis names to be saved in the axes attribute
+    """
+    if 'entry_1/data_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('data_1')
+    data1 = cxi_file['entry_1/data_1']
+
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    det1 = i1['detector_1']
+
+    if isinstance(data, t.Tensor):
+        data = data.detach().cpu().numpy()
+
+    det1.create_dataset('data', data=data, compression=compression,
+                        chunks=chunks)
+    data1['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
+
+    if axes is not None:
+        if isinstance(axes, list):
+            axes_str = ':'.join(axes)
+        else:
+            axes_str = str(axes)
+        det1['data'].attrs['axes'] = np.string_(axes_str)
+
+        
+def add_shot_to_shot_info(cxi_file, data, field_name):
+    """Adds a specified dataset of shot-to-shot information to the cxi file
+
+    The data is assumed to be in the form of an array, with one dimension
+    being the number of patterns being stored in the dataset. This is
+    helpful for storing additional readback data on the shot-to-shot
+    level that may be important but doens't have a clearly defined
+    place to be stored in the .cxi file specification. Such data includes
+    shot-to-shot probe intensity measurements, polarizer positions, etc.
+
+    This function is also used internally to store the translations
+    associated with a ptychography experiment
+    
+    It will store this data in 3 places:
+
+    1) The entry_1/sample_1/geometry_1/<field_name> path
+    2) A softlink at entry_1/data_1/<field_name>
+    3) A softlink at entry_1/instrument_1/detector_1/<field_name>
+
+    The geometry and detector paths may not always be relevant, but this
+    ensures that the data is always available in any of the places that
+    an eventual reader may go to look for, e.g., the translations.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the translations to
+    data : array
+        The data to be saved
+    field_name : str
+        The field name to save the data under
+    """
+    
+    if 'entry_1/sample_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('sample_1')
+    s1 = cxi_file['entry_1/sample_1']
+
+    if 'geometry_1' not in s1:
+        s1.create_group('geometry_1')
+    g1 = s1['geometry_1']
+
+    if 'entry_1/data_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('data_1')
+    data1 = cxi_file['entry_1/data_1']
+
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    det1 = i1['detector_1']
+
+
+    if isinstance(data, t.Tensor):
+        data = data.detach().cpu().numpy()
+
+
+    g1.create_dataset(field_name, data=data)
+    data1[field_name] = h5py.SoftLink('/entry_1/sample_1/geometry_1/'
+                                      + field_name)
+    det1[field_name] = h5py.SoftLink('/entry_1/sample_1/geometry_1/'
+                                      + field_name)
+
+
+def add_ptycho_translations(cxi_file, translations):
+    """Adds the specified translations to the cxi file
+
+    It will add the translations to the file, negating them to conform to
+    the standard in cxi files that the translations refer to the object's
+    translation.
+
+    It will store them in 3 places:
+
+    1) The entry_1/sample_1/geometry_1/translation path
+    2) A softlink at entry_1/data_1/translation
+    3) A softlink at entry_1/instrument_1/detector_1/translation
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the translations to
+    translations : array
+        The translations to be saved
+    """
+    # accounting for the different definition between cxi files and
+    # CDTools
+    translations = -translations
+
+    add_shot_to_shot_info(cxi_file, translations, 'translation')

src/CDTools/tools/image_processing/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.image_processing.image_processing import *

src/CDTools/tools/image_processing/image_processing.py

@@ -0,0 +1,319 @@
+"""Contains functions for basic image processing needs
+
+This module contains two kinds of image processing tools. The first type
+is specific tools for calculating commonly needed metrics (such as the
+centroid of an image), directly on complex-valued torch tensors. The second
+kind of tools perform common image manipulations on torch tensors, in such
+a way that it is safe to include them in automatic differentiation models.
+"""
+
+import numpy as np
+import torch as t
+from CDTools.tools import propagators
+
+__all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift',
+           'find_subpixel_shift', 'find_pixel_shift', 'find_shift',
+           'convolve_1d', 'fourier_upsample']
+
+
+def centroid(im, dims=2):
+    """Returns the centroid of an image or a stack of images
+
+    By default, the last two dimensions are used in the calculation
+    and the remainder of the dimensions are passed through.
+
+    Beware that the meaning of the centroid is not well defined if your
+    image contains values less than 0
+
+    Parameters
+    ----------
+    im : torch.Tensor
+        An image or stack of images to calculate from
+    dims : int 
+        Default 2, how many trailing dimensions to calculate with
+    
+    Returns
+    -------
+    centroid : torch.Tensor
+        An (i,j) index or stack of indices
+    """
+    # For some reason this needs to be a list
+    indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)]
+    indices = t.meshgrid(*indices, indexing='ij')
+
+    use_dims = [-dims+i for i in range(dims)]
+    divisor = t.sum(im, dim=use_dims)
+    centroids = [t.sum(index * im, dim=use_dims) / divisor
+                 for index in indices]
+
+    return t.stack(centroids,dim=-1)
+
+
+def centroid_sq(im, dims=2, comp=False):
+    """Returns the centroid of the square of an image or stack of images
+
+    By default, the last two dimensions are used in the calculation
+    and the remainder of the dimensions are passed through.
+
+    If the "comp" flag is set, it will be assumed that the last dimension
+    represents the real and imaginary part of a complex number, and the
+    centroid will be calculated for the magnitude squared of those numbers
+
+    Parameters
+    ----------
+    im : torch.Tensor
+        An image or stack of images to calculate from
+    dims : int 
+        Default 2, how many trailing dimensions to calculate for
+    comp : bool
+        Default is False, whether the data represents complex numbers
+
+    Returns
+    -------
+    centroid: torch.Tensor
+        An (i,j) index or stack of indices
+    """
+    if comp:
+        im_sq = t.abs(im)**2
+    else:
+        im_sq = im**2
+
+    return centroid(im_sq, dims=dims)
+
+
+def sinc_subpixel_shift(im, shift):
+    """Performs a subpixel shift with sinc interpolation on the given tensor
+
+    The subpixel shift is done circularly via a multiplication with a linear
+    phase mask in Fourier space.
+
+    Parameters
+    ----------
+    im : torch.Tensor
+        A complex-valued tensor to perform the subpixel shift on
+    shift : array
+        A length-2 array describing the shift to perform, in pixels
+
+    Returns
+    -------
+    shifted_im : torch.Tensor
+        The subpixel shifted tensor
+    """
+
+    i = t.arange(im.shape[0]) - im.shape[0]//2
+    j = t.arange(im.shape[1]) - im.shape[1]//2
+    I,J = t.meshgrid(i,j, indexing='ij')
+    I = 2 * np.pi * I.to(t.float32) / im.shape[0]
+    J = 2 * np.pi * J.to(t.float32) / im.shape[1]
+    I = I.to(dtype=im.dtype,device=im.device)
+    J = J.to(dtype=im.dtype,device=im.device)
+
+    fft_im = t.fft.fftshift(t.fft.fft2(im),dim=(-2,-1))
+    shifted_fft_im = fft_im * t.exp(1j * (-shift[0]*I - shift[1]*J))
+    return t.fft.ifft2(t.fft.ifftshift(shifted_fft_im, dim=(-2,-1)))
+
+
+
+def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
+    """Calculates the subpixel shift between two images by maximizing the autocorrelation
+
+    This function only searches in a 2 pixel by 2 pixel box around the
+    specified search_around parameter. The calculation is done using the
+    approach outlined in "Efficient subpixel image registration algorithms",
+    Optics Express (2008) by Manual Guizar-Sicarios et al.
+
+    Parameters
+    ----------
+    im1 : torch.Tensor
+        The first real or complex-valued torch tensor
+    im2 : torch.Tensor
+        The second real or complex-valued torch tensor
+    search_around : array
+        Default (0,0), the shift to search in the vicinity of
+    resolution : int
+        Default is 10, the fraction of a pixel to calculate to
+
+    Returns
+    -------
+    shift : torch.Tensor
+        The relative shift (i,j) needed to best map im1 onto im2
+    """
+
+    #
+    # Here's my approach, perhaps it's a little unconventional. I will first
+    # calculate the phase correlation function as found in ____ (cite a paper
+    # defining it). This is strongly peaked, so I can take a small window
+    # of say, 10x10 pixels, and then do a sinc interpolation of that area
+    # using an FFT with upsampling by a factor of resolution in reciprocal
+    # space
+    #
+    cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
+
+    # Not sure if this is more or less stable than just the correlation
+    # maximum - requires some testing
+    cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
+
+
+    # Now, I need to shift the array to pull out a contiguous window
+    # around the correlation maximum
+    try:
+        search_around = search_around.cpu()
+    except:
+        search_around = t.as_tensor(search_around)
+
+    window_size = 15
+    shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
+    cor_window = t.roll(cor, shift_zero, dims=(-2,-1))[...,:2*window_size,:2*window_size]
+
+    # Now we upsample this window
+    cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1))
+    upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape) * resolution),
+                        dtype=cor.dtype,device=cor.device)
+
+    upsampled[...,:2*window_size,:2*window_size] = cor_window_fft
+    upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1))
+    upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2,
+                       (-window_size*resolution,-window_size*resolution),
+                       dims=(0,1))
+
+
+    # And we extract the shift from the window
+    sh = t.as_tensor(upsampled.shape, device=upsampled.device)
+    cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
+                                rounding_mode='floor'),
+                          t.argmax(upsampled) % sh[1]],
+                         device=upsampled.device)
+
+    sh_over_2 = t.div(sh,2,rounding_mode='floor')
+    subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
+
+    return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
+        subpixel_shift / resolution
+
+
+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. This corresponds
+    to the amount that im1 would have to be shifted by to line up best with
+    im2
+
+    Parameters
+    ----------
+    im1 : torch.Tensor
+        The first real or complex-valued torch tensor
+    im2 : torch.Tensor
+        The second real or complex-valued torch tensor
+
+    Returns
+    -------
+    shift : torch.Tensor
+        The integer-valued shift (i,j) that best maps im1 onto im2 
+    """
+    cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
+
+    # Not sure if this is more or less stable than just the correlation
+    # maximum - requires some testing
+    cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft)))
+
+
+    sh = t.as_tensor(cor.shape,device=im1.device)
+    cormax = t.tensor([t.div(t.argmax(cor),sh[1],rounding_mode='floor'),
+                       t.argmax(cor) % sh[1]]).to(device=im1.device)
+
+    sh_over_2 = t.div(sh,2,rounding_mode='floor')
+    return (cormax + sh_over_2) % sh - sh_over_2
+
+
+
+def find_shift(im1, im2, resolution=10):
+    """Calculates the shift between two images by maximizing the autocorrelation
+
+    This function starts by calculating the maximum shift to integer
+    pixel resolution, and then searchers the nearby area to calculate a
+    subpixel shift
+
+    Parameters
+    ----------
+    im1 : torch.Tensor
+        The first real or complex-valued torch tensor
+    im2 : torch.Tensor
+        The second real or complex-valued torch tensor
+    resolution : int
+        Default is 10, the fraction of a pixel to calculate to
+
+    Returns
+    -------
+    shift : torch.Tensor
+        The relative shift (i,j) needed to best map im1 onto im2
+    """
+    integer_shift = find_pixel_shift(im1,im2)
+    subpixel_shift = find_subpixel_shift(im1, im2, search_around=integer_shift,
+                                         resolution=resolution)
+
+    return subpixel_shift
+
+
+def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
+    """Convolves an image with a 1d kernel along a specified dimension
+
+    The convolution is a circular convolution calculated using a Fourier
+    transform. The calculation is done so the input remains differentiable
+    with respect to the output.
+
+    If the image has a final dimension of 2, it is assumed to be complex.
+    Otherwise, the image is assumed to be real. The image and kernel
+    must either both be real or both be complex.
+    
+    Parameters
+    ----------
+    image : torch.Tensor
+        The image to convolve
+    kernel : torch.Tensor
+        The 1d kernel to convolve with
+    dim : int
+        Default 0, the dimension to convolve along
+    fftshift_kernel : bool
+        Default True, whether to fftshift the kernel first.
+    
+    Returns
+    -------
+    convolved_im : torch.Tensor
+        The convolved image
+    """
+
+    
+    if fftshift_kernel:
+        kernel = t.fft.ifftshift(kernel,dim=(-1,))
+
+    # We have to transpose the relevant dimension to -1 before using the fft,
+    # which expects to operate on the final dimension
+    trans_im = t.transpose(image, dim, -1)
+        
+    # Take a correlation
+    fft_im = t.fft.fft(trans_im)
+    fft_kernel = t.fft.fft(kernel)
+    trans_conv = t.fft.ifft(fft_im * fft_kernel)
+
+    conv_im = t.transpose(trans_conv, dim, -1)
+
+    return conv_im
+
+
+def fourier_upsample(ims, preserve_mean=False):
+    # If preserve_mean is true, it preserves the mean pixel intensity
+    # otherwise, it preserves the total summed intensity
+    upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]),
+                           dtype=ims.dtype,
+                           device=ims.device)
+    left = [ims.shape[-2]//2,ims.shape[-1]//2]
+    right = [ims.shape[-2]//2+ims.shape[-2],
+             ims.shape[-1]//2+ims.shape[-1]]
+    
+    upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims)
+    if preserve_mean:
+        upsampled *= 2
+    return propagators.inverse_far_field(upsampled)
+
+

src/CDTools/tools/initializers/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.initializers.initializers import *

src/CDTools/tools/initializers/initializers.py

@@ -0,0 +1,659 @@
+"""Contains functions to sensibly initialize reconstructions
+
+The functions in this module both do the geometric calculations needed to
+initialize the reconstrucions, and the heuristic calculations for
+geierating sensible initializations for the probe guess.
+"""
+
+import math
+import numpy as np
+import torch as t
+from CDTools.tools.propagators import *
+from CDTools.tools.analysis import orthogonalize_probes
+from CDTools.tools import image_processing
+from scipy.fftpack import next_fast_len
+from scipy.sparse import linalg as spla
+from torch.nn.functional import pad
+import numpy as np
+from functools import *
+
+__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
+           'gaussian_probe', 'SHARP_style_probe', 'STEM_style_probe',
+           'RPI_spectral_init',
+           'generate_subdominant_modes'] 
+
+def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1):
+    """Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
+    
+    It takes in the parameters for a given detector - the basis defining
+    the pixel pitch and the shape, as well as the wavelength and propagation
+    distance. Optionally, it accepts a defined "center", the zero-frequency
+    pixel's location. It then will automatically define a larger detector
+    if necessary, define the exit wave basis associated with a far-field
+    diffraction experiment, and return that basis, shape, and detector slice
+    
+    Parameters
+    ----------
+    det_basis : torch.Tensor
+        The detector basis, as defined elsewhere
+    det_shape : torch.Size
+        The (i,j) shape of the detector
+    wavelength : float)
+        The wavelength of light for the experiment, in m
+    distance : float
+        The sample-detector distance, in m
+    center : torch.Tensor
+        If defined, the location of the zero frequency pixel
+    opt_for_fft : bool
+        Default is true, whether to increase detector size to improve fft performance
+    padding : int
+        Default is 0, the size of an extra border of nonphysical pixels around the detector
+    oversampling : int
+        Default is 1, the amount to multiply the exit wave shape by.
+    
+    Returns
+    -------
+    basis : torch.Tensor
+        The exit wave basis
+    shape : torch.Tensor
+        The exit wave's shape
+    slice : slice
+        The slice corresponding to the physical detector
+    """
+    
+    det_shape = t.as_tensor(tuple(det_shape), dtype=t.int32)
+    det_basis = t.as_tensor(det_basis)
+    # First, set the center if it's not already specified
+    # This definition matches the center pixel of an fftshifted array
+    if center is None:
+        # this is center//2, but in pytorch 1.9.0 it throws a warning
+        # if you just do that
+
+        center = t.div(det_shape,2,rounding_mode='floor')# // 2
+    else:
+        center = t.as_tensor(center, dtype=t.int32)
+        
+    # Then, calculate the required detector size from the centering
+    # This is a bit opaque but was worth doing accurately
+    min_left = center * 2
+    min_right = (det_shape - center) * 2 - 1
+    full_shape = t.max(min_left,min_right).to(t.int32) + 2 * padding
+
+    # In some edge cases this shape can be smaller than the detector shape
+    full_shape = t.max(full_shape, det_shape)
+
+    
+    if opt_for_fft:
+        full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape],
+                                 dtype=t.int32)
+    
+    # Then, generate a slice that pops the actual detector from the full
+    # detector shape
+    # this is full_center//2, but in pytorch 1.9.0 it throws a warning
+    # if you just do that
+    full_center = t.div(full_shape,2, rounding_mode='floor')# // 2
+    det_slice = np.s_[int(full_center[0]-center[0]):
+                      int(full_center[0]-center[0]+det_shape[0]),
+                      int(full_center[1]-center[1]):
+                      int(full_center[1]-center[1]+det_shape[1])]
+
+    
+    # Finally, generate the basis for the exit wave in real space
+
+    # This method should work for a general parallelogram
+    # shaped detector
+    det_shape = det_basis * full_shape.to(t.float32)
+    pinv_basis = t.tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32)
+    real_space_basis = pinv_basis * wavelength * distance
+
+    # This is definitely correct, but less simple. Included here
+    # So future me can check that both versions are consistent.
+    #oop_dir = np.cross(det_basis[:,0],det_basis[:,1])
+    #oop_dir /= np.linalg.norm(oop_dir)
+    #full_basis = np.array([np.array(det_basis[:,0]),np.array(det_basis[:,1]),oop_dir]).transpose()
+    #inv_basis = t.tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32)
+    #real_space_basis = inv_basis*wavelength * distance / \
+    #    full_shape.to(t.float32)
+
+    # Finally, convert the shape back to a torch.Size
+    full_shape = t.Size([dim * oversampling for dim in full_shape]) 
+
+
+    return real_space_basis, full_shape, det_slice
+
+
+def calc_object_setup(probe_shape, translations, padding=0):
+    """Returns an object shape and minimum pixel translation
+
+    Based on the given pixel-space translations, it will calculate the
+    required size for an object array and calculate the pixel translation
+    that corresponds to a shift by (0,0) of the probe. 
+    
+    Optionally a small extra border can be defined via the padding
+    attribute. If this is done, the calculated pixel translation will
+    correspond to (padding,padding)
+    
+    Parameters
+    ----------
+    probe_shape : torch.Size
+        The size of the probe array
+    translations : torch.Tensor
+        A Jx2 stack of pixel-valued (i,j) translations
+    padding : int
+        Optional, the size of an extra border to include
+    
+    Returns
+    -------
+    obj_shape : torch.Size
+        The minimum required size of the object array
+    min_translation : torch.Tensor
+        The minimum pixel-valued translation
+    """
+
+    # First we look at the translations to find the minimum translation
+    # and the range of translations
+    min_translation = t.min(translations, dim=0)[0]
+    translation_range = t.max(translations, dim=0)[0] - min_translation
+    
+    # Calculate the required shape
+    translation_range = t.ceil(translation_range).numpy().astype(np.int32)
+    shape = translation_range + np.array(probe_shape) + 2 * padding
+    shape = t.Size(shape)
+    
+    # And the minimum translation
+    min_translation = min_translation - padding
+    
+    return shape, min_translation
+    
+    
+
+def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
+    """Returns an array with a centered Gaussian
+
+    Takes in the shape and standard deviation of a gaussian
+    and returns a complex torch tensor (trailing dimension is 2) with
+    values corresponding to a two-dimensional gaussian function
+
+    Note that [0, 0] is taken to be at the upper left corner of the array.
+    Default is centered at ((shape[0]-1)/2, (shape[1]-1)/2)) because x and y are zero-indexed.
+    By default, the phase is uniformly 0, however a curvature can be
+    specified to simulate a probe that has been propagated a known distance
+    from it's focal point. The curvature is implemented by adding a quadratic
+    phase phi = exp(i*curvature/2 r^2) to the Gaussian
+
+    Parameters
+    ----------
+    shape : array
+        A 1x2 array specifying the dimensions of the output array in the form (i shape, j shape)
+    sigma : array
+        A 1x2 array specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
+    amplitude : float
+        Default 1, the amplitude the gaussian to simulate
+    center : array
+        Optional, a 1x2 array specifying the location of the center of the gaussian (i center, j center)
+    curvature : array
+        Optional, a complex part to add to the gaussian coefficient
+
+    Returns
+    -------
+    torch.Tensor 
+        The complex-style tensor storing the Gaussian
+    """
+    if center is None:
+        center = ((shape[0]-1)/2, (shape[1]-1)/2)
+        
+    i, j = np.mgrid[:shape[0], :shape[1]]
+    isq = (i - center[0])**2
+    jsq = (j - center[1])**2
+    result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \
+        (1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq)
+    return t.as_tensor(amplitude*result,dtype=t.complex64)
+
+
+
+def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0, polarized=False, left_polarized=True):
+    """Initializes a gaussian probe based on experimental parameters
+
+    This function generates a gaussian probe initialization which has a
+    total fluence matching the order of magnitude of the intensity in
+    the observed dataset, provided the object function is of order 1.
+    
+    The initialization is done using parameters defined in physical units,
+    such as sigma (in meters) and the propagation distance (in meters).
+    The internal conversion to pixel space is done with a provided probe
+    basis and probe shape.
+    
+    TODO: Should be updated to accept a mask
+    
+    Sigma can be provided either as a scalar for a uniform beam, or as 
+    an iterable of length 2 with [sigma_i, sigma_j] being the components
+    of sigma in the directions parallel to the i and j basis vectors of
+    the probe basis
+    
+    Parameters
+    ----------
+    dataset : Ptycho_2D_Dataset
+        The dataset whose intensity we want to match
+    basis : array
+        The real space basis for exit waves in our experiment
+    shape : array
+        The shape of the simulated real space arrays
+    sigma : array
+        The standard deviation of the probe at it's focus
+    propagation_distance : float
+        Default 0, a distance to propagate the gaussian from it's focus
+    
+    Returns
+    -------
+    torch.Tensor
+        The complex-style tensor storing the Gaussian
+    """
+    # First, we want to generate the parameters (sigma and curvature) for the
+    # propagated gaussian. Ignore the purely z-dependent phases
+    wavelength = dataset.wavelength
+    if propagation_distance is  None:
+        propagation_distance = 0
+    
+    z = propagation_distance # for shorthand
+    
+    sigma = np.array(sigma)
+    k = 2 * np.pi / wavelength
+    zr = k * sigma**2
+    sigmaz = sigma * np.sqrt(1 + (z / zr)**2)
+    curvature = -k * z / (z**2 + zr**2)
+
+    # The conversion must then be done to pixel space
+    sigma_pix = sigmaz / np.array([np.linalg.norm(basis[:,0]),
+                                   np.linalg.norm(basis[:,1])])
+    curvature_pix = curvature * np.array([np.linalg.norm(basis[:,0]),
+                                          np.linalg.norm(basis[:,1])])**2
+
+    # Then we can generate the gaussian array
+    probe = gaussian(shape, sigma=sigma_pix, curvature=curvature_pix)
+        
+    # Finally, we should calculate the average pattern intensity from the
+    # dataset and normalize the gaussian probe. This should be done by
+    if not polarized:
+        avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
+
+    else:
+        polarizer = dataset.polarizer.tolist()
+        analyzer = dataset.analyzer.tolist()
+        factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if (abs(polarizer[idx] - analyzer[idx]) > 5)]
+        avg_intensities = [t.sum(dataset[idx][1]) / factor[idx] for idx in range(len(dataset))]
+
+    avg_intensity = t.mean(t.tensor(avg_intensities))
+    probe_intensity = t.sum(t.abs(probe)**2)
+    probe = t.sqrt(avg_intensity / probe_intensity) * probe
+
+    if polarized:
+        if left_polarized == True:
+            x = 1j
+        else:
+            x = -1j
+        probe = t.stack((probe.to(dtype=t.cfloat), x * probe.to(dtype=t.cfloat)), dim=-4)
+        return probe[..., None, :, :]
+
+    return probe
+
+
+
+
+def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, oversampling=1, polarized=False, left_polarized=True):
+    """Generates a SHARP style probe guess from a dataset
+
+    What we call the "SHARP" style probe guess is to take a mean of all
+    the diffraction patterns and use that as an initial guess of the
+    Fourier space distribution of the probe. We set all the phases to
+    zero, which would for many simple beams (like a zone plate) generate
+    a first guess of the probe that is very close to the focal spot of
+    the probe beam.
+
+    If the probe is simulated in higher resolution than the detector,
+    a common occurence, these undefined pixels are set to zero for the
+    purposes of defining the guess
+
+    We make a small tweak to this procedure to lower the central pixel of
+    the probe generated this way, which can often overwhelm the rest of the
+    probe if there is significant noise on the detector
+        
+    Parameters
+    ----------
+    dataset : Ptycho_2D_Dataset
+        The dataset to work from
+    shape : torch.Size
+        The size of the probe array to simulate
+    det_slice : slice
+        A slice or tuple of slices corresponding to the detector region in Fourier space
+    propagation_distance : float
+        Default is no propagation, an amount to propagate the guessed probe from it's focal point
+    oversampling : int 
+        Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
+ 
+    Returns
+    -------
+    torch.Tensor
+        The complex-style tensor storing the probe guess
+    """
+
+    # NOTE: I don't love the way np and torch are mixed here, I think this
+    # function deserves some love.
+
+    # to use the mask or not?
+    intensities = np.zeros([dim // oversampling for dim in shape])
+    
+    if polarized:
+        factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if abs(polarizer[idx] - analyzer[idx]) > 5]
+    else:
+        factors = [1 for idx in range(len(dataset))]
+
+
+    for params, im in dataset:
+        if hasattr(dataset,'mask') and dataset.mask is not None:
+            intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy() / factors[params[0]]
+        else:
+            intensities[det_slice] += im.cpu().numpy() / params[factors[0]]       
+    intensities /= len(dataset)
+
+
+
+    # Subtract off a known background if it's stored
+    if hasattr(dataset, 'background') and dataset.background is not None:
+        intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None)
+    
+    probe_fft = t.tensor(np.sqrt(intensities)).to(dtype=t.complex64)
+
+    probe_guess = inverse_far_field(probe_fft).numpy()
+    # Now we remove the central pixel
+    #center = np.array(probe_guess.shape) // 2
+    
+    # I'm always unsure whether to use this modification:
+    
+    #probe_guess[center[0], center[1]]=np.mean([
+    #    probe_guess[center[0]-1, center[1]],
+    #    probe_guess[center[0]+1, center[1]],
+    #    probe_guess[center[0], center[1]-1],
+    #    probe_guess[center[0], center[1]+1]])
+
+    probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
+
+    if propagation_distance is not None:
+        # First generate the propagation array
+
+        probe_shape = t.as_tensor(tuple(probe_guess.shape))
+
+        # Start by recalculating the probe basis from the given information
+        det_basis = t.as_tensor(dataset.detector_geometry['basis'])
+        basis_dirs = det_basis / t.norm(det_basis, dim=0)
+        distance = dataset.detector_geometry['distance']
+        probe_basis = basis_dirs * dataset.wavelength * distance / \
+            (probe_shape * t.norm(det_basis,dim=0))
+
+        # Then package everything as it's needed
+        probe_spacing = t.norm(probe_basis,dim=0).numpy()
+        probe_shape = probe_shape.numpy().astype(np.int32)
+
+        # And generate the propagator
+        AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
+
+        probe_guess = near_field(probe_guess,AS_prop)
+    
+    # Finally, place this probe in a full-sized array if there is oversampling
+    final_probe = t.zeros(shape,dtype=t.complex64)
+    left = shape[0]//2 - probe_guess.shape[0] // 2
+    top = shape[1]//2 - probe_guess.shape[1] // 2 
+    final_probe[left:left+probe_guess.shape[0],
+                top:top+probe_guess.shape[1]] = probe_guess
+
+    if polarized:
+        if left_polarized:
+            x = 1j
+        else:
+            x = -1j
+        final_probe = t.stack((final_probe.to(dtype=t.cfloat), final_probe.to(dtype=t.cfloat)), dim=-4)
+        final_probe = final_probe[..., None, :, :]
+    
+    return final_probe
+
+
+def STEM_style_probe(dataset, shape, det_slice, convergence_semiangle, propagation_distance=None, oversampling=1):
+    """Generates a STEM style probe guess from a dataset
+
+    What we call the "STEM" style probe guess is a probe generated by
+    a uniform aperture in Fourier space, with no optical aberrations.
+    This is the kind of probe than an ideal, aberration-corrected STEM
+    with perfect coherence would produce, so it is a good starting
+    guess for STEM datasets.
+
+    We set the initial intensity of the probe with relation to the
+    diffraction patterns so that a typical object will have an intensity
+    of around 1. We also set the initial phase ramp of the probe such that
+    the undiffracted probe has a centroid on the detector which matches the
+    centroid of the diffraction patterns
+    
+    Parameters
+    ----------
+    dataset : Ptycho_2D_Dataset
+        The dataset to work from
+    shape : torch.Size
+        The size of the probe array to simulate
+    det_slice : slice
+        A slice or tuple of slices corresponding to the detector region in Fourier space
+    convergence_angle : float
+        The convergence angle of the probe, in mrad.
+    propagation_distance : float
+        Default is no propagation, an amount to propagate the guessed probe from it's focal point
+    oversampling : int 
+        Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
+ 
+    Returns
+    -------
+    torch.Tensor
+        The complex-style tensor storing the probe guess
+    """
+
+    # The basis in the dataset should describe the basis of the probe's
+    # Fourier transform, even if the probe is simulated on larger stage than
+    # the detector. The only issue is if the probe is oversampled in
+    # Fourier space (simulated on a larger stage in real space). That factor
+    # is defined by oversampling.
+
+    probe_basis = (t.as_tensor(dataset.detector_geometry['basis'],
+                               dtype=t.float32)
+                   / oversampling)
+    
+    mean_im = t.mean(dataset.patterns,dim=0)
+    center = image_processing.centroid(mean_im)
+
+    Is = t.arange(shape[0], dtype=t.float32) - center[0]
+    Js = t.arange(shape[1], dtype=t.float32) - center[1]
+    Is,Js = t.meshgrid(Is,Js)
+
+    Rs = t.tensordot(probe_basis,t.stack([Is,Js]),dims=1)
+    forward = t.Tensor([0,0,1])
+    Rs = (forward * dataset.detector_geometry['distance'])[:,None,None] + Rs
+    dirs = Rs / t.linalg.norm(Rs,dim=(0,))
+    angles = t.acos(t.tensordot(forward,dirs,dims=1)) * 1000 #in mrad
+
+    probe_fft = t.zeros(shape,dtype=t.complex64)
+    probe_fft[angles<convergence_semiangle] = 1
+
+    probe_mean = t.mean(t.abs(probe_fft)**2)
+    diff_mean = t.mean(mean_im)
+    probe_fft = probe_fft * t.sqrt(diff_mean / probe_mean)
+    
+    probe_guess = inverse_far_field(probe_fft)
+
+    if propagation_distance is not None:
+        # It's probably worth checking if this is correct when oversampling
+        # is not 1
+
+        probe_shape = t.as_tensor(tuple(probe_guess.shape))
+
+        # Start by recalculating the probe basis from the given information
+        det_basis = t.as_tensor(dataset.detector_geometry['basis'])
+        basis_dirs = det_basis / t.norm(det_basis, dim=0)
+        distance = dataset.detector_geometry['distance']
+        probe_basis = basis_dirs * dataset.wavelength * distance / \
+            (probe_shape * t.norm(det_basis,dim=0))
+
+        # Then package everything as it's needed
+        probe_spacing = t.norm(probe_basis,dim=0).numpy()
+        probe_shape = probe_shape.numpy().astype(np.int32)
+
+        # And generate the propagator
+        AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
+
+        probe_guess = near_field(probe_guess,AS_prop)
+    
+    # Finally, place this probe in a full-sized array if there is oversampling
+    final_probe = t.zeros(shape,dtype=t.complex64)
+    left = shape[0]//2 - probe_guess.shape[0] // 2
+    top = shape[1]//2 - probe_guess.shape[1] // 2 
+    final_probe[left:left+probe_guess.shape[0],
+                top:top+probe_guess.shape[1]] = probe_guess
+    
+    return final_probe
+
+
+def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None):
+
+    # First, check if the probe is a single mode or many modes.
+    # If the probe is many modes, orthogonalize it and use the top mode
+    # for initialization
+    if probe.dim() == 4:
+        probe = orthogonalize_probes(probe)[0]
+    
+    pad0l = (probe.shape[-2] - obj_shape[0])//2
+    pad0r = probe.shape[-2] - obj_shape[0] - pad0l
+    pad1l = (probe.shape[-1] - obj_shape[1])//2
+    pad1r = probe.shape[-1] - obj_shape[1] - pad1l
+    
+    def a_dagger(im):
+        im = t.tensor(im.reshape(obj_shape)).to(dtype=t.complex64)
+        im = inverse_far_field(pad(far_field(im), (pad1l,pad1r,pad0l,pad0r)))
+        exit_wave = probe * im
+        return far_field(exit_wave).numpy().ravel()
+
+    def a(measured):
+        measured = t.tensor(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.complex64)
+        im = inverse_far_field(measured)
+        multiplied = t.conj(probe) * im
+        backplane = far_field(multiplied)
+        clipped = backplane[pad0l:pad0l+obj_shape[0],
+                            pad1l:pad1l+obj_shape[1]]
+        return (inverse_far_field(clipped)).numpy().ravel()
+
+    patsize = pattern.shape[0]*pattern.shape[1]
+    imsize = obj_shape[0]*obj_shape[1]
+    probesize = probe.shape[0]*probe.shape[1]
+    A_dagger = spla.LinearOperator((patsize, imsize),matvec=a_dagger)
+    A = spla.LinearOperator((imsize,patsize),matvec=a)
+
+    # Correct the pattern for the background and mask
+    np_pattern = pattern.numpy()
+    if background is not None:
+        np_pattern = np_pattern - background.numpy()
+    if mask is not None:
+        np_pattern = np_pattern * mask.numpy()
+    np_pattern = np.abs(np_pattern.ravel())
+        
+    realspace_intensities = A * A_dagger * np.ones(imsize)
+    vec = A_dagger * realspace_intensities
+    def y(measured):
+        #return measured * np_pattern
+        # This normalizes the intensities to account for the fact
+        # that some pixels draw from way more spots on the detector than
+        # others
+        return measured * np_pattern / np.abs(vec)
+
+    Y = spla.LinearOperator((patsize, patsize),matvec=y)
+    eigval, z0 = spla.eigs(A * Y * A_dagger, k=n_modes, which='LM')
+    z0 = z0.transpose().reshape(n_modes,obj_shape[0], obj_shape[1])
+
+    # Now we set the overall scale and relative weights of the guess
+    scale_factor = np.sqrt(np.sum(np_pattern) /
+                           t.sum(t.abs(probe)**2).numpy())
+    relative_weights = eigval / np.sum(eigval**2)
+
+    z0 = z0 * (scale_factor * relative_weights[:,None,None])
+    # Now we have to normalize the modes by their eigenvalues
+
+    return t.as_tensor(z0, dtype=t.complex64)
+
+
+def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
+    """Generates guesses of subdominant modes based on spatial derivatives
+    
+    The idea here is that vibration is an extremely common cause of
+    spatial incoherence. This typically presents as subdominant modes that
+    look like the derivative of the dominant mode along various axes. 
+    
+    Therefore, it is a good starting guess if the guess of the dominant
+    mode is reasonable. The output probes will all be normalized to have
+    the same intensity as the input dominant probe
+
+    Parameters
+    ----------
+    dominant_mode : array
+        The complex-valued dominant mode to work from
+    n_modes : int
+        The number of additional modes to create
+    circular : bool
+        Default True, whether to use circular modes (x+iy) or linear (x,y)
+    Returns
+    -------
+    torch.Tensor
+        The complex-style tensor storing the probe guesses
+    """
+
+    # This generates a list of tuples, where each tuple (a,b) corresponds
+    # a term (kx^a)*(ky^b), or (kx+iky)^a*(kx-iky)^b for circular modes
+    def make_orders(n_orders):
+        if n_orders >1024:
+            raise KeyError('Are you sure you want this many orders?')
+        i = 1
+        n = 0
+        total = 0
+        while True:
+            if total >= n_orders:
+                break
+            
+            yield (n,i-n)
+            total += 1
+            if n<i:
+                n+=1
+            else:
+                n = 0
+                i += 1
+
+    # Gotta check and convert to pytorch if it's numpy
+
+    # First step is to take the FFT of the probe
+    dominant_fft = far_field(dominant_mode)
+
+    shape = dominant_mode.shape
+    center = ((shape[-2])//2, (shape[-1])//2)
+        
+    i, j = np.mgrid[:shape[-2], :shape[-1]]
+    i = t.tensor(i - center[0]).to(dtype=dominant_fft.dtype,
+                                    device=dominant_fft.device)
+    j = t.tensor(j - center[1]).to(dtype=dominant_fft.dtype,
+                                    device=dominant_fft.device)
+
+    if circular:
+        a = i + 1j * j
+        b = i - 1j * j
+    else:
+        a = i
+        b = j
+
+    probe_norm = t.sum(t.abs(dominant_mode)**2)
+    # Then we need to multiply that FFT by various phase masks and IFFT
+    probes = []
+    for a_order,b_order in make_orders(n_modes):
+        mask = reduce(t.mul,[a]*a_order+[b]*b_order)
+        new_probe = inverse_far_field(mask * dominant_fft)
+        probes.append(new_probe * (probe_norm / t.sum(t.abs(new_probe))))
+
+    return t.stack(probes)

src/CDTools/tools/interactions/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.interactions.interactions import *

src/CDTools/tools/interactions/interactions.py

@@ -0,0 +1,633 @@
+""" This module contains various to simulate stages in the probe-sample interaction
+
+All the tools here are designed to work with automatic differentiation. Each
+function simulates some aspect of an interaction model that can be used
+for ptychographic reconstruction.
+"""
+
+import torch as t
+import numpy as np
+from CDTools.tools import propagators, image_processing, polarization
+
+__all__ = ['translations_to_pixel', 'pixel_to_translations',
+           'project_translations_to_sample',
+           'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc',
+           'RPI_interaction']
+
+
+def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
+    """Takes real space translations and outputs them in pixel space
+
+    This works for any 2D ptychography geometry. It takes in
+    A set of translations in (x,y) space and outputs the same translations
+    in internal pixel units perpendicular to the detector.
+
+    It uses information on the wavefield basis and, if defined, the
+    sample normal, to perform the conversion.
+
+    The assumed geometry is incoming radiation with a wavevector parallel
+    to the +z axis, [0,0,1]. The default sample orientation has a surface
+    normal parallel to this direction
+
+    Parameters
+    ----------
+    basis : torch.Tensor
+        The real space basis the wavefields are defined in
+    translations : torch.Tensor
+        A Jx3 stack of real-space translations, or a single translation
+    surface_normal : torch.Tensor
+        Optional, the sample's surface normal
+
+    Returns
+    -------
+    pixel_translations : torch.Tensor
+        A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
+    """
+
+    projection_1 = t.Tensor([[1,0,0],
+                             [0,1,0],
+                             [0,0,0]]).to(device=translations.device,dtype=translations.dtype)
+    projection_2 = t.inverse(t.Tensor([[1,0,0],
+                                       [0,1,0],
+                                       -surface_normal/
+                                       surface_normal[2]])).to(device=translations.device,dtype=translations.dtype)
+    basis_vectors_inv = t.pinverse(basis).to(device=translations.device,
+                                             dtype=translations.dtype)
+    projection = t.mm(basis_vectors_inv,
+                      t.mm(projection_2,projection_1))
+    projection = projection.t()
+
+    single_translation = False
+    if len(translations.shape) == 1:
+        translations = translations[None,:]
+        single_translation = True
+
+    pixel_translations = t.mm(translations, projection)
+
+    if single_translation:
+        return pixel_translations[0]
+    else:
+        return pixel_translations
+
+
+def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,0,1])):
+    """Takes pixel-space translations and outputs them in real space
+
+    This works for any 2D ptychography geometry. It takes in
+    A set of internal pixel unit translations in (i,j) space and
+    outputs the same translations real (x,y) space
+
+    It uses information on the wavefield basis and, if defined, the
+    sample normal, to perform the conversion.
+
+    The assumed geometry is incoming radiation with a wavevector parallel
+    to the +z axis, [0,0,1]. The default sample orientation has a surface
+    normal parallel to this direction. Because of this, the z direction
+    translation is always set to zero in the conversion
+
+    Parameters
+    ----------
+    basis : torch.Tensor
+        The real space basis the wavefields are defined in
+    translations : torch.Tensor
+        A Jx2 stack of pixel-space translations, or a single translation
+    surface_normal : torch.Tensor
+        Optional, the sample's surface normal
+
+    Returns
+    -------
+    real_translations : torch.Tensor
+        A Jx3 stack of real-space translations, or a single translation
+    """
+    projection_1 = t.Tensor([[1,0,0],
+                             [0,1,0],
+                             [0,0,0]]).to(device=basis.device,dtype=basis.dtype)
+    projection_2 = t.inverse(t.Tensor([[1,0,0],
+                                       [0,1,0],
+                                       -surface_normal/
+                                       surface_normal[2]])).to(device=basis.device,dtype=basis.dtype)
+    basis_vectors_inv = t.pinverse(basis)
+    projection = t.mm(basis_vectors_inv,
+                      t.mm(projection_2,projection_1))
+    # Literally just need the pseudoinverse of the projection we used to go
+    # the other way
+    projection = t.pinverse(projection).t()
+
+    single_translation = False
+    if len(pixel_translations.shape) == 1:
+        pixel_translations = pixel_translations[None,:]
+        single_translation = True
+
+    translations = t.mm(pixel_translations, projection)
+
+    if single_translation:
+        return translations[0]
+    else:
+        return translations
+
+
+
+def project_translations_to_sample(sample_basis, translations):
+    """Takes real space translations and outputs them in pixels in a sample basis
+
+    This projection function is designed for the Bragg2DPtycho class. More
+    broadly, it works to take a set of translations in the lab frame and
+    convert each one into two values. First, an (i,j) value in pixels
+    describing the location of the probe's intersection with the sample
+    plane, assuming the basis found in sample_basis is used. Second, an
+    amount that the probe must be propagated along the z-axis to reach
+    the sample plane at the given location. This includes both the effect of
+    the tilted sample plane and any explicitly defined motion along the z axis
+    of the probe-forming optic as included in the input translations.
+
+
+    Note that because of the sign convention (that this function returns th
+    relative amount the probe needs to be propagated to reach any given
+    location), a positive motion along the z-axis of the probe forming optics
+    will lead to a negative propagation distance.
+
+    The assumed geometry is incoming radiation with a wavevector parallel
+    to the +z axis, [0,0,1].
+
+    Parameters
+    ----------
+    sample_basis : torch.Tensor
+        The real space basis the wavefields are defined in
+    translations : torch.Tensor
+        A Jx3 stack of real-space translations, or a single translation
+
+    Returns
+    -------
+    pixel_translations : torch.Tensor
+        A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
+    propagations : torch.Tensor
+        A length-J vector of propagation distances, in meters
+    """
+
+    # Must do this all in pytorch unfortunately
+    # First we calculate the surface normal for the projection onto the sample
+    surface_normal =  t.cross(sample_basis[:,1],sample_basis[:,0])
+    surface_normal /= t.norm(surface_normal)
+
+
+    # Then we calculate a matrix which can do the projection
+
+    propagation_dir = t.Tensor(np.array([0,0,1])).to(
+        device=surface_normal.device,
+        dtype=surface_normal.dtype)
+
+    I = t.eye(3).to(
+        device=surface_normal.device,
+        dtype=surface_normal.dtype)
+
+    # Here we're setting up a matrix-vector equation mat*answer=input
+    # At some point ger will need to be replaced by outer, but for now
+    # outer many places still don't have new enough versions of torch.
+    mat = t.cat((I - t.ger(propagation_dir,propagation_dir),
+                 surface_normal.unsqueeze(0)))
+
+    # And we invert the matrix to do the projection
+    projector = t.pinverse(mat)[:,:3].to(device=translations.device,
+                                         dtype=translations.dtype)
+
+    # Finally, we need to get the result in the right basis, so this will
+    # do that conversion for us
+    basis_vectors_inv = t.pinverse(sample_basis).to(device=translations.device,
+                                                    dtype=translations.dtype)
+
+    # Same for the propagation direction to extract how far it must propagate
+    propagation_dir_inv = propagation_dir.unsqueeze(0).to(
+        device=translations.device,
+        dtype=translations.dtype)
+
+
+    sample_projection = t.mm(basis_vectors_inv, projector).t()
+    prop_projection = t.mm(propagation_dir_inv, projector).t()
+
+
+    single_translation = False
+    if len(translations.shape) == 1:
+        translations = translations[None,:]
+        single_translation = True
+
+    pixel_translations = t.mm(translations, sample_projection)
+    propagations = t.mm(translations, prop_projection) \
+        - t.mm(translations,propagation_dir[:,None])
+
+    if single_translation:
+        return pixel_translations[0], propagations[0]
+    else:
+        return pixel_translations, propagations
+
+
+
+
+def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False):
+    """Returns a stack of exit waves without 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 by the rounded value of the translation
+
+    If multiple_modes is set to False, any additional dimensions in the
+    ptycho_2D_round function will be assumed to correspond to the translation
+    index. If multiple_modes is set to true, the (-4th) dimension of the probe
+    will always be assumed to be defining a set of (P) incoherently mixing
+    modes to be broadcast all translation indices. If any additional dimensions
+    closer to the start exist, they will be assumed to be translation indices
+
+
+    Parameters
+    ----------
+    probe : torch.Tensor
+        A (P)xMxL probe function to illuminate the object
+    object : torch.Tensor
+        The object function to be probed
+    translations : torch.Tensor
+        The (N)x2 array of (i,j) translations to simulate
+    multuple_modes : bool
+        Default False, whether to assume the probe contains multiple modes
+
+    Returns
+    -------
+    exit_waves : torch.Tensor
+        An (N)x(P)xMxL tensor of the calculated exit waves
+    """
+
+    single_translation = False
+    if translations.dim() == 1:
+        translations = translations[None,:]
+        single_translation = True
+
+
+    integer_translations = t.round(translations).to(dtype=t.int32)
+
+    if upsample_obj:
+        selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2,
+                                  tr[1]:tr[1]+probe.shape[-1]//2]
+                              for tr in integer_translations])
+        selections = image_processing.fourier_upsample(selections,
+                                                       preserve_mean=True)
+
+    else:
+        selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
+                              tr[1]:tr[1]+probe.shape[-1]]
+                          for tr in integer_translations])
+
+
+    if multiple_modes:
+        # if the probe dimension is 4, then this hasn't yet been broadcast
+        # over the translation dimensions
+        output = probe * selections[:,None,:,:]
+    else:
+        output = probe * selections
+
+    if single_translation:
+        return output[0]
+    else:
+        return output
+
+
+
+def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
+    """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 linear
+    interpolation to combine the results
+
+    If shift_probe is True, it applies the subpixel shift to the probe,
+    otherwise the subpixel shift is applied to the object
+
+    Parameters
+    ----------
+    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
+        Default True, Whether to subpixel shift the probe or object
+
+    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 = []
+    if shift_probe:
+        for tr, sp in zip(integer_translations,
+                          subpixel_translations):
+            # This isn't perfectly symmetric but I think it's okay for now
+            # It should get the job done
+            # Basically, we shift the probe's position by a subpixel (i,j),
+            # rolling the edges of the array, and use that to multiply
+            # by the object
+            sel00 = probe[:,:]
+            sel01 = t.cat((probe[:,-1:],probe[:,:-1]),dim=1)
+            sel10 = t.cat((probe[-1:,:],probe[:-1,:]),dim=0)
+            sel11 = t.cat((sel01[-1:,:],sel01[:-1,:]),dim=0)
+
+            selection = sel00 * (1-sp[0])*(1-sp[1]) + \
+                sel10 * sp[0]*(1-sp[1]) + \
+                sel01 * (1-sp[0])*sp[1] + \
+                sel11 * sp[0]*sp[1]
+
+            obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
+                            tr[1]:tr[1]+probe.shape[1]]
+
+            exit_waves.append(selection * obj_slice)
+    else:
+        for tr, sp in zip(integer_translations,
+                          subpixel_translations):
+            #
+            # Here we subpixel shift the object by (-i,-j) after
+            # slicing out the correct translation of the probe
+            #
+
+            sel00 = obj[tr[0]:tr[0]+probe.shape[0],
+                        tr[1]:tr[1]+probe.shape[1]]
+
+            sel01 = obj[tr[0]:tr[0]+probe.shape[0],
+                        tr[1]+1:tr[1]+1+probe.shape[1]]
+
+            sel10 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
+                        tr[1]:tr[1]+probe.shape[1]]
+
+            sel11 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
+                        tr[1]+1:tr[1]+1+probe.shape[1]]
+
+            selection = sel00 * (1-sp[0])*(1-sp[1]) + \
+                sel01 * (1-sp[0])*sp[1] + \
+                sel10 * sp[0]*(1-sp[1]) + \
+                sel11 * sp[0]*sp[1]
+
+            exit_waves.append(probe * selection)
+
+    if single_translation:
+        return exit_waves[0]
+    else:
+        return t.stack(exit_waves)
+
+
+def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None):
+    """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 [not yet implemented]
+
+    If multiple_modes is set to False, any additional dimensions in the
+    ptycho_2D_round function will be assumed to correspond to the translation
+    index. If multiple_modes is set to true, the (-4th) dimension of the probe
+    will always be assumed to be defining a set of (P) incoherently mixing
+    modes to be broadcast all translation indices. If any additional dimensions
+    closer to the start exist, they will be assumed to be translation indices
+
+    Parameters
+    ----------
+    probe : torch.Tensor
+        An (P)xMxL probe function for the exit waves
+    object : torch.Tensor
+        The object function to be probed
+    translations : torch.Tensor
+        The (N)x2 array of translations to simulate
+    shift_probe : bool
+        Default True, Whether to subpixel shift the probe or object
+    multuple_modes : bool
+        Default False, whether to assume the probe contains multiple modes
+
+    Returns
+    -------
+    exit_waves : torch.Tensor
+        An (N)x(P)xMxL 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)
+    if not polarized:
+        selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
+                                  tr[1]:tr[1]+probe.shape[-1]]
+                              for tr in integer_translations])
+    else:
+        selections = t.stack([obj[:, :, tr[0]:tr[0]+probe.shape[-2],
+                                  tr[1]:tr[1]+probe.shape[-1]]
+                      for tr in integer_translations])
+        # Nx2x2xMxL tensor
+
+    exit_waves = []
+    if shift_probe:
+        i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
+            - probe.shape[-2]//2
+        j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
+            - probe.shape[-1]//2
+        I,J = t.meshgrid(i,j, indexing='ij')
+        I = 2 * np.pi * I / probe.shape[-2]
+        J = 2 * np.pi * J / probe.shape[-1]
+        phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
+                                -subpixel_translations[:,1,None,None]*J))
+        if polarized:
+            phase_masks = phase_masks[..., None, :, :]
+            # Nx2x1xMxL tensor
+            # probe is (N)(P)x2xMxL tensor
+        fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
+        if multiple_modes: # Multi-mode probe
+            if polarized:
+                shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:]
+            else:
+                shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
+        else:
+            shifted_fft_probe = fft_probe * phase_masks
+        shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
+                                                    dim=(-1,-2)))
+        if not polarized:
+            if multiple_modes: # Multi-mode probe
+                output = shifted_probe * selections[...,None,:,:]
+            else:
+                output = shifted_probe * selections
+        # selections: Nx2x2xMxL
+        # probe: Nx(P)x2x1xMxL
+        else:
+            output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes)
+
+    else:
+        raise NotImplementedError('Object shift not yet implemented')
+    if single_translation:
+        return output[0]
+    else:
+        return output
+
+
+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. NOTE: I believe
+    this change has now been made
+
+    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[-2]) - probe.shape[-2]//2
+        j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
+        I,J = t.meshgrid(i,j, indexing='ij')
+        I = 2 * np.pi * I.to(t.float32) / probe.shape[-2]
+        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 = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2))
+            shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J))
+            shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
+                                                         dim=(-1,-2)))
+
+            s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B,
+                                      tr[1]:tr[1]+probe.shape[-1]+2*B]
+
+
+            output = t.zeros([probe.shape[-2]+2*B,probe.shape[-1]+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[-2],j:j+probe.shape[-1]] += \
+                        shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]]
+
+            exit_waves.append(output)
+
+    else:
+        raise NotImplementedError('Object shift not yet implemented')
+
+    if single_translation:
+        return exit_waves[0]
+    else:
+        return t.stack(exit_waves)
+
+
+def RPI_interaction(probe, obj):
+    """Returns an exit wave from a high-res probe and a low-res obj
+
+    In this interaction, the probe and object arrays are assumed to cover
+    the same physical region of space, but with the probe array sampling that
+    region of space more finely. Thus, to do the interaction, the object
+    is first upsampled by padding it in Fourier space (equivalent to a sinc
+    interpolation) before being multiplied with the probe. This function is
+    called RPI_interaction because this interaction is central to the RPI
+    method and is not commonly used elsewhere.
+
+    This also works with object functions that have an extra first dimension
+    for an incoherently mixing model.
+
+
+    Parameters
+    ----------
+    probe : torch.Tensor
+        An MxL probe function for simulating the exit waves
+    obj : torch.Tensor
+        An M'xL' or NxM'xL' object function for simulating the exit waves
+
+    Returns
+    -------
+    exit_wave : torch.Tensor
+        An MxL tensor of the calculated exit waves
+    """
+
+    # TODO: The upsampling only works for arrays of even dimension!
+
+    # The far-field propagator is just a 2D FFT but with an fftshift
+    fftobj = propagators.far_field(obj)
+    # We calculate the padding that we need to do the upsampling
+    # This is carefully set up to keep the zero-frequency pixel in the correct
+    # location as the overall shape changes. Don't mess with this without
+    # having thought about this carefully.
+    pad2l = probe.shape[-2]//2 - obj.shape[-2]//2
+    pad2r = probe.shape[-2] - obj.shape[-2] - pad2l
+    pad1l = probe.shape[-1]//2 - obj.shape[-1]//2
+    pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
+        
+    fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r))
+
+    # Again, just an inverse FFT but with an fftshift
+    upsampled_obj = propagators.inverse_far_field(fftobj)
+
+    if obj.dim() >= 3:
+        return probe[None,...] *  upsampled_obj
+    else:
+        return probe * upsampled_obj

src/CDTools/tools/losses/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.losses.losses import *

src/CDTools/tools/losses/losses.py

@@ -0,0 +1,174 @@
+"""Contains various loss functions to be used for optimization
+
+It exposes three losses, one returning the mean squared amplitude error, one
+that returns the mean squared intensity error, and one that returns the
+maximum likelihood metric for a system with Poisson statistics.
+
+"""
+
+import torch as t
+
+__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll']
+
+
+def amplitude_mse(intensities, sim_intensities, mask=None):
+    """ Returns the mean squared error of a simulated dataset's amplitudes
+
+    Calculates the mean squared error between a given set of 
+    measured diffraction intensities and a simulated set.
+
+    This function calculates the mean squared error between their
+    associated amplitudes. Because this is not well defined for negative
+    numbers, make sure that all the intensities are >0 before using this
+    loss. Note that this is actually a sum-squared error, because this
+    formulation makes it vastly simpler to compare error calculations
+    between reconstructions with different minibatch size. I hope to
+    find a better way to do this that is more honest with this
+    cost function, though.
+
+    It can accept intensity and simulated intensity tensors of any shape
+    as long as their shapes match, and the provided mask array can be
+    broadcast correctly along them.
+
+    This is empirically the most useful loss function for most cases
+
+    Parameters
+    ----------
+    intensities : torch.Tensor
+        A tensor with measured detector values
+    sim_intensities : torch.Tensor
+        A tensor of simulated detector intensities
+    mask : torch.Tensor
+        A mask with ones for pixels to include and zeros for pixels to exclude
+
+    Returns
+    -------
+    loss : torch.Tensor
+        A single value for the mean amplitude mse
+    """
+
+    # I know it would be more efficient if this function took in the
+    # amplitudes instead of the intensities, but I want to be consistent
+    # with all the errors working off of the same inputs
+
+    if mask is None:
+        return t.sum((t.sqrt(sim_intensities) -
+                      t.sqrt(intensities))**2)
+    else:
+        masked_intensities = intensities.masked_select(mask)
+        return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
+                      t.sqrt(masked_intensities))**2)
+
+
+    
+def intensity_mse(intensities, sim_intensities, mask=None):
+    """ Returns the mean squared error of a simulated dataset's intensities
+
+    Calculates the summed mean squared error between a given set of 
+    diffraction intensities - the measured set of detector intensities -
+    and a simulated set of diffraction intensities. This function
+    calculates the mean squared error between the intensities.
+
+    It can accept intensity and simulated intensity tensors of any shape
+    as long as their shapes match, and the provided mask array can be
+    broadcast correctly along them.
+
+    Parameters
+    ----------
+    intensities : torch.Tensor
+        A tensor with measured detector intensities.
+    sim_intensities : torch.Tensor
+        A tensor of simulated detector intensities
+    mask : torch.Tensor
+        A mask with ones for pixels to include and zeros for pixels to exclude
+
+    Returns
+    -------
+    loss : torch.Tensor
+        A single value for the mean intensity mse
+
+    """
+    if mask is None:
+        return t.sum((sim_intensities - intensities)**2) \
+            / intensities.view(-1).shape[0]
+    else:
+        masked_intensities = intensities.masked_select(mask)
+        return t.sum((sim_intensities.masked_select(mask) -
+                      masked_intensities)**2) \
+                      / masked_intensities.shape[0]
+
+
+    
+def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6, subtract_min=False):
+    """ Returns the Poisson negative log likelihood for a simulated dataset's intensities
+
+    Calculates the overall Poisson maximum likelihood metric using
+    diffraction intensities - the measured set of detector intensities -
+    and a simulated set of intensities. This loss would be appropriate
+    for detectors in a single-photon counting mode, with their output
+    scaled to number of photons
+
+    Note that this calculation ignores the log(intensities!) term in the
+    full expression for Poisson negative log likelihood. This term doesn't
+    change the calculated gradients so isn't worth taking the time to compute
+
+    It can accept intensity and simulated intensity tensors of any shape
+    as long as their shapes match, and the provided mask array can be
+    broadcast correctly along them.
+
+    The default value of eps is 1e-6 - a nonzero value here helps avoid
+    divergence of the log function near zero.
+
+    Parameters
+    ----------
+    intensities : torch.Tensor
+        A tensor with measured detector intensities.
+    sim_intensities : torch.Tensor
+        A tensor of simulated detector intensities
+    mask : torch.Tensor
+        A mask with ones for pixels to include and zeros for pixels to exclude
+    eps : float
+        Optional, a small number to add to the simulated intensities
+    
+    Returns
+    -------
+    loss : torch.Tensor
+        A single value for the poisson negative log likelihood
+
+    """
+    #When x.logy gets into the regular build, add it by uncommenting!
+    if mask is None:
+
+        nll = t.sum(sim_intensities+eps -
+                    intensities * t.log(sim_intensities+eps)) \
+                    / intensities.view(-1).shape[0]
+        #nll = t.sum(sim_intensities+epsa -
+        #            t.xlogy(intensities,sim_intensities+eps)) \
+        #            / intensities.view(-1).shape[0]
+        if subtract_min:
+            nll -= t.nansum(intensities - intensities*t.log(intensities))\
+                / intensities.view(-1).shape[0]
+            #nll -= t.sum(intensities - t.xlogy(intensities,intensities))
+            # We don't need to include the log factorial part here, because
+            # it will get subtracted off in the min anyway.
+            
+        return nll
+    else:
+        masked_intensities = intensities.masked_select(mask)
+        masked_sims = sim_intensities.masked_select(mask)
+
+        nll = t.sum(masked_sims + eps - \
+                    masked_intensities * t.log(masked_sims+eps)) \
+                    / masked_intensities.shape[0]
+
+        #nll = t.sum(masked_sims + eps - \
+        #            t.xlogy(masked_intensities, masked_sims+eps)) \
+        #            / masked_intensities.shape[0]
+        
+        
+        if subtract_min:
+            nll -= t.nansum(masked_intensities -
+                            masked_intensities*t.log(masked_intensities)) \
+                            / masked_intensities.shape[0]
+
+        return nll

src/CDTools/tools/measurements/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.measurements.measurements import *

src/CDTools/tools/measurements/measurements.py

@@ -0,0 +1,159 @@
+"""This module contains tools to simulate various measurement models
+
+All the measurements here are safe to use in an automatic differentiation
+model. There exist tools to simulate detectors with finite saturation
+thresholds, backgrounds, and more.
+"""
+
+import torch as t
+import numpy as np
+from torch.nn.functional import avg_pool2d
+
+#
+# This file will host tools to turn a propagated wavefield into a measured
+# intensity pattern on a detector
+#
+
+__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
+
+
+def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
+    """Returns the intensity of a wavefield
+    
+    The intensity is defined as the magnitude squared of the
+    wavefront. If a detector slice is given, the returned array
+    will only include that slice from the simulated wavefront.
+    
+    Parameters
+    ----------
+    wavefield : torch.Tensor
+        A JxMxN stack of complex-valued wavefields
+    detector_slice : slice
+        Optional, a slice or tuple of slices defining a section of the simulation to return
+    saturation : float
+        Optional, a maximum saturation value to clamp the resulting intensities to
+    oversampling : int
+        Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
+
+    Returns
+    -------
+    sim_patterns : torch.Tensor
+        A real MxN array storing the wavefield's intensities
+    """
+    output = t.abs(wavefield)**2
+
+    # Now we apply oversampling
+    if oversampling != 1:
+        if wavefield.dim() == 2:
+            output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
+        else:
+            output = avg_pool2d(output, oversampling)
+
+    # Then we grab the detector slice
+    if detector_slice is not None:
+        if wavefield.dim() == 2:
+            output = output[detector_slice]
+        else:
+            output = output[(np.s_[:],) + detector_slice]
+
+    # And now saturation    
+    if saturation is None:
+        return output + epsilon
+    else:
+        return t.clamp(output + epsilon,0,saturation)
+            
+            
+def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
+    """Returns the incoherent sum of the intensities of the wavefields
+    
+    The intensity is defined as the sum of the magnitudes squared of
+    the wavefields. If a detector slice is given, the returned array
+    will only include that slice from the simulated wavefronts.
+
+    The (-3) index is the set of incoherently adding patterns, and any
+    indexes further to the front correspond to the set of diffraction patterns
+    to measure. The (-2) and (-1) indices are the wavefield, and the final
+    index is the complex index
+    
+    Parameters
+    ----------
+    wavefields : torch.Tensor
+        An LxJxMxNx stack of complex wavefields
+    detector_slice : slice
+        Optional, a slice or tuple of slices defining a section of the simulation to return
+    saturation : float
+        Optional, a maximum saturation value to clamp the resulting intensities to
+    oversampling : int
+        Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
+
+    Returns
+    -------
+    sim_patterns : torch.Tensor 
+        A real LXMxN array storing the incoherently summed intensities
+    """
+
+    output = t.sum(t.abs(wavefields)**2,dim=-3) 
+
+    # Now we apply oversampling
+    if oversampling != 1:
+        if wavefields.dim() == 3:
+            output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
+        else:
+            output = avg_pool2d(output, oversampling)
+            
+    # Then we grab the detector slice
+    if detector_slice is not None:
+        if wavefields.dim() == 3:
+            output = output[detector_slice]
+        else:
+            output = output[(np.s_[:],) + detector_slice]
+
+    if saturation is None:
+        return output + epsilon
+    else:
+        return t.clamp(output + epsilon,0,saturation)
+
+
+def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
+    """Returns the intensity of a wavefield plus a background
+    
+    The intensity is calculated via the given measurment function 
+    Then, the square of the given background is added. This kind
+    of background model is commonly used to enforce positivity of the
+    background model.
+
+    Parameters
+    ----------
+    wavefield : torch.Tensor
+        A JxMxN stack of complex-valued wavefields
+    background : torch.Tensor
+        An tensor storing the square root of the detector background
+    detector_slice : slice
+        Optional, a slice or tuple of slices defining a section of the simulation to return
+    measurement : function
+        Default is measurements.intensity, the measurement function to use.
+    saturation : float
+        Optional, a maximum saturation value to clamp the resulting intensities to
+    oversampling : int
+        Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
+
+    Returns
+    -------
+    sim_patterns : torch.Tensor
+        A real MxN array storing the wavefield's intensities
+    """
+    
+    if detector_slice is None:
+        output = measurement(wavefield, *args, epsilon=epsilon,
+                             oversampling=oversampling) + background**2
+    else:
+        output = measurement(wavefield, *args, detector_slice=detector_slice,
+                             epsilon=epsilon, oversampling=oversampling) \
+                             + background**2
+
+    # This has to be done after the background is added, hence we replicate
+    # it here
+    if saturation is None:
+        return output
+    else:
+        return t.clamp(output,0,saturation)

src/CDTools/tools/plotting/__init__.py

@@ -0,0 +1,2 @@
+from CDTools.tools.plotting.plotting import *
+

src/CDTools/tools/plotting/plotting.py

@@ -0,0 +1,782 @@
+"""This module contains functions for plotting various important metrics
+
+All the plotting functions here can accept torch input or numpy input,
+to facilitate their use both for live inspection of running reconstructions
+and for after-the-fact analysis. Utilities for plotting complex valued
+images exist, as well as plotting scan patterns and nanomaps
+"""
+
+import torch as t
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import hsv_to_rgb
+from matplotlib.widgets import Slider
+from matplotlib import ticker, patheffects
+
+
+__all__ = ['colorize', 'plot_amplitude', 'plot_phase',
+           'plot_colorized', 'plot_translations', 'get_units_factor',
+           'plot_nanomap', 'plot_real', 'plot_imag',
+           'plot_nanomap_with_images']
+
+
+def colorize(z):
+    """ Returns RGB values for a complex color plot given a complex array
+    This function returns a set of RGB values that can be used directly
+    in a call to imshow based on an input complex numpy array (not a
+    torch tensor representing a complex field)
+
+    Parameters
+    ----------
+    z : array
+        A complex-valued array
+    Returns
+    -------
+    rgb : list(array)
+        A list of arrays for the R,G, and B channels of an image
+    """
+
+    amp = np.abs(z)
+    rmin = 0
+    rmax = np.max(amp)
+    amp = np.where(amp < rmin, rmin, amp)
+    amp = np.where(amp > rmax, rmax, amp)
+    ph = np.angle(z, deg=1) + 90
+    # HSV are values in range [0,1]
+    h = (ph % 360) / 360
+    s = 0.85 * np.ones_like(h)
+    v = (amp - rmin) / (rmax - rmin)
+
+    return hsv_to_rgb(np.dstack((h,s,v)))
+
+
+def get_units_factor(units):
+    """Gets the multiplicative factor associated with a length unit
+
+    Parameters
+    ----------
+    units : str
+        The abbreviation for the unit type
+
+    Returns
+    -------
+    factor : float
+        The factor meters / (unit)
+    """
+
+    u = units.lower()
+    if u=='m':
+        factor=1
+    if u=='cm':
+        factor=1e2
+    if u=='mm':
+        factor=1e3
+    if u=='um' or u=="$\\mu$m":
+        factor=1e6
+    if u=='nm':
+        factor=1e9
+    if u=='a':
+        factor=1e10
+    if u=='pm':
+        factor=1e12
+    return factor
+
+
+def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs):
+    """Plots an image with a colorbar and on an appropriate spatial grid
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Finally, if a function is passed to the plot_func argument, this function
+    will be called on each slice of data before it is plotted. This is used
+    internally to enable the plot_real, plot_image, plot_phase, etc. functions.
+
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    plot_func : callable
+        A function which maps numpy arrays to the image to be plotted
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    cmap_label : str
+        What to label the colorbar when plotting
+    interpolation : str
+        What interpolation to use for imshow
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+
+    # convert to numpy
+    if isinstance(im, t.Tensor):
+        # If final dimension is 2, assume it is a complex array. If not,
+        # assume it represents a real array
+        if im.shape[-1] == 2:
+            im = im.detach().cpu().numpy()
+        else:
+            im = im.detach().cpu().numpy()
+
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+
+    # This nukes everything and updates either the appropriate image from the
+    # stack of images, or the only image if only a single image has been
+    # given
+    def make_plot(idx):
+        plt.figure(fig.number)
+        title = plt.gca().get_title()
+        fig.clear()
+
+
+        # If im only has two dimensions, this reshape will add a leading
+        # dimension, and update will be called on index 0. If it has 3 or more
+        # dimensions, then all the leading dimensions will be compressed into
+        # one long dimension which can be scrolled through.
+        s = im.shape
+        reshaped_im = im.reshape(-1,s[-2],s[-1])
+        num_images = reshaped_im.shape[0]
+        fig.plot_idx = idx % num_images
+
+        to_plot = plot_func(reshaped_im[fig.plot_idx])
+
+        #Plot in a basis if it exists, otherwise dont
+        if basis is not None:
+            if isinstance(basis,t.Tensor):
+                np_basis = basis.detach().cpu().numpy()
+            else:
+                np_basis = basis
+            # This fails if the basis is not rectangular
+            basis_norm = np.linalg.norm(np_basis, axis = 0)
+            basis_norm = basis_norm * get_units_factor(units)
+
+            extent = [0, im.shape[-1]*basis_norm[1], 0,
+                      im.shape[-2]*basis_norm[0]]
+        else:
+            extent=None
+
+        plt.imshow(to_plot, cmap = cmap, extent = extent, interpolation=interpolation)
+        cbar = plt.colorbar()
+        if cmap_label is not None:
+            cbar.set_label(cmap_label)
+
+        if basis is not None:
+            plt.xlabel('X (' + units + ')')
+            plt.ylabel('Y (' + units + ')')
+        else:
+            plt.xlabel('j (pixels)')
+            plt.ylabel('i (pixels)')
+
+
+        plt.title(title)
+
+        if len(im.shape) >= 3:
+            plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure)
+        return fig
+
+    if hasattr(fig, 'plot_idx'):
+        result = make_plot(fig.plot_idx)
+    else:
+        result = make_plot(0)
+
+    update = make_plot
+
+
+    def on_action(event):
+        if not hasattr(event, 'button'):
+            event.button = None
+        if not hasattr(event, 'key'):
+            event.key = None
+
+        if event.key == 'up' or event.button == 'up':
+            update(fig.plot_idx - 1)
+        elif event.key == 'down' or event.button == 'down':
+            update(fig.plot_idx + 1)
+        plt.draw()
+
+    if len(im.shape) >=3:
+        if not hasattr(fig,'my_callbacks'):
+            fig.my_callbacks = []
+
+        for cid in fig.my_callbacks:
+            fig.canvas.mpl_disconnect(cid)
+        fig.my_callbacks = []
+        fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
+        fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
+
+    return result
+
+
+def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
+    """Plots the real part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    cmap_label : str
+        What to label the colorbar when plotting
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    plot_func = lambda x: np.real(x)
+    return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
+                      units=units, cmap=cmap, cmap_label=cmap_label,
+                      **kwargs)
+
+
+
+def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
+    """Plots the imaginary part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    cmap_label : str
+        What to label the colorbar when plotting
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    plot_func = lambda x: np.imag(x)
+    return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
+                      units=units, cmap=cmap, cmap_label=cmap_label,
+                      **kwargs)
+
+
+def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs):
+    """Plots the amplitude of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates.
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    cmap_label : str
+        What to label the colorbar when plotting
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    plot_func = lambda x: np.absolute(x)
+    return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
+                      units=units, cmap=cmap, cmap_label=cmap_label,
+                      **kwargs)
+
+
+def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', cmap_label='Phase (rad)', **kwargs):
+    """ Plots the phase of a complex array with dimensions NxMx2
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    cmap_label : str
+        What to label the colorbar when plotting
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    if cmap == 'auto':
+        if 'twilight' in plt.colormaps():
+            cmap = 'twilight'
+        elif 'hsv' in plt.colormaps():
+            cmap = 'hsv'
+        else:
+            raise AttributeError('Neither twilight or hsv colormap exists in this screwed up matplotlib install')
+
+    plot_func = lambda x: np.angle(x)
+    return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
+                      units=units, cmap=cmap, cmap_label=cmap_label,
+                      **kwargs)
+
+
+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
+
+    The darkness corresponds to the intensity of the image, and the color
+    corresponds to the phase.
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    plot_func = lambda x: colorize(x)
+    return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
+                      units=units, **kwargs)
+
+
+def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
+    """Plots a set of probe translations in a nicely formatted way
+
+    Parameters
+    ----------
+    translations : array
+        An Nx2 or Nx3 set of translations in real space
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    units : str
+        Default is um, units to report in (assuming input in m)
+    lines : bool
+        Whether to plot lines indicating the path taken
+    invert_xaxis : bool
+        Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+
+    factor = get_units_factor(units)
+
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    if isinstance(translations, t.Tensor):
+        translations = translations.detach().cpu().numpy()
+
+    translations = translations * factor
+    plt.plot(translations[:,0], translations[:,1],'k.')
+    if invert_xaxis:
+        plt.gca().invert_xaxis()
+        
+    if lines:
+        plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
+    plt.xlabel('X (' + units + ')')
+    plt.ylabel('Y (' + units + ')')
+
+    return fig
+
+
+def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True):
+    """Plots a set of nanomap data in a flexible way
+
+    Parameters
+    ----------
+    translations : array
+        An Nx2 or Nx3 set of translations in real space
+    values : array
+        A length-N object of values associated with the translations
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    units : str
+        Default is um, units to report in (assuming input in m)
+    convention : str
+        Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
+    invert_xaxis : bool
+        Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+
+    if fig is None:
+        fig = plt.figure()
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    factor = get_units_factor(units)
+
+    bbox = fig.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+    if isinstance(translations, t.Tensor):
+        trans = translations.detach().cpu().numpy()
+    else:
+        trans = np.array(translations)
+
+    if isinstance(values, t.Tensor):
+        values = values.detach().cpu().numpy()
+    else:
+        values = np.array(values)
+
+    if convention.lower() != 'probe':
+        trans = trans * -1
+
+    s = bbox.width * bbox.height / trans.shape[0] * 72**2 #72 is points per inch
+    s /= 4 # A rough value to make the size work out
+
+    plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
+    if invert_xaxis:
+        plt.gca().invert_xaxis()
+    
+    plt.gca().set_facecolor('k')
+    plt.xlabel('Translation x (' + units + ')')
+    plt.ylabel('Translation y (' + units + ')')
+    plt.colorbar()
+
+    return fig
+
+
+def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
+    """Plots a nanomap, with an image or stack of images for each point
+
+    In many situations, ptychography data or the output of ptychography
+    reconstructions is formatted as a set of images associated with various
+    points in real space. This function is designed to allow for browsing
+    through this kind of data, by making it possible to visualize a
+
+    """
+
+    # This should pull heavily from the dataset.inspect function
+    # In fact, I should be able to replace most of that function with a
+    # call to this function once it's built
+    # We start by making the figure and axes
+
+    # The key will be writing this so it works okay when called in "update"
+    # mode, i.e. on a figure that already has this thing showing.
+
+    if fig is None:
+        fig = plt.figure(figsize=(8,5.3))
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+        if hasattr(fig, 'nanomap_cids'):
+            for cid in fig.nanomap_cids:
+                fig.canvas.mpl_disconnect(cid)
+
+    # Does figsize work with the fig.subplots, or just for plt.subplots?
+    axes = fig.subplots(1,2)
+
+    fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
+    plt.subplots_adjust(wspace=0.25) #avoids overlap of labels with plots
+    axslider = plt.axes([0.15,0.06,0.75,0.03])
+
+    # This gets the set of sizes for the points in the nanomap
+    def calculate_sizes(idx):
+        bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+        s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
+        s0 /= 4 # A rough value to make the size work out
+        s = np.ones(translations.shape[0]) * s0
+
+        s[idx] *= 4
+        return s
+
+
+    def update_colorbar(im):
+        #
+        # This solves the problem of the colorbar being changed
+        # when the forward and back buttons are used!!!
+        #
+        if hasattr(im, 'norecurse') and im.norecurse:
+            im.norecurse=False
+            return
+
+        im.norecurse=True
+        # This is needed to update the colorbar
+        # only change limits if array contains multiple values
+        if np.min(im.get_array()) != np.max(im.get_array()):
+            im.set_clim(vmin=np.min(im.get_array()),
+                        vmax=np.max(im.get_array()))
+
+    #
+    # The meatiest part of this program, here we just go through and
+    # set up the plot how we want it
+    #
+
+    # First we set up the left-hand plot, which shows an overview map
+    axes[0].set_title('Relative Displacement Map')
+
+    translations = translations.detach().cpu().numpy()
+
+    if convention.lower() != 'probe':
+        translations = translations * -1
+
+    s = calculate_sizes(0)
+
+    nanomap_units_factor = get_units_factor(nanomap_units)
+    nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
+                              nanomap_units_factor * translations[:,1],
+                              s=s,c=values, picker=True)
+
+    axes[0].invert_xaxis()
+    axes[0].set_facecolor('k')
+    axes[0].set_xlabel('Translation x ('+nanomap_units+')', labelpad=1)
+    axes[0].set_ylabel('Translation y ('+nanomap_units+')', labelpad=1)
+    cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',
+                       format='%.2e',
+                       ticks=ticker.LinearLocator(numticks=5),
+                       pad=0.17,fraction=0.1)
+    cb1.ax.set_title(nanomap_colorbar_title, size="medium", pad=5)
+    cb1.ax.tick_params(labelrotation=20)
+    if values is None:
+        # This seems to do a good job of leaving the appropriate space
+        # where the colorbar should have been to avoid stretching the
+        # nanomap plot, while still not showing the (now useless) colorbar.
+        cb1.remove()
+
+    # Now we set up the second plot, which shows the individual
+    # diffraction patterns
+    axes[1].set_title(image_title)
+    #Plot in a basis if it exists, otherwise dont
+    if basis is not None:
+        axes[1].set_xlabel('X (' + image_units + ')')
+        axes[1].set_ylabel('Y (' + image_units + ')')
+
+        example_im = get_image_func(0)
+        if isinstance(example_im, t.Tensor):
+            example_im = example_im.cpu().numpy()
+
+        if isinstance(basis,t.Tensor):
+            np_basis = basis.detach().cpu().numpy()
+        else:
+            np_basis = basis
+        # This fails if the basis is not rectangular
+        basis_norm = np.linalg.norm(np_basis, axis = 0)
+        basis_norm = basis_norm * get_units_factor(image_units)
+
+        extent = [0, example_im.shape[-1]*basis_norm[1], 0,
+                  example_im.shape[-2]*basis_norm[0]]
+    else:
+        axes[1].set_xlabel('j (pixels)')
+        axes[1].set_ylabel('i (pixels)')
+        extent=None
+
+    im=get_image_func(0)
+    if len(im.shape) >= 3:
+        im_idx=0
+        axes[1].image_idx = im_idx
+        im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
+        axes[1].text_box = axes[1].text(0.98, 0.98, str(im_idx), color='w',
+                                        fontsize=14,
+                                        horizontalalignment='right',
+                                        verticalalignment='top',
+                                        transform=axes[1].transAxes)
+        axes[1].text_box.set_path_effects(
+            [patheffects.Stroke(linewidth=2, foreground='black'),
+             patheffects.Normal()])
+
+    meas = axes[1].imshow(im, extent=extent, cmap=cmap)
+
+    cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
+                       format='%.2e',
+                       ticks=ticker.LinearLocator(numticks=5),
+                       pad=0.17,fraction=0.1)
+    cb2.ax.tick_params(labelrotation=20)
+    cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
+    cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
+
+    # This function handles all the updating, except for moving the
+    # slider value. This is done because the slider widget is
+    # ultimately responsible for triggering an update, so all other
+    # updates are done by changing the slider widget value which
+    # then triggers this
+    def update(idx, im_idx=None):
+        # We have to explicitly make it an integer because the slider will
+        # output floats (even if they are still integer-valued)
+        idx = int(idx)
+
+        # Get the new data for this index
+        im = get_image_func(idx)
+        if len(im.shape) >= 3:
+            if im_idx == None and hasattr(axes[1],'image_idx'):
+                im_idx = axes[1].image_idx
+            elif im_idx == None:
+                im_idx=0
+            axes[1].image_idx = im_idx
+            axes[1].text_box.set_text(str(im_idx))
+            im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
+
+        # Now we resize the nanomap to show the new selection
+        axes[0].collections[0].set_sizes(calculate_sizes(idx))
+
+        # And we update the data in the image as well
+
+        ax_im = axes[1].images[-1]
+        ax_im.set_data(im)
+        update_colorbar(ax_im)
+
+
+    #
+    # Now we define the functions to handle various kinds of events
+    # that can be thrown our way
+    #
+
+    # We start by creating the slider here, so it can be used
+    # by the update hooks.
+    slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
+
+    # This handles scroll wheel and keypress events
+    def on_action(event):
+        im = get_image_func(0)
+        if event.inaxes is axes[1] and len(im.shape) >=3:
+            # This is triggered if the data to display has more than 2
+            # dimensions (i.e. is an image stack) and the event originates
+            # while the mouse is within the image display
+            im = im.reshape(-1,im.shape[-2],im.shape[-1])
+            im_idx = axes[1].image_idx
+
+            if event.key == 'up' or event.button == 'up' \
+               or event.key == 'left':
+                im_idx = (im_idx - 1) % im.shape[0]
+            if event.key == 'down' or event.button == 'down' \
+               or event.key == 'right':
+                im_idx = (im_idx + 1) % im.shape[0]
+
+            axes[1].image_idx=im_idx
+            slider.set_val(slider.val)#update(slider.val,im_idx=im_idx)
+            return # This prevents the rest from also happening
+
+        # Otherwise the if statements can throw errors when the
+        # event type isn't right, this way they just don't trigger
+        if not hasattr(event, 'button'):
+            event.button = None
+        if not hasattr(event, 'key'):
+            event.key = None
+
+        if event.key == 'up' or event.button == 'up' or event.key == 'left':
+            idx = slider.val - 1
+        elif event.key == 'down' or event.button == 'down' or event.key == 'right':
+            idx = slider.val + 1
+        else:
+            # This prevents errors from being thrown on irrelevant key
+            # or mouse input
+            return
+
+        # Handle the wraparound and trigger the update
+        idx = int(idx) % translations.shape[0]
+        slider.set_val(idx)
+
+    # This handles "pick" events in the nanomap
+    def on_pick(event):
+        # If we don't filter on type of event, this will also capture,
+        # for example, scroll events that happen over the nanomap
+        if event.mouseevent.button == 1:
+            slider.set_val(event.ind[0])
+
+
+    # Here we connect the various update functions
+    cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
+    cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
+    cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
+    # It's so dumb that matplotlib doesn't automatically track this for you
+    fig.nanomap_cids = [cid1,cid2,cid3]
+    slider.on_changed(update)
+
+    # Throw an extra update into the mix just to get rid of any things
+    # (like the nanomap dot sizes) that otherwise would change on the
+    # first update
+    update(0)
+
+    return fig

src/CDTools/tools/polarization/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.polarization.polarization import *

src/CDTools/tools/polarization/polarization.py

@@ -0,0 +1,217 @@
+import numpy as numpy
+import torch as t
+# Abe - again, we don't need math here. replace with torch-native functions
+# in all the definitions here. We have to use torch here if we want to be
+# able to calculate derivatives w.r.t. the polarizer angle, say.
+import math
+from math import sin
+from math import cos
+
+__all__ = ['apply_linear_polarizer',
+           'apply_phase_retardance',
+           'apply_half_wave_plate',
+           'apply_quarter_wave_plate',
+           'apply_circular_polarizer',
+           'apply_jones_matrix',
+           'generate_linear_polarizer',
+           'generate_birefringent_obj']
+
+
+def generate_linear_polarizer(pol_angle):
+    single_angle = False
+
+    pol_angle = t.as_tensor(pol_angle).to(dtype=t.float32)
+    if pol_angle.dim() == 0:
+        pol_angle = t.unsqueeze(pol_angle,0)
+        single_angle = True
+        
+    pol_angle_rad = t.deg2rad(pol_angle)
+    a = t.cos(pol_angle_rad) ** 2
+    b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)
+    c = b
+    d = t.sin(pol_angle_rad) ** 2
+    ab = t.stack((a, b), dim=-1)
+    cd = t.stack((c, d), dim=-1)
+    jones_matrices = t.stack((ab, cd), dim=-2)
+    if single_angle:
+        return jones_matrices[0].to(dtype=t.cfloat)
+    else:
+        return jones_matrices.to(dtype=t.cfloat)
+
+
+def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True):
+    """
+    Applies a linear polarizer to the probe
+
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe
+        The angle between the fast-axis of the linear polarizer and the horizontal axis
+    polarizer: t.Tensor
+        A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1))
+
+    Returns:
+    --------
+    linearly polarized probe: t.Tensor
+        (N)(P)x2x1xMxL
+    """
+    jones_matrices = generate_linear_polarizer(polarizer)
+    return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
+
+
+def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True):
+    # print('probe', probe.shape, 'jones matrix', jones_matrix.shape)
+    # if jones_matrix.shape == t.Size([5, 2, 2, 2, 2]):
+    #     print('probe', probe.shape, 'jones', jones_matrix.shape)
+    #     print('JONES', jones_matrix)
+    """
+    Applies a given Jones matrix to the probe
+
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (N)(P)x2xMxL tensor representing the probe
+    jones_matrix: t.tensor
+        (N)x2x2x(M)x(L)
+
+    Returns:
+    --------
+    a probe with the jones matrix applied: t.Tensor
+        (N)(P)x2xMxL
+
+    """
+    if transpose:
+        if jones_matrix.dim() < 4:
+            jones_matrix = jones_matrix[..., None, None]
+        if multiple_modes:
+            jones_matrix = jones_matrix.unsqueeze(-5)
+        probe = probe[..., None, :, :]
+        # if jones matrices do not differ from pattern to pattern
+        if probe.dim() > jones_matrix.dim():
+            jones_matrix = jones_matrix.unsqueeze(0)
+        # vice versa
+        elif jones_matrix.dim() > probe.dim():
+            probe = probe.unsqueeze(0)
+        # print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix)
+        jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
+        probe = probe.transpose(-1, -3).transpose(-2, -4)
+        output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
+
+    else:
+        raise NotImplementedError
+
+    return output
+
+
+def apply_phase_retardance(probe, phase_shift, multiple_modes=True):
+    """
+    Shifts the y-component of the field wrt the x-component by a given phase shift
+
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (...)x2x1xMxL tensor representing the probe
+    phase_shift: float
+        phase shift in degrees
+
+    Returns:
+    --------
+    probe: t.Tensor
+        (...)x2x1xMxL
+    """
+    theta = t.as_tensor(phase_shift, dtype=t.float32)
+    theta = t.deg2rad(theta)
+    probe = probe.to(dtype=t.cfloat)
+    jones_matrix = t.tensor([[1, 0], [0, t.exp(phase_shift)]]).to(dtype=t.cfloat)
+    polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
+
+    return polarized
+
+def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
+    """
+    Applies a circular polarizer to the probe
+
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (...)x2xMxL tensor representing the probe
+    left_polarizd: bool
+        True for the left-polarization, False for the right
+
+    Returns:
+    --------
+    circularly polarized probe: t.Tensor
+        (...)x2xMxL
+    """
+    probe = probe.to(dtype=t.cfloat)
+    if left_polarized:
+        jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat)
+    else:
+        jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat)
+    polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
+    return polarized
+
+def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
+    """
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
+    fast_axis_angle: float
+        The angle between the fast-axis of the polarizer and the horizontal axis
+
+    Returns:
+    --------
+    polarized probe: t.Tensor
+        (...)x2x1xMxL
+    """
+    probe = probe.to(dtype=t.cfloat)
+    theta = math.radians(fast_axis_angle)
+    exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
+    jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat)
+    out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
+
+    return out
+
+def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
+    """
+    Parameters:
+    ----------
+    probe: t.Tensor
+        A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
+    fast_axis_angle: float
+        The angle between the fast-axis of the polarizer and the horizontal axis
+
+    Returns:
+    --------
+    polarized probe: t.Tensor
+        (...)x2x1xMxL
+    """
+    probe = probe.to(dtype=t.cfloat)
+    theta = math.radians(fast_axis_angle)
+    exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
+    jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat)
+    out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
+
+    return out
+
+def generate_birefringent_obj(fast_axis=90, phase_ret=10, atten_fast=1, atten_ret=1, global_phase=0):
+    def to_rad(angle):
+        angle = t.as_tensor(angle, dtype=t.float32)
+        angle = t.deg2rad(angle)
+        return angle
+
+    fast_axis = to_rad(fast_axis)
+    phase_ret = to_rad(phase_ret)
+    global_phase = to_rad(global_phase)
+
+    def coord_rot(angle):
+        a = t.stack((t.cos(angle), t.sin(angle)), dim=-1)
+        b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1)
+        return t.stack((a, b), dim=-2).to(dtype=t.cfloat)
+
+    r1 = coord_rot(-fast_axis)
+    r2 = coord_rot(fast_axis)
+    p = t.exp(global_phase * 1j) * t.as_tensor([[atten_fast, 0], [0, atten_ret * t.exp(phase_ret*1j)]], dtype=t.cfloat)
+    return t.matmul(r1, t.matmul(p, r2))

src/CDTools/tools/propagators/__init__.py

@@ -0,0 +1 @@
+from CDTools.tools.propagators.propagators import *

src/CDTools/tools/propagators/propagators.py

@@ -0,0 +1,646 @@
+"""This module contains various propagators for light fields
+
+All the functions here are designed for use in an automatic differentiation
+ptychography model. Each function implements a different propagator.
+"""
+
+import torch as t
+from torch.nn.functional import grid_sample
+from scipy import fftpack
+import numpy as np
+from matplotlib import pyplot as plt
+
+__all__ = ['far_field', 'near_field',
+           'generate_angular_spectrum_propagator',
+           'inverse_far_field', 'inverse_near_field',
+           'generate_high_NA_k_intensity_map',
+           'high_NA_far_field',
+           'generate_generalized_angular_spectrum_propagator']
+
+def far_field(wavefront):
+    """Implements a far-field propagator in torch
+
+    This accepts a torch tensor, where the last dimension
+    represents the real and imaginary components of the wavefield,
+    and returns the far-field propagated version of it assuming it matches the
+    detector dimensions. It assumes that the
+    propagation is purely far-field, without checking that the geometry
+    is consistent with that assumption.
+
+
+    It also assumes that the real space wavefront is stored in an array
+    [i,j] where i corresponds to the y-axis and j corresponds to the
+    x-axis, with the origin following the CS standard of being in the
+    upper right. The zero frequency component of the propagated wavefield is
+    shifted to the center of the array.
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The JxNxMx2 stack of complex wavefronts to be propagated
+    
+    Returns
+    -------
+    propagated : torch.Tensor
+        The JxNxMx2 propagated wavefield
+    """
+    
+    shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
+    propagated = t.fft.fft2(shifted, norm='ortho')
+    return t.fft.fftshift(propagated, dim=(-1,-2))
+
+
+def inverse_far_field(wavefront):
+    """Implements the inverse of the far-field propagator in torch
+
+    This accepts a torch tensor, where the last dimension
+    represents the real and imaginary components of the propagated wavefield,
+    and returns the un-propagated array.
+
+    It assumes that the real space wavefront is stored in an array
+    [i,j] where i corresponds to the y-axis and j corresponds to the
+    x-axis, with the origin following the CS standard of being in the
+    upper right. The zero frequency component of the propagated wavefield is
+    assumed to be the center of the array.
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The JxNxMx2 stack of complex wavefronts propagated to the far-field
+    
+    Returns
+    -------
+    propagated : torch.Tensor
+        The JxNxMx2 exit wavefield
+    """
+    shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
+    propagated = t.fft.ifft2(shifted, norm='ortho')
+    return t.fft.fftshift(propagated, dim=(-1,-2))
+
+
+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
+    linear map between location on the detector and location in k-space
+    starts to break down. In addition, at angles above roughly 15 degrees,
+    a correction is needed to account for the decreasing solid angle
+    intersected by each pixel on the detector.
+    
+    This function generates a map which can be used to apply both corrections
+    via the high_NA_far_field propagator. The k-map which is output is
+    defined as a map between pixel location on the detector and locations
+    in the output of the standard, low-NA far-field propagated wavefield.
+    The output coordinate system is defined to run from -1 to 1 in both
+    directions - this allows for compatibility with pytorch's grid_sample
+    function. Some detector pixels may be mapped to values outside the
+    rangel [-1,1], depending on the respective sample and detector geometries.
+    This is most likely to be the case if the major correction is due to a
+    tilted sample.
+
+    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
+        The 3x2 sample basis, in real space
+    det_basis: array
+        The 3x2 detector basis, in real space
+    det_shape: array
+        The length-2 shape of the detector array, (N,M)
+    distance: float
+        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
+    -------
+    k_map : torch.Tensor
+        An NxMx2 tensor mapping detector pixels to locations in the low-NA propagated wavefield
+    intensity_map : torch.Tensor
+        An NxM tensor encoding the high-NA intensity correction
+    """
+
+    # First we convert things to numpy
+    try:
+        sample_basis = sample_basis.detach().cpu().numpy()
+    except:
+        pass
+
+    try:
+        det_basis = det_basis.detach().cpu().numpy()
+    except:
+        pass
+
+    det_shape = np.array(tuple(det_shape))
+
+    try:
+        distance = distance.detach().cpu().numpy()[0]
+    except:
+        pass
+
+    try:
+        wavelength = wavelength.detach().cpu().numpy()[0]
+    except:
+        pass
+
+    # The next order of business is to calculate the k values associated with
+    # each pixel.
+
+    i_arr = np.arange(det_shape[0])
+    i_arr = i_arr - np.mean(i_arr)
+    j_arr = np.arange(det_shape[1])
+    j_arr = j_arr - np.mean(j_arr)
+    Is, Js = np.meshgrid(i_arr,j_arr,indexing='ij')
+    samp_det_vec = np.cross(det_basis[:,0],det_basis[:,1])
+    samp_det_vec *= distance / np.linalg.norm(samp_det_vec)
+
+    # This could potentially correct for a mistake in the implied
+    # propagation direction (e.g. choosing e^ikx instead of e^-ikx)
+    # This appears to be correct, based on empirical evidence from
+    # a grazing incidence reflection experiment at 10 degrees
+    # on the optical table
+    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)
+
+    
+    # This accounts for the implied phase ramp along the exit wave direction
+    # In other words, it prevents the diffraction pattern from sliding off the
+    # detector when the sample is tilted but represented by an object with
+    # 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
+    # basis is the real space basis). In fact, because we want to return
+    # values scaled to the overall size of the k-space window, we don't
+    # even need the shape of the sample array
+
+    k_map = np.tensordot(2*sample_basis.transpose()[::-1,:] / (2*np.pi),Ks,axes=1)
+    k_map = t.Tensor(np.moveaxis(k_map,0,2)).to(*args, **kwargs)
+    
+    # Potentially we need a correction to account for the discrete nature
+    # of the FFT
+    
+    # And finally, we need to convert the results to pytorch
+    
+    return k_map, intensity_map
+
+
+
+
+def high_NA_far_field(wavefront, k_map, intensity_map=None):
+    """Performs a far-field propagation step including a correction for high-NA scenarios
+
+    Two major corrections need to be performed when propagating light fields
+    into the far field at high numerical aperture or when the sample is
+    tilted as compared to the detector. The first correction is a deviation
+    from the linear relationship between detector position and spatial
+    frequency in the near field. This is accounted for with the k_map
+    argument, as generated by the generate_high_NA_k_intensity_map 
+    function.
+
+    The second correction is the change in the solid angle which each pixel
+    subtends at high NA. This is accounted for with an optional intensity
+    map. This is kept optional because some detectors - specifically, those
+    for penetrating radiation - may either not need a correction or need
+    a different correction due to the volumetric nature of the pixels.
+    
+    If the k-map maps any pixels on the detector to pixels outside of the
+    k-space range of the wavefront, these will be set to zero. This is in
+    keeping with the typical assumption that the sample is band-limited to
+    the Nyquist frequency for the array on which it is sampled.
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The JxNxMx2 stack of complex wavefronts propagated to the far-field
+    k_map : torch.Tensor
+        The NxMx2 map accounting for high NA distortion, as generated by generate_high_NA_k_intensity_map
+    intensity_map : torch.Tensor
+        The optional NxM tensor accounting for the intensity variation across the detector
+
+    
+    Returns
+    -------
+    propagated : torch.Tensor
+        The JxNxMx2 exit wavefield
+
+
+    """
+    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):
+        # grid_sample doesn't work on complex-valued wavefields
+        real_output = grid_sample(low_NA_wavefield[None,:,:,:].real,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
+        imag_output = grid_sample(low_NA_wavefield[None,:,:,:].imag,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
+
+        result = real_output[0,:,:,:] + 1j * imag_output[0,:,:,:]
+    
+        if intensity_map is not None:
+            result = result * intensity_map[None,:,:]
+
+        return result
+
+    original_dim = wavefront.dim()
+    if original_dim == 2:
+        result = process_wavefield_stack(low_NA_wavefield[None,:,:])
+        return result[0,:,:]
+    if original_dim == 3:
+        result = process_wavefield_stack(low_NA_wavefield)
+        return result
+    if original_dim == 4:
+        result = []
+        for i in range(low_NA_wavefield.size()[0]):
+            result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:]))
+        return t.stack(result)
+    else:
+        raise IndexError('Wavefield had incorrect number of dimensions')
+
+
+
+def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, bandlimit=None, **kwargs):
+    """Generates an angular-spectrum based near-field propagator from experimental quantities
+    
+    This function generates an angular-spectrum based near field
+    propagator that will work on torch Tensors. The function is structured
+    this way - to generate the propagator first - because the
+    generation of the propagation mask is a bit expensive and if this
+    propagator is used in a reconstruction program, then it will be best
+    to calculate this mask once and reuse it.
+
+    Formally, this propagator is the the fourier transform of the convolution
+    kernel for light propagation in free space
+    
+    If the optional bandlimit parameter is set, the propagator will be set
+    to zero beyond an explicit bandlimiting frequency. This is helpful if the
+    propagator will be used in a repeated multiply/propagate framework such
+    as a multislice algorithm, where it helps to prevent aliasing.
+    
+    Parameters
+    ----------
+    shape : array
+        The shape of the arrays to be propagated
+    spacing : array
+        The pixel size in each dimension of the arrays to be propagated
+    wavelength : float
+        The wavelength of light to simulate propagation of
+    z : float
+        The distance to simulate propagation over
+    remove_z_phase : bool
+        Default False, whether to remove the dominant z-direction phase dependence
+    bandlimit : float
+        Optional, a fraction of the full detector radius beyond which to set the propagator to zero.
+
+    Returns
+    -------
+    propagator : torch.Tensor
+        A phase mask which accounts for the phase change that each plane wave will undergo.
+    """
+    
+    # Internally, the generalized propagation function is used, so we start
+    # by creating an appropriate basis
+    # This creates a real-valued tensor which matches the kind of complex
+    # number dtype requested in **kwargs
+    if 'dtype' in kwargs:
+        basis = t.real(t.zeros([3,2], **kwargs))
+    else:
+        basis = t.zeros([3,2], dtype=t.float32)
+    spacing = t.as_tensor(spacing, dtype=basis.dtype)
+
+    basis[1,0] = -spacing[0]
+    basis[0,1] = -spacing[1]
+    # And similarly, the offset is just z along the z direction
+    #offset = t.tensor([0,0,z], dtype=basis.dtype)
+    offset = t.zeros(3, dtype=basis.dtype)
+    offset[2] = z
+    
+    # And we call the generalized function! 
+    propagator = generate_generalized_angular_spectrum_propagator(shape, basis,
+                            wavelength, offset, 
+                            propagate_along_offset=remove_z_phase, **kwargs)
+    if z < 0:
+        propagator = t.conj(propagator)
+
+
+    # Bandlimiting is not implemented in the generalized function, because it
+    # has a less clear meaning in that setting, so we apply it here instead
+    if bandlimit is not None:
+        # No need to multiply by 2pi
+        ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
+        kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
+        Ki, Kj = t.meshgrid(ki,kj, indexing='ij')
+        min_radius = min(t.max(ki),t.max(kj))
+        Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2)
+        propagator = propagator * (Rs < bandlimit)
+        
+    return propagator
+
+
+
+def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset_vector, *args, propagation_vector=None, propagate_along_offset=False, **kwargs):
+    """Generates an angular-spectrum based near-field propagator from experimental quantities
+
+    This function generates an angular-spectrum based near field
+    propagator that will work on torch Tensors. The function is structured
+    this way - to generate the propagator first - because the
+    generation of the propagation mask is a bit expensive and if this
+    propagator is used in a reconstruction program, it will be best
+    to calculate this mask once and then reuse it it.
+
+    Formally, this propagator is the fourier transform of the convolution
+    kernel for light propagation in free space. It will map a light field
+    at an input plane, with the size and shape defined by the shape and basis
+    inputs, and map it to a plane of the same size and shape offset by the
+    offset vector. It is designed to work on any wavefield defined on an
+    array of parallelograms.
+
+
+    In addition, if the propagation_vector is set, there is an assumed
+    phase ramp  applied to the wavefield before propagation, defined such
+    that a feature with uniform phase will propagate along the direction of
+    the propagation vector. This will also remove the phase variation along
+    the propagation direction, because it makes the most physical sense to
+    regard this choice as removing the dominant phase variation in 3D, allowing
+    for the generation of a smoothly varying wavefield over 3D volumes.
+    This decision provides the best numerical stability and allows for the
+    simple setup of light fields copropagating with the coordinate system.
+    
+    If the propagate_along_offset option is set to True, then the propagation
+    vector will be set equal to the offset vector. This overrides the
+    propagation_vector option
+
+    Note that, unlike in the case of the simple angular spectrum propagator,
+    the direction of "forward propagation" is defined by the propagation
+    vector, or (if the propagation vector is not defined), the offset
+    vector. Therefore, in the simple case of a perpendicular offset,
+    there will be no difference between using an offset vector or the
+    negative of the  offset vector. This is because, for the light
+    propagation problem to be well posed, the assumption must be made that
+    light only passes through the plane of the known wavefield in one
+    direction. We always assume that light passes through the initial plane
+    travelling in the direction of the final plane.
+
+    Practically, if one wants to simulate inverse propagation, there are 
+    two possible approaches. First, one can use the inverse_near_field 
+    function, which simulates the inverse propagation problem and therefore
+    will naturally simulate propagation in the opposite direction. Second,
+    one can explicitly include a propagation_vector argument, in the direction
+    opposite to the offset vector. Note for both cases that this function
+    will only return propagators which suppress evanescent waves. Thus,
+    propagating forward and then backward by either of these two methods
+    will lead to a supression of evanescant waves. If you need a propagator
+    that will cause evanescent waves to undergo exponential growth, good
+    for you, but this function will not provide it for you.    
+
+    Parameters
+    ----------
+    shape : array
+        The shape of the arrays to be propagated
+    basis : array
+        The (2x3) set of basis vectors describing the array to be propagated
+    wavelength : float
+        The wavelength of light to simulate propagation of
+    offset_vector : array
+        The displacement to propagate the wavefield along.
+    propagation_vector : array
+        The vector along which to include an implied phase ramp to propagate uniform phase features along, if set
+    propagate_along_offset : bool
+        Overrides propagation_vector, sets the propagation vector to equal the offset vector if set.
+
+    Returns
+    -------
+    propagator : torch.Tensor
+        A phase mask which accounts for the phase change that each plane wave will undergo.
+    """
+
+    # make sure everything is in pytorch, and set the propagation vector
+    # appropriately if propagate_along_offset is chosen
+    basis = t.as_tensor(basis)
+    offset_vector = t.as_tensor(offset_vector, dtype=basis.dtype)
+
+    if propagate_along_offset:
+        propagation_vector = offset_vector
+        
+    if propagation_vector is not None:
+        propagation_vector = t.as_tensor(propagation_vector, dtype=basis.dtype)
+
+    #
+    # In this section, we calculate the wavevectors associated with each
+    # pixel in Fourier space. This is the meat of the function
+    #
+    
+    # First we calculate a dual basis for the real space grid
+    inv_basis =  t.linalg.pinv(basis).transpose(0,1)
+
+    # Then we calculate the frequencies in (i,j) space
+    ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype)
+    kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype)
+    K_ij = t.stack(t.meshgrid(ki,kj, indexing='ij'))
+    
+    # Now we convert these to frequencies in reciprocal space
+    # These frequencies span the 2D plane of the input wavefield,
+    # hence K_ip for "in-plane"
+    K_ip = t.tensordot(inv_basis, K_ij, dims=1)
+
+
+    # Now, we need to generate the out-of-plane direction, so we can
+    # expand these Ks to the full Ks in 3D reciprocal space.
+    
+    # This is broken down into 2 steps to avoid floating point underflow
+    # which was a real problem that showed up for electron ptycho
+    b1_dir = basis[:,0] / t.linalg.norm(basis[:,0])
+    b2_dir = basis[:,1] / t.linalg.norm(basis[:,1])
+    perpendicular_dir = t.cross(b1_dir,b2_dir)
+    # Note that we cannot use in-place operations if we want to be able to
+    # use automatic differentiation successfully
+    perpendicular_dir = perpendicular_dir / t.linalg.norm(perpendicular_dir)
+
+    # We set the sign of the propagation direction appropriately
+    if propagation_vector is not None:
+        perpendicular_dir = perpendicular_dir \
+            * t.sign(t.dot(perpendicular_dir,propagation_vector))
+    else:
+        pass
+        perpendicular_dir = perpendicular_dir * \
+            t.sign(t.dot(perpendicular_dir,offset_vector))
+    
+    # Then, if we have a propagation vector, we shift the in-plane
+    # components of all the pixels to be centered around the in-plane
+    # component of the propagation vector.
+    if propagation_vector is not None:
+        prop_dir = (propagation_vector /
+                    t.linalg.norm(propagation_vector))
+        K_0 = 2*np.pi / wavelength * prop_dir
+        K_0_ip = K_0 - t.dot(perpendicular_dir,K_0) * perpendicular_dir
+        K_ip += K_0_ip[:,None,None]
+    else:
+        K_0 = t.zeros_like(offset_vector)
+
+    # Now, we have accurate in-plane values for K, so we can calculate the
+    # out-of-plane part. We start by calculating it's squared magnitude
+    K_oop_squared = (2*np.pi/wavelength)**2 - t.linalg.norm(K_ip,dim=0)**2
+    
+    # Then, we take the square root and assign it the appropriate direction,
+    # adding to get the full 3D wavevectors. Note that we convert to complex
+    # before the square root to appropriately map negative numbers to
+    # complex frequencies
+    K = K_ip + perpendicular_dir[:,None,None] \
+        * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
+
+    # 
+    # In this section, we take the inner product of the calcualted
+    # wavevectors with the offset vector, to get the phase shift
+    # experienced by each plane wave.
+    #
+    
+    # We need to convert to complex because K is complex
+    offset_vector = t.complex(offset_vector,t.zeros_like(offset_vector))
+
+    # We then subtract off K_0, essentially setting the phase offset
+    # experienced by K_0 to 0. K_0 will already have been set to 0 if
+    # there was no propagation vector set.
+    K_m_K_0 = K - K_0[:,None,None]
+    # We actually calculate the phase mask
+    phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1)
+
+    # And we take the conjugate if needed, which makes sure that the
+    # imaginary part is always positive (supressing evanescent waves)
+    if propagation_vector is not None \
+       and t.sign(t.dot(propagation_vector,offset_vector.real)) == -1:
+        phase_mask = t.conj(phase_mask)
+
+    # We return the final mask, in the form asked for!
+    return t.exp(1j*phase_mask).to(*args,**kwargs)
+    
+
+def near_field(wavefront, angular_spectrum_propagator):
+    """ Propagates a wavefront via the angular spectrum method
+
+    This function accepts an 3D torch tensor, where the last dimension
+    represents the real and imaginary components of the wavefield, and
+    returns the near-field propagated version of it. It does this
+    using the supplied angular spectrum propagator, which is a premade
+    phase mask representing the Fourier transform of the kernel for
+    light propagation in the desired geometry.
+
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The (Leading Dims)xNxM stack of complex wavefronts to be propagated
+    angular_spectrum_propagator : torch.Tensor
+        The NxM phase mask to be applied during propagation
+
+    Returns
+    -------
+    propagated : torch.Tensor
+        The propagated wavefront 
+    """
+    return t.fft.ifft2(angular_spectrum_propagator * t.fft.fft2(wavefront))
+
+
+
+def inverse_near_field(wavefront, angular_spectrum_propagator):
+    """ Inverse propagates a wavefront via the angular spectrum method
+
+    This function accepts an 3D torch tensor, where the last dimension
+    represents the real and imaginary components of the wavefield, and
+    returns the near-field propagated version of it. It does this
+    using the supplied angular spectrum propagator, which is a premade
+    phase mask.
+
+    It propagates the wave using the complex conjugate of the supplied
+    phase mask. This corresponds to propagation backward across the original
+    propagation region - however, the treatment of evanescent waves is such
+    that evanescent waves will decay both during the forward propagation and
+    inverse propagation. This is done for reasons of numerical stability,
+    as the choice to magnify evanescent waves during the inverse propagation
+    process will quickly lead to magnification of any small amount of noise at 
+    frequencies larger than k_0, and in most typical situations will even
+    lead to overflow of the floating point range. If evanescent waves need
+    to be treated appropriately for any reason, it is recommended to use the
+    "magnify_evanescent" option in the appropriate helper function used to
+    generate the propagation phase mask. In this case, evanescent waves will
+    be magnified both when used with the forward and inverse near field
+    functions
+
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The JxNxMx2 stack of complex wavefronts to be propagated
+    angular_spectrum_propagator : torch.Tensor
+        The NxM phase mask to be applied during propagation
+
+    Returns
+    -------
+    propagated : torch.Tensor
+        The inverse propagated wavefront
+    """
+    return t.fft.ifft2(t.fft.fft2(wavefront)
+                       * t.conj(angular_spectrum_propagator))
+
+
+
+# I think it would be worthwhile to implement an FFT-DI based strategy as
+# well, especially for probe initialization where the propagation distance
+# can be large relative to what the angular spectrum method can reliably handle
+
+