Create a base version of the models

CDTools/datasets/polarized_ptycho_2d_dataset.py

@@ -115,7 +115,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
     # 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):
+    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
@@ -125,19 +125,22 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
         ----------
         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)
+                return cls.from_cxi(f, cut_zeros=cut_zeros)
 
         # Generate a base dataset
-        dataset = Ptycho2DDataset.from_cxi(cxi_file)
+        dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros)
 
         # Mutate the class to this subclass (PolarizedPtycho2DDataset)
         dataset.__class__ = cls

CDTools/datasets/ptycho_2d_dataset.py

@@ -118,7 +118,7 @@ class Ptycho2DDataset(CDataset):
     # 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):
+    def from_cxi(cls, cxi_file, cut_zeros=True):
         """Generates a new Ptycho2DDataset from a .cxi file directly
 
         This generates a new Ptycho2DDataset from a .cxi file storing
@@ -128,6 +128,8 @@ class Ptycho2DDataset(CDataset):
         ----------
         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
         -------
@@ -145,11 +147,10 @@ class Ptycho2DDataset(CDataset):
         dataset.__class__ = cls
 
         # Load the data that is only relevant for this class
-        patterns, axes = cdtdata.get_data(cxi_file)
+        patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
         translations = cdtdata.get_ptycho_translations(cxi_file)
         # And now re-do the stuff from __init__
         dataset.translations = t.tensor(translations, dtype=t.float32)
-
         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')

CDTools/models/__init__.py

@@ -22,7 +22,7 @@ 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']
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI']
 
 from CDTools.models.base import CDIModel
 from CDTools.models.simple_ptycho import SimplePtycho
@@ -32,6 +32,8 @@ 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

CDTools/models/base.py

@@ -586,7 +586,6 @@ class CDIModel(t.nn.Module):
 
             inputs, output = dataset[idx]
             sim_data = self.forward(*inputs).detach().cpu().numpy()
-            sim_data = sim_data
             meas_data = output.detach().cpu().numpy()
             if hasattr(self, 'mask') and self.mask is not None:
                 mask = self.mask.detach().cpu().numpy()

CDTools/models/fancy_ptycho.py

@@ -67,12 +67,11 @@ class FancyPtycho(CDIModel):
 
         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)
+                shape = self.probe[0][self.detector_slice].shape
             else:
-                background = 1e-6 * t.ones(self.probe[0].shape,
-                                           dtype=t.float32)
+                shape = self.probe[0].shape
+            background = 1e-6 * t.ones([s//oversampling for s in shape],
+                                       dtype=t.float32)
 
         self.background = t.nn.Parameter(background)
 

CDTools/models/polarized_fancy_ptycho.py

@@ -294,6 +294,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         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(

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,
+                }

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}
+

CDTools/tools/analysis/analysis.py

@@ -546,7 +546,6 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
                                         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)

CDTools/tools/data/data.py

@@ -317,7 +317,7 @@ def get_dark(cxi_file):
     return darks
 
 
-def get_data(cxi_file, cut_zeroes = True):
+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
@@ -333,6 +333,8 @@ def get_data(cxi_file, cut_zeroes = True):
     ----------
     cxi_file : h5py.File
         A file object to be read
+    cut_zeros : bool
+            Default True, whether to set all negative data to zero
 
     Returns
     -------
@@ -351,9 +353,8 @@ def get_data(cxi_file, cut_zeroes = True):
         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_zeroes:
+    if cut_zeros:
         np.maximum(data,0,data)
 
     if 'axes' in cxi_file[pull_from].attrs: