Finally finally finally add an RPI model. Makes it soooo much easier...

CDTools/models/__init__.py

@@ -52,7 +52,7 @@ from matplotlib.widgets import Slider
 from matplotlib import ticker
 import numpy as np
 
-__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho']
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI']
 
 
 class CDIModel(t.nn.Module):
@@ -111,7 +111,8 @@ class CDIModel(t.nn.Module):
     def save_results(self):
         raise NotImplementedError()
 
-    def AD_optimize(self, iterations, data_loader,  optimizer, scheduler=None):
+    def AD_optimize(self, iterations, data_loader,  optimizer,\
+                    scheduler=None, regularization_factor=None):
         """Runs a round of reconstruction using the provided optimizer
         
         This is the basic automatic differentiation reconstruction tool
@@ -130,13 +131,14 @@ class CDIModel(t.nn.Module):
             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
         """
         # First, calculate the normalization
         normalization = 0
         for inputs, patterns in data_loader:
             normalization += t.sum(patterns).cpu().numpy()
             
-            
         for it in range(iterations):
             loss = 0
             N = 0
@@ -150,6 +152,10 @@ class CDIModel(t.nn.Module):
                     else:
                         loss = self.loss(patterns,sim_patterns)
 
+                    if regularization_factor is not None \
+                       and hasattr(self, 'regularizer'):
+                        loss += self.regularizer(regularization_factor)
+                        
                     loss.backward()
                     return loss
 
@@ -162,7 +168,9 @@ class CDIModel(t.nn.Module):
             yield loss
 
 
-    def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False):
+    def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005,
+                      schedule=False, amsgrad=False, subset=None,
+                      regularization_factor=None):
         """Runs a round of reconstruction using the Adam optimizer
         
         This is generally accepted to be the most robust algorithm for use
@@ -182,7 +190,18 @@ class CDIModel(t.nn.Module):
             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
         """
+
+        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)
@@ -197,11 +216,14 @@ class CDIModel(t.nn.Module):
         else:
             scheduler = None
 
-        return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler)
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                scheduler=scheduler,
+                                regularization_factor=regularization_factor)
 
 
     def LBFGS_optimize(self, iterations, dataset, batch_size=None,
-                       lr=0.1,history_size=2):
+                       lr=0.1,history_size=2, subset=None,
+                       regularization_factor=None):
         """Runs a round of reconstruction using the L-BFGS optimizer
         
         This algorithm is often less stable that Adam, however in certain
@@ -221,7 +243,16 @@ class CDIModel(t.nn.Module):
             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
         """
+        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:
@@ -235,12 +266,13 @@ class CDIModel(t.nn.Module):
         optimizer = t.optim.LBFGS(self.parameters(),
                                   lr = lr, history_size=history_size)
 
-        return self.AD_optimize(iterations, data_loader, optimizer)
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                regularization_factor=regularization_factor)
 
 
     def SGD_optimize(self, iterations, dataset, batch_size=None,
                      lr=0.01, momentum=0, dampening=0, weight_decay=0,
-                     nesterov=False):
+                     nesterov=False, subset=None, regularization_factor=None):
         """Runs a round of reconstruction using the SGDoptimizer
         
         This algorithm is often less stable that Adam, but it is simpler
@@ -258,8 +290,18 @@ class CDIModel(t.nn.Module):
             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
         """
-        
+
+        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,
@@ -275,7 +317,8 @@ class CDIModel(t.nn.Module):
                                 weight_decay=weight_decay,
                                 nesterov=nesterov)
 
-        return self.AD_optimize(iterations, data_loader, optimizer)
+        return self.AD_optimize(iterations, data_loader, optimizer,
+                                regularization_factor=regularization_factor)
 
 
     # By default, the plot_list is empty
@@ -474,3 +517,4 @@ from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
 from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
 from CDTools.models.s_matrix_ptycho import SMatrixPtycho
 from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
+from CDTools.models.rpi import RPI

CDTools/models/bragg_2d_ptycho.py

@@ -196,7 +196,7 @@ class Bragg2DPtycho(CDIModel):
 
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False, opt_for_fft=False):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -223,7 +223,7 @@ class Bragg2DPtycho(CDIModel):
                                                    distance,
                                                    center=center,
                                                    padding=padding,
-                                                   opt_for_fft=False,
+                                                   opt_for_fft=opt_for_fft,
                                                    oversampling=oversampling)
         # now we grab the sample surface normal
         if hasattr(dataset, 'sample_info') and \

CDTools/models/fancy_ptycho.py

@@ -93,10 +93,10 @@ class FancyPtycho(CDIModel):
             self.obj_support = t.ones_like(self.obj)
 
         self.oversampling = oversampling
-        
+
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -123,7 +123,7 @@ class FancyPtycho(CDIModel):
                                                    distance,
                                                    center=center,
                                                    padding=padding,
-                                                   opt_for_fft=False,
+                                                   opt_for_fft=opt_for_fft,
                                                    oversampling=oversampling)
 
 

CDTools/models/rpi.py

@@ -0,0 +1,368 @@
+from __future__ import division, print_function, absolute_import
+
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import cmath
+from CDTools.tools import plotting as p
+from CDTools.tools.interactions import RPI_interaction
+from CDTools.tools import initializers
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+#
+# 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):
+    
+    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):
+
+        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[-2]/obj_guess.shape[-2],
+                                 probe.shape[-3]/obj_guess.shape[-3]])
+        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.BoolTensor(mask)
+        
+            
+        self.probe = probe.to(t.float32)
+
+        if obj_guess.dim() == 3:
+            obj_guess = obj_guess[None,:,:,:]
+            
+        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+        
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape[:-1])
+                
+        self.background = t.Tensor(background).to(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):
+        
+        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 = cmath.complex_to_torch(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(cmath.cabssq(probe))
+            obj_guess = scale * cmath.expi(2 * 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"')
+            
+            
+        # Maybe put something here to initialize an object support based on
+        # a probe threshold?
+        obj_support=None
+
+        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)
+
+
+    def random_init(self, pattern):
+        scale = t.sum(pattern) / t.sum(cmath.cabssq(self.probe))
+        self.obj.data = scale * cmath.expi(
+            2 * np.pi * t.rand(self.obj.shape[:-1])).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 = []
+        for i in range(self.probe.shape[0]):
+            pr = self.probe[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[None,:,:] * 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(1).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(cmath.cabssq(self.obj[0,:,:,:])) \
+            + factors[1] * t.sum(cmath.cabssq(self.obj[1:,:,:,:]))
+        
+    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(cmath.cabssq(self.probe).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 = cmath.torch_to_complex(self.probe.detach().cpu())
+        # Provide the option to save out the subdominant objects or
+        # just the dominant one
+        if full_obj:
+            obj = cmath.torch_to_complex(self.obj.detach().cpu())
+        else:
+            obj = cmath.torch_to_complex(self.obj[0].detach().cpu())
+        background = self.background.detach().cpu().numpy()**2
+        
+        return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
+                'probe': probe,'obj': obj,
+                'background': background}
+

CDTools/tools/initializers.py

@@ -9,11 +9,14 @@ import numpy as np
 import torch as t
 
 __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
-           'gaussian_probe', 'SHARP_style_probe'] 
+           'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init'] 
 
 from CDTools.tools import cmath
-from CDTools.tools.propagators import inverse_far_field, generate_angular_spectrum_propagator, near_field
+from CDTools.tools.propagators import *
+from CDTools.tools.analysis import orthogonalize_probes
 from scipy.fftpack import next_fast_len
+from scipy.sparse import linalg as spla
+from torch.nn.functional import pad
 import numpy as np
 
 
@@ -368,4 +371,68 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
     
     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]
     
+    pad0 = (probe.shape[-3] - obj_shape[0])//2
+    pad1 = (probe.shape[-2] - obj_shape[1])//2
+
+    def a_dagger(im):
+        im = cmath.complex_to_torch(im.reshape(obj_shape)).to(dtype=t.float32)
+        im = inverse_far_field(pad(far_field(im), (0,0,pad1,pad1,pad0,pad0)))
+        exit_wave = cmath.cmult(probe,im)
+        farfield = cmath.torch_to_complex(far_field(exit_wave))
+        return farfield.ravel()
+
+    def a(measured):
+        measured = cmath.complex_to_torch(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.float32)
+        im = inverse_far_field(measured)
+        multiplied = cmath.cmult(cmath.cconj(probe), im)
+        backplane = far_field(multiplied)
+        clipped = backplane[pad0:pad0+obj_shape[0],
+                            pad1:pad1+obj_shape[1],:]
+        return cmath.torch_to_complex(inverse_far_field(clipped)).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(cmath.cabssq(probe)).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 cmath.complex_to_torch(z0).to(dtype=t.float32)
+

CDTools/tools/interactions.py

@@ -10,11 +10,12 @@ from __future__ import division, print_function, absolute_import
 from CDTools.tools.cmath import *
 import torch as t
 import numpy as np
-
+from CDTools.tools import propagators 
 
 __all__ = ['translations_to_pixel', 'pixel_to_translations',
            'project_translations_to_sample',
-           'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc']
+           'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc',
+           'RPI_interaction']
 
 
 
@@ -524,6 +525,56 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
         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
+    pad0 = (probe.shape[-3] - obj.shape[-3])//2
+    pad1 = (probe.shape[-2] - obj.shape[-2])//2
+    
+    if obj.dim() == 3:
+        fftobj = t.nn.functional.pad(fftobj, (0, 0, pad1, pad1, pad0, pad0))
+    elif obj.dim() == 4:
+        fftobj = t.nn.functional.pad(fftobj,
+                                     (0, 0, pad1, pad1, pad0, pad0, 0, 0))
+    else:
+        raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.')
+        
+    # Again, just an inverse FFT but with an fftshift
+    upsampled_obj = propagators.inverse_far_field(fftobj)
+
+    if obj.dim() == 4:
+        return cmult(probe[None,...], upsampled_obj)
+    else:
+        return cmult(probe, upsampled_obj)

examples/transmission_RPI.py

@@ -0,0 +1,48 @@
+from __future__ import division, print_function, absolute_import
+
+import CDTools
+from matplotlib import pyplot as plt
+import pickle
+from torch.utils.data import Subset
+
+# First, we load an example dataset from a .cxi file
+
+ss_filename = 'example_data/Optical_Data_ss.cxi'
+
+with open('example_data/Optical_ptycho.pickle', 'rb') as f:
+    ptycho_results = pickle.load(f)
+
+probe = ptycho_results['probe']
+background = ptycho_results['background']
+
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename)
+
+# Next, we create a ptychography model from the dataset
+# Note that we explicitly as for two incoherent probe modes
+model = CDTools.models.RPI.from_dataset(dataset, probe, [800,800],
+                                        background=background, n_modes=2)
+
+
+# Let's do this reconstruction on the GPU, shall we? 
+model.to(device='cuda')
+dataset.get_as(device='cuda')
+
+# Note that the inspect step takes the vast majority of the time
+# The regularization is an L2 regularizer that empirically helps accelerate
+# convergence
+for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)):
+    model.inspect(dataset)
+    print(i,loss)
+
+# Now we use the regularizer to damp all but the top modes
+for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=0.4, regularization_factor=[0.001,0.1])):
+    model.inspect(dataset)
+    print(i,loss)
+
+results = model.save_results()
+
+
+# Finally, we plot the results
+model.inspect(dataset)
+model.compare(dataset)
+plt.show()