diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py
index 1063afc..e34350f 100644
--- a/CDTools/datasets/polarized_ptycho_2d_dataset.py
+++ b/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
diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py
index c7a9661..d684477 100644
--- a/CDTools/datasets/ptycho_2d_dataset.py
+++ b/CDTools/datasets/ptycho_2d_dataset.py
@@ -19,7 +19,8 @@ class Ptycho2DDataset(CDataset):
     It should save and load files compatible with most reconstruction
     programs, although it is only tested against SHARP.
     """
-    def __init__(self, translations, patterns, axes=None, *args, **kwargs):
+    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
@@ -52,23 +53,27 @@ class Ptycho2DDataset(CDataset):
         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, dtype=t.float32)
+        self.translations = t.tensor(translations)
         
-        self.patterns = t.as_tensor(patterns, dtype=t.float32)
-        if self.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')
+        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]
 
@@ -118,7 +123,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 +133,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 +152,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')
@@ -158,6 +164,12 @@ class Ptycho2DDataset(CDataset):
         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
 
 
@@ -187,6 +199,9 @@ class Ptycho2DDataset(CDataset):
             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
@@ -234,5 +249,5 @@ class Ptycho2DDataset(CDataset):
         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)
+        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)
 
diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py
index 49dcca2..5e005a1 100644
--- a/CDTools/models/__init__.py
+++ b/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
@@ -31,6 +31,9 @@ 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
diff --git a/CDTools/models/base.py b/CDTools/models/base.py
index 8eabaa6..d7f1203 100644
--- a/CDTools/models/base.py
+++ b/CDTools/models/base.py
@@ -11,7 +11,7 @@ simulate_to_dataset
     Creates a CDataset from the simulation defined in the model
 save_results
     Saves out a dictionary with the recovered parameters
-    
+
 
 Simulation
 ----------
@@ -55,8 +55,10 @@ class CDIModel(t.nn.Module):
 
     def __init__(self):
         super(CDIModel,self).__init__()
+
         self.loss_train = []
-        
+        self.iteration_count = 0
+
     def from_dataset(self, dataset):
         raise NotImplementedError()
 
@@ -79,11 +81,11 @@ class CDIModel(t.nn.Module):
 
     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.
         """
@@ -99,7 +101,7 @@ class CDIModel(t.nn.Module):
 
     def simulate_to_dataset(self, args_list):
         raise NotImplementedError()
-    
+
     def save_results(self):
         raise NotImplementedError()
 
@@ -107,10 +109,10 @@ class CDIModel(t.nn.Module):
                     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.
 
@@ -135,7 +137,7 @@ class CDIModel(t.nn.Module):
         normalization = 0
         for inputs, patterns in data_loader:
             normalization += t.sum(patterns).cpu().numpy()
-            
+
         def run_iteration(stop_event=None):
             loss = 0
             N = 0
@@ -144,7 +146,7 @@ class CDIModel(t.nn.Module):
                 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]),
@@ -152,7 +154,7 @@ class CDIModel(t.nn.Module):
                     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
@@ -168,7 +170,7 @@ class CDIModel(t.nn.Module):
                             loss = self.loss(pats,sim_patterns)
 
                         loss.backward()
-                    
+
                         total_loss += loss.detach()
 
                     if regularization_factor is not None \
@@ -178,7 +180,7 @@ class CDIModel(t.nn.Module):
                     return total_loss
 
                 loss += optimizer.step(closure).detach().cpu().numpy()
-                
+
             loss /= normalization
             if scheduler is not None:
                 scheduler.step(loss)
@@ -197,7 +199,7 @@ class CDIModel(t.nn.Module):
                     # 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)
@@ -205,10 +207,10 @@ class CDIModel(t.nn.Module):
                     calc.start()
                     while calc.is_alive():
                         if hasattr(self, 'figs'):
-                            self.figs[0].canvas.start_event_loop(0.01) 
+                            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.')
@@ -232,7 +234,7 @@ class CDIModel(t.nn.Module):
                       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
@@ -266,7 +268,7 @@ class CDIModel(t.nn.Module):
             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)
@@ -289,17 +291,17 @@ class CDIModel(t.nn.Module):
                                 calculation_width=calculation_width)
 
 
-    def LBFGS_optimize(self, iterations, dataset, 
+    def LBFGS_optimize(self, iterations, dataset,
                        lr=0.1,history_size=2, subset=None,
                        regularization_factor=None, thread=True,
-                       calculation_width=10):
+                       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
 
@@ -319,14 +321,14 @@ class CDIModel(t.nn.Module):
             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))
@@ -334,10 +336,11 @@ class CDIModel(t.nn.Module):
 
         # Define the optimizer
         optimizer = t.optim.LBFGS(self.parameters(),
-                                  lr = lr, history_size=history_size)
+                                  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,
@@ -349,7 +352,7 @@ class CDIModel(t.nn.Module):
                      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.
 
@@ -381,7 +384,7 @@ class CDIModel(t.nn.Module):
             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,
@@ -417,28 +420,28 @@ class CDIModel(t.nn.Module):
                   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), 
+        ( '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 
+        Where the third element in the tuple (a function that returns
         True if the plot is relevant) is not required.
 
         Parameters
@@ -447,9 +450,20 @@ class CDIModel(t.nn.Module):
             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
@@ -475,7 +489,7 @@ class CDIModel(t.nn.Module):
             plotter = plots[1]
 
             if figs is None:
-                fig = plt.figure()                
+                fig = plt.figure()
                 self.figs.append(fig)
             else:
                 fig = figs[idx]
@@ -483,13 +497,13 @@ class CDIModel(t.nn.Module):
             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
 
@@ -497,15 +511,15 @@ class CDIModel(t.nn.Module):
                 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.
 
@@ -519,7 +533,7 @@ class CDIModel(t.nn.Module):
         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
@@ -527,7 +541,7 @@ class CDIModel(t.nn.Module):
         extention : strategy
             Default is .eps, the file extension to save with.
         """
-        
+
         if hasattr(self, 'figs') and self.figs:
             figs = self.figs
         else:
@@ -537,21 +551,21 @@ class CDIModel(t.nn.Module):
             fig.savefig(prefix + fig.axes[0].get_title() + extension,
                         bbox_inches = 'tight')
 
-            
-    def compare(self, dataset):
+
+    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)
@@ -563,7 +577,7 @@ class CDIModel(t.nn.Module):
             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()))
 
@@ -571,16 +585,19 @@ class CDIModel(t.nn.Module):
             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()
-            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()
             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')
@@ -599,7 +616,7 @@ class CDIModel(t.nn.Module):
                 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]
@@ -613,8 +630,8 @@ class CDIModel(t.nn.Module):
                 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")
@@ -625,7 +642,7 @@ class CDIModel(t.nn.Module):
                 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':
@@ -636,5 +653,3 @@ class CDIModel(t.nn.Module):
         fig.canvas.mpl_connect('key_press_event',on_action)
         fig.canvas.mpl_connect('scroll_event',on_action)
         update(0)
-        
-
diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py
index 7a8553a..ca4996c 100644
--- a/CDTools/models/fancy_ptycho.py
+++ b/CDTools/models/fancy_ptycho.py
@@ -33,18 +33,19 @@ class FancyPtycho(CDIModel):
                  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 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'])
+        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)
 
@@ -76,12 +77,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)
 
@@ -115,6 +115,14 @@ class FancyPtycho(CDIModel):
 
         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'):
@@ -144,7 +152,8 @@ class FancyPtycho(CDIModel):
                      opt_for_fft=False,
                      fourier_probe=False,
                      loss='amplitude mse',
-                     units='um'
+                     units='um',
+                     simulate_probe_translation=False
                      ):
 
         wavelength = dataset.wavelength
@@ -248,6 +257,10 @@ class FancyPtycho(CDIModel):
             # 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:
@@ -277,7 +290,8 @@ class FancyPtycho(CDIModel):
                    probe_support=probe_support,
                    fourier_probe=fourier_probe,
                    oversampling=oversampling,
-                   loss=loss, units=units)
+                   loss=loss, units=units,
+                   simulate_probe_translation=simulate_probe_translation)
 
 
     def interaction(self, index, translations, *args):
@@ -305,8 +319,15 @@ class FancyPtycho(CDIModel):
             basis_prs = tools.propagators.inverse_far_field(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 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:
@@ -317,12 +338,24 @@ class FancyPtycho(CDIModel):
             # 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
 
 
@@ -352,16 +385,20 @@ class FancyPtycho(CDIModel):
         self.wavelength = self.wavelength.to(*args, **kwargs)
         # move the detector geometry too
         det_geo = self.detector_geometry
-        if hasattr(det_geo, 'distance'):
+        if 'distance' in det_geo:
             det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
-        if hasattr(det_geo, 'basis'):
+        if 'basis' in det_geo:
             det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
-        if hasattr(det_geo, 'corner'):
+        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)
@@ -369,7 +406,7 @@ class FancyPtycho(CDIModel):
         self.surface_normal = self.surface_normal.to(*args, **kwargs)
 
 
-    def sim_to_dataset(self, args_list):
+    def sim_to_dataset(self, args_list, calculation_width=None):
         # In the future, potentially add more control
         # over what metadata is saved (names, etc.)
 
@@ -395,9 +432,23 @@ class FancyPtycho(CDIModel):
         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(indices, translations)
+        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,
@@ -411,11 +462,15 @@ class FancyPtycho(CDIModel):
     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
+        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):
@@ -571,7 +626,10 @@ class FancyPtycho(CDIModel):
          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()
diff --git a/CDTools/models/multimode_rpi.py b/CDTools/models/multimode_rpi.py
new file mode 100644
index 0000000..2db9e21
--- /dev/null
+++ b/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}
+
diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py
index dada686..bbb2bb2 100644
--- a/CDTools/models/polarized_fancy_ptycho.py
+++ b/CDTools/models/polarized_fancy_ptycho.py
@@ -3,6 +3,7 @@ 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
@@ -21,22 +22,22 @@ class PolarizedFancyPtycho(FancyPtycho):
                  detector_slice=None,
                  surface_normal=np.array([0.,0.,1.]),
                  min_translation = t.Tensor([0,0]),
-                 background = None, translation_offsets=None, 
+                 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, obj_support=None, oversampling=1,
+                 probe_support = None, oversampling=1,
                  loss='amplitude mse',units='um'):
-        
-        super(FancyPtycho, self).__init__(wavelength, detector_geometry,
+
+        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 = None, translation_scale = 1, saturation=None,
-                 probe_support = None, obj_support=None, oversampling=1,
+                 weights = weights, translation_scale = 1, saturation=None,
+                 probe_support = None, oversampling=1,
                  loss='amplitude mse',units='um')
 
         if polarizer_offsets is None:
@@ -49,28 +50,224 @@ class PolarizedFancyPtycho(FancyPtycho):
         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):
-        
-        model = FancyPtycho.from_dataset(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', 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 
+        # Mutate the class to its subclass
         model.__class__ = cls
 
         if left_polarized:
             x = 1j
         else:
             x = -1j
-        model.probe.data = t.stack((model.probe.data.to(dtype=t.cfloat), x * model.probe.data.to(dtype=t.cfloat)), dim=-3)
-        obj = t.stack((model.obj.data, model.obj.data), dim=-3)
-        model.obj.data = t.stack((obj, obj), dim=-4)        
+
+        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
 
-    
-    # WHAT IS INDEX?
+    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
@@ -83,9 +280,9 @@ class PolarizedFancyPtycho(FancyPtycho):
         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 
+        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
@@ -96,17 +293,17 @@ class PolarizedFancyPtycho(FancyPtycho):
             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(
-            prs, self.obj_support * self.obj,pix_trans,
+            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, :, :, :]
@@ -122,7 +319,7 @@ class PolarizedFancyPtycho(FancyPtycho):
     def backward_propagator(self, wavefields):
         return tools.propagators.inverse_far_field(wavefields)
 
-    
+
     def measurement(self, wavefields):
         wavefields_x = wavefields[..., 0, :, :]
         wavefields_y = wavefields[..., 1, :, :]
@@ -145,15 +342,15 @@ class PolarizedFancyPtycho(FancyPtycho):
 
     # 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',
@@ -166,16 +363,16 @@ class PolarizedFancyPtycho(FancyPtycho):
         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)
 
@@ -187,14 +384,14 @@ class PolarizedFancyPtycho(FancyPtycho):
                                  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:
@@ -205,18 +402,18 @@ class PolarizedFancyPtycho(FancyPtycho):
         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
@@ -228,11 +425,11 @@ class PolarizedFancyPtycho(FancyPtycho):
             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()
@@ -248,7 +445,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         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
@@ -263,14 +460,14 @@ class PolarizedFancyPtycho(FancyPtycho):
             # 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)
 
@@ -287,63 +484,63 @@ class PolarizedFancyPtycho(FancyPtycho):
                 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, 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, 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, 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)),
+         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)),
+         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),
+         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)),
+         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)),
+         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)),
+         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()
@@ -352,7 +549,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         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,
diff --git a/CDTools/models/rpi.py b/CDTools/models/rpi.py
index ae0602c..a2e0e51 100644
--- a/CDTools/models/rpi.py
+++ b/CDTools/models/rpi.py
@@ -4,6 +4,7 @@ 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
 
@@ -49,11 +50,18 @@ 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):
+                 obj_support=None, oversampling=1, weight_matrix=False):
 
         super(RPI, self).__init__()
 
@@ -72,7 +80,7 @@ class RPI(CDIModel):
 
         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.obj_basis = self.probe_basis * scale_factor
         self.detector_slice = detector_slice
 
         # Maybe something to include in a bit
@@ -84,10 +92,10 @@ class RPI(CDIModel):
             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, :, :]
 
@@ -96,6 +104,13 @@ class RPI(CDIModel):
         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))
 
@@ -120,7 +135,7 @@ class RPI(CDIModel):
 
 
     @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):
+    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']
@@ -211,16 +226,24 @@ class RPI(CDIModel):
         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
+
+        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)
+                   obj_support=obj_support, oversampling=oversampling,
+                   weight_matrix=weight_matrix)
 
 
     def random_init(self, pattern):
@@ -251,8 +274,12 @@ class RPI(CDIModel):
 
         
         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 = self.probe[i]
+            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,
@@ -275,7 +302,6 @@ class RPI(CDIModel):
             output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
         except TypeError:
             pass
-        
         return output
 
 
@@ -303,8 +329,11 @@ class RPI(CDIModel):
         #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)
+        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)
@@ -336,7 +365,7 @@ class RPI(CDIModel):
     plot_list = [
         ('Root Sum Squared Amplitude of all Probes',
          lambda self, fig: p.plot_amplitude(
-             np.sqrt(np.sum((t.abs(self.probe)**2).cpu().numpy(),axis=0)),
+             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,
diff --git a/CDTools/models/time_resolved_ptycho_calibration.py b/CDTools/models/time_resolved_ptycho_calibration.py
new file mode 100644
index 0000000..eecadc4
--- /dev/null
+++ b/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,
+                }
diff --git a/CDTools/models/time_resolved_rpi.py b/CDTools/models/time_resolved_rpi.py
new file mode 100644
index 0000000..c8b52fb
--- /dev/null
+++ b/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}
+
diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py
index bef729b..e28d2d2 100644
--- a/CDTools/tools/analysis/analysis.py
+++ b/CDTools/tools/analysis/analysis.py
@@ -15,7 +15,8 @@ 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_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):
@@ -456,7 +457,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
     return cor
 
 
-def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
+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
@@ -479,6 +480,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
         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
     -------
@@ -507,14 +510,14 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
                           (im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
 
     if nbins is None:
-        nbins = np.max(im1[im_slice].shape) // 4
+        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)
 
-    cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \
-        t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2))
-
-    F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2
-    F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2
+    F1 = t.abs(f1)**2
+    F2 = t.abs(f2)**2
 
 
     di = np.linalg.norm(basis[:,0])
@@ -526,21 +529,41 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
     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)
 
-    numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy())
-    denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy())
-    denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy())
-    n_pix, bins = np.histogram(Rs,bins=nbins)
-
-    frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2))
+    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
 
-    threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
+    # 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)
@@ -581,6 +604,7 @@ def calc_vn_entropy(matrix):
         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
     
@@ -610,3 +634,326 @@ def calc_top_mode_fraction(matrix):
         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
+
+
diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py
index cd2da1a..be9785e 100644
--- a/CDTools/tools/data/data.py
+++ b/CDTools/tools/data/data.py
@@ -93,7 +93,7 @@ def get_sample_info(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']
 
@@ -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:
@@ -642,7 +643,8 @@ def add_dark(cxi_file, dark):
     d1.create_dataset('data_dark',data=dark)
 
 
-def add_data(cxi_file, data, axes=None):
+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:
@@ -673,7 +675,8 @@ def add_data(cxi_file, data, axes=None):
     if isinstance(data, t.Tensor):
         data = data.detach().cpu().numpy()
 
-    det1.create_dataset('data', data=data)
+    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:
diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py
index 5d6791c..c3afe16 100644
--- a/CDTools/tools/initializers/initializers.py
+++ b/CDTools/tools/initializers/initializers.py
@@ -633,7 +633,7 @@ def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
     dominant_fft = far_field(dominant_mode)
 
     shape = dominant_mode.shape
-    center = ((shape[-2]-1)//2, (shape[-1]-1)//2)
+    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,
diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py
index c50e095..5e850fd 100644
--- a/CDTools/tools/interactions/interactions.py
+++ b/CDTools/tools/interactions/interactions.py
@@ -17,14 +17,14 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
 
 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. 
-    
+    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
@@ -33,7 +33,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
     ----------
     basis : torch.Tensor
         The real space basis the wavefields are defined in
-    translations : torch.Tensor 
+    translations : torch.Tensor
         A Jx3 stack of real-space translations, or a single translation
     surface_normal : torch.Tensor
         Optional, the sample's surface normal
@@ -43,7 +43,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
     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)
@@ -68,18 +68,18 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
         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
@@ -96,7 +96,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
 
     Returns
     -------
-    real_translations : torch.Tensor 
+    real_translations : torch.Tensor
         A Jx3 stack of real-space translations, or a single translation
     """
     projection_1 = t.Tensor([[1,0,0],
@@ -129,7 +129,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
 
 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
@@ -145,7 +145,7 @@ def project_translations_to_sample(sample_basis, translations):
     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].
 
@@ -153,7 +153,7 @@ def project_translations_to_sample(sample_basis, translations):
     ----------
     sample_basis : torch.Tensor
         The real space basis the wavefields are defined in
-    translations : torch.Tensor 
+    translations : torch.Tensor
         A Jx3 stack of real-space translations, or a single translation
 
     Returns
@@ -171,7 +171,7 @@ def project_translations_to_sample(sample_basis, translations):
 
 
     # 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)
@@ -179,13 +179,13 @@ def project_translations_to_sample(sample_basis, translations):
     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)
@@ -200,7 +200,7 @@ def project_translations_to_sample(sample_basis, translations):
         device=translations.device,
         dtype=translations.dtype)
 
-    
+
     sample_projection = t.mm(basis_vectors_inv, projector).t()
     prop_projection = t.mm(propagation_dir_inv, projector).t()
 
@@ -218,9 +218,9 @@ def project_translations_to_sample(sample_basis, translations):
         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
@@ -229,15 +229,15 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
     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
@@ -251,7 +251,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
 
     Returns
     -------
-    exit_waves : torch.Tensor 
+    exit_waves : torch.Tensor
         An (N)x(P)xMxL tensor of the calculated exit waves
     """
 
@@ -260,9 +260,9 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
         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]
@@ -292,7 +292,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
 
 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
@@ -322,7 +322,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
     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)
@@ -342,15 +342,15 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
             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,
@@ -359,16 +359,16 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
             # 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]]
 
@@ -387,7 +387,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
 
 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
@@ -423,26 +423,27 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
     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,:]
+        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)
-
-    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
-                              tr[1]:tr[1]+probe.shape[-1]]
-                          for tr in integer_translations])
-    if polarized:
-        selections = t.stack([obj[:, :,tr[0]:tr[0]+probe.shape[-2],
-                          tr[1]:tr[1]+probe.shape[-1]]
+    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) \
@@ -454,15 +455,11 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
         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,:,:,:]
@@ -470,11 +467,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
                 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,:,:]
@@ -483,12 +477,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
         # selections: Nx2x2xMxL
         # probe: Nx(P)x2x1xMxL
         else:
-            # print('F', shift_probe)
-            output = polarization.apply_jones_matrix(shifted_probe, selections)
+            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:
@@ -497,7 +489,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
 
 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
@@ -513,7 +505,7 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
     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 
+    this change has now been made
 
     Parameters
     ----------
@@ -537,17 +529,17 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
     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
@@ -556,14 +548,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
         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]
 
@@ -572,14 +564,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
                                   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')
 
@@ -587,11 +579,11 @@ 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
@@ -601,7 +593,7 @@ def RPI_interaction(probe, obj):
     method and is not commonly used elsewhere.
 
     This also works with object functions that have an extra first dimension
-    for an incoherently mixing model. 
+    for an incoherently mixing model.
 
 
     Parameters
@@ -618,17 +610,20 @@ def RPI_interaction(probe, obj):
     """
 
     # 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
-    pad0l = (probe.shape[-2] - obj.shape[-2])//2
-    pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
-    pad1l = (probe.shape[-1] - obj.shape[-1])//2
+    # 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, pad0l, pad0r))
-        
+    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)
 
diff --git a/CDTools/tools/plotting/__init__.py b/CDTools/tools/plotting/__init__.py
index dd4e060..3ed18c3 100644
--- a/CDTools/tools/plotting/__init__.py
+++ b/CDTools/tools/plotting/__init__.py
@@ -1 +1,2 @@
 from CDTools.tools.plotting.plotting import *
+from CDTools.tools.plotting.polarized_plotting import *
diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py
index bee128c..07177fb 100644
--- a/CDTools/tools/plotting/plotting.py
+++ b/CDTools/tools/plotting/plotting.py
@@ -84,7 +84,7 @@ def get_units_factor(units):
 
 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.
 
@@ -94,7 +94,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
     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
     ----------
@@ -122,7 +122,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
     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,
@@ -144,7 +144,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
         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
@@ -153,9 +153,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
         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):
@@ -183,7 +183,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
             plt.xlabel('j (pixels)')
             plt.ylabel('i (pixels)')
 
-            
+
         plt.title(title)
 
         if len(im.shape) >= 3:
@@ -196,14 +196,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
         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':
@@ -219,9 +219,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
         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
@@ -258,7 +258,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
     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):
@@ -306,7 +306,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
 
     If a basis is explicitly passed, the image will be plotted in real-space
     coordinates.
-    
+
     Parameters
     ----------
     im : array
@@ -418,7 +418,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
                       units=units, **kwargs)
 
 
-def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **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
@@ -431,6 +431,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
         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)
 
@@ -455,6 +457,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
 
     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 + ')')
@@ -463,7 +468,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
     return fig
 
 
-def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'):
+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
@@ -478,6 +483,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
         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
     -------
@@ -511,7 +518,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
     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 + ')')
@@ -522,12 +531,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
 
 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
@@ -560,11 +569,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         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
@@ -573,29 +582,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         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],
@@ -616,7 +625,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         # 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)
@@ -636,7 +645,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         # 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:
@@ -657,9 +666,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         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),
@@ -667,7 +676,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
     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
@@ -677,7 +686,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         # 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:
@@ -688,22 +697,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
             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")
@@ -717,7 +726,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
             # 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]
@@ -735,7 +744,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
             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':
@@ -744,7 +753,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
             # 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)
@@ -755,17 +764,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
         # 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] 
+    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
diff --git a/CDTools/tools/plotting/polarized_plotting.py b/CDTools/tools/plotting/polarized_plotting.py
new file mode 100644
index 0000000..e17cc16
--- /dev/null
+++ b/CDTools/tools/plotting/polarized_plotting.py
@@ -0,0 +1,342 @@
+import torch as t
+from matplotlib import pyplot as plt
+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
+from CDTools.tools.polarization import polarization
+
+all = ['visualize_components_amplitudes',
+        'visualize_attenuations',
+        'visualize_fast_axes',
+        'visualize_global_axes',
+        'visualize_phase_ret',
+        'visualize_probe']
+
+def iterator(a, b):
+    '''
+    A helper function we can use to facilitate to process of iterating over 2D arrays
+    '''
+    x, y = t.arange(a), t.arange(b)
+    x, y = t.meshgrid(x, y)
+    x, y = t.ravel(x), t.ravel(y)
+    return (x, y)
+
+def plot_probe_ellipse(a=1, b=1, phase_ret=0, scale=1, x0=4, y0=5):
+    '''
+    Given a probe vector at a point (x0, y0), visualizes ellipticity
+    of its polarization
+
+    Parameters:
+    -----------
+    a : 1D np.array
+        An amplitude of the horizontal component of the probe vector
+    b : 1D np.array
+        An amplitude of the vertical component of the probe vector
+    phase_ret : 1D np.array
+        A phase difference in radians bettween the phases of the y and x components
+    scale : int
+        A scaling factor
+    x0, y0: 1D np.array or float
+        Defines the location of the vector to be plotted
+
+    Returns:
+    --------
+    x, y set of points to plot a single ellipse
+    '''
+    theta = np.linspace(0, 2*np.pi, 20)
+    x = x0 + scale * a * np.real(np.exp(1j * theta))
+    y = y0 + scale * b * np.real(np.exp(1j * (theta + phase_ret)))
+    return x, y
+
+def plot_attenuations(atten_slow=1, atten_fast=1, fast_ax_angle=0, scale=1, x0=4, y0=4):
+    '''
+    Plots attenuations along the fast and slow axes
+
+    Parameters:
+    -----------
+    atten_slow: 1D np.array
+        Attenuation along the slow axis
+    atten_fast: 1D np.array
+        Attenuation along the fast axis
+    fast_ax_angle: 1D np.array
+        An angle between the fast horizontal and fast axes
+    phase_ret: 1D np.array
+        A difference in phases gained by the slow and the fast components
+        The most clockwise axis is always considered to be the fast one
+    scale: 1D np.array
+        A scaling factor
+    x0, y0: 1D np.array or float
+        Defines the location to be plotted at
+
+    Returns:
+    --------
+    x, y set of points to plot a single Jones matrix of the object
+    '''
+    angle = fast_ax_angle
+    theta = np.linspace(0, 2*np.pi, 20)
+    # collection of points to plot a fast axis
+    print('angle', np.rad2deg(angle))
+    x = x0 + scale * atten_fast * np.real(np.exp(1j * theta))
+    x_f = x0 + (x - x0) * np.cos(angle)
+    y_f = y0 + (x - x0) * np.sin(angle)
+    # collection of point to plot a slow axis
+    y = y0 + scale * atten_slow * np.real(np.exp(1j * theta))
+    x_s = x0 - (y - y0) * np.sin(angle)
+    y_s = y0 + (y - y0) * np.cos(angle)
+    return x_f, y_f, x_s, y_s
+
+def plot_fast_axis(fast_ax_angle=0, scale=1, x0=4, y0=4):
+    '''
+    Plots directions of the fast axes only
+    '''
+    xf, yf, xs, ys = plot_attenuations(fast_ax_angle=fast_ax_angle, atten_fast=1, atten_slow=0, scale=scale, x0=x0, y0=y0)
+    return xf, yf
+
+def plot_figures(shape, num_of_el_along_x=20, num_of_el_along_y=20,
+                 phases=None, fast_ax_angles=None,
+                 atten_fast=None, atten_slow=None, scale=5,
+                 probe=False, attenuations=False, fast_axes=False):
+    """
+    All the parameters - np.arrays of shape (shape)
+    """
+    x_centers = np.linspace(1, shape[0] - 1, num_of_el_along_x)
+    y_centers = np.linspace(1, shape[1] - 1, num_of_el_along_y)
+    X, Y = np.meshgrid(x_centers, y_centers)
+    X, Y = np.ravel(X), np.ravel(Y)
+    xs, ys = np.array([]), np.array([])
+    for x, y in zip(X, Y):
+        k, m = int(x), int(y)
+        if probe:
+            xx, yy = plot_probe_ellipse(a=atten_fast[k, m], b=atten_slow[k, m],
+                                        phase_ret=phases[k, m], scale=scale, x0=x, y0=y)
+            title = 'Polarized Probe'
+        elif attenuations:
+            xf, yf, xs, ys = plot_attenuations(atten_slow=atten_slow[k, m], atten_fast=atten_fast[k, m],
+                                       fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
+        elif fast_axes:
+            xx, yy = plot_fast_axis(fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
+            title = 'Fast Axes'
+
+        if attenuations:
+            plt.plot(xf, yf, c='b')
+            plt.plot(xs, ys, c='b')
+            plt.axis('equal')
+            plt.title('Attenuations')
+        else:
+            plt.plot(xx, yy, c='b')
+            plt.axis('equal')
+            plt.title(title)
+
+"""
+VISUALIZATION FUNCTIONS (and helper functions)
+"""
+def determine_fast_axis(angle0, angle1):
+    # in radians
+    w0, w1 = t.as_tensor(angle0, dtype=t.float32), t.as_tensor(angle1, dtype=t.float32)
+    diff = (w0 - w1) % (2 * np.pi)
+    if t.tensor(0, dtype=t.float32) <= diff and diff <= t.tensor(np.pi, dtype=t.float32):
+        fast = 1
+        phase_ret = (w0 - w1) % (2 * np.pi)
+        gl_phase = w1
+    else:
+        fast = 0
+        phase_ret = (w1 - w0) % (2 * np.pi)
+        gl_phase = w0
+
+    if t.allclose((w0 - w1) % (2 * np.pi), phase_ret):
+        gl_phase = w1
+        fast = 1
+
+    return fast, phase_ret, gl_phase
+
+def retrieve_obj_info(obj):
+    obj_t = obj.transpose(-1, -3).transpose(-2, -4)
+    w, v = t.linalg.eig(obj_t)
+    b = obj.shape[-1]
+    a = obj.shape[-2]
+    eigenvectors = t.empty(a, b, 2, dtype=t.cfloat)
+    ret_phases = t.empty(a, b, dtype=t.float32)
+    global_phases = t.empty(a, b, dtype=t.float32)
+    atten_fast = t.empty(a, b, dtype=t.float32)
+    atten_slow = t.empty(a, b, dtype=t.float32)
+    fast_ax_angles = t.empty(a, b, dtype=t.float32)
+
+    for k, m in zip(*iterator(a, b)):
+        angle0, angle1 = t.angle(w[k, m, 0]).to(dtype=t.float32), t.angle(w[k, m, 1]).to(dtype=t.float32)
+        fst, phase_ret, gl_phase = determine_fast_axis(angle0, angle1)
+        atten0, atten1 = t.abs(w[k, m, 0]), t.abs(w[k, m, 1])
+        ret_phases[k, m] = phase_ret
+        def fast(b):
+            if b == 0:
+                eigenvectors[k, m, :] = v[k, m, 0, :]
+                atten_fast[k, m] = atten0
+                atten_slow[k, m] = atten1
+                global_phases[k, m] = angle0
+            elif b == 1:
+                eigenvectors[k, m, :] = v[k, m, 1, :]
+                atten_fast[k, m] = atten1
+                atten_slow[k, m] = atten0
+                global_phases[k, m] = angle1
+
+        if t.allclose(angle0, angle1):
+        # then it's a linear polarizer, fast ax - the one for which attenuation is bigger
+            if atten0 > atten1:
+                fst = 0
+            else:
+                fst = 1
+        fast(fst)
+
+        cos = np.abs(eigenvectors[k, m, 0])
+        sin = np.abs(eigenvectors[k, m, 1])
+        if t.allclose(cos, t.zeros(1, dtype=t.float32)):
+            fast_ax_angles[k, m] = 90
+        else:
+            fast_ax_angles[k, m] = t.atan(sin/cos)
+
+    fast_ax_angles = np.asarray(fast_ax_angles, dtype=np.float32)
+    ret_phases = np.asarray(ret_phases, dtype=np.float32)
+    global_phases = np.asarray(global_phases, dtype=np.float32)
+    atten_fast = np.asarray(atten_fast, dtype=np.float32)
+    atten_slow = np.asarray(atten_slow, dtype=np.float32)
+
+    return fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow
+
+
+def visualize_components_amplitudes(obj, rot_angle=0, logarithmic=False):
+    # coord_rot angle is the only angle in degrees here
+    def coord_rot(angle):
+            angle = t.as_tensor(angle, dtype=t.float32)
+            angle = t.deg2rad(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)
+    for k, m in zip(*iterator(obj.shape[-2], obj.shape[-1])):
+        obj[:, :, k, m] = t.matmul(coord_rot(rot_angle), obj[:, :, k, m])
+    components = [obj[i, j, :, :] for i, j in zip(*iterator(2, 2))]
+    if logarithmic:
+        components = [np.log(comp)/np.log(10) for comp in components]
+    titles = ['Amptlitudes of the a components', 'Amplitudes of the b components',
+              'Amplitudes of the c components', 'Amplitudes of the d components']
+
+    for i in range(4):
+        amplitude = np.abs(components[i])
+        plt.imshow(module)
+        plt.colorbar()
+        plt.title(titles[i])
+
+def visualize_phase_ret(obj, logarithmic=False):
+    print('DHSBCUYLIWGBCGLWIYV')
+    fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
+    if logarithmic:
+        ret_phases = np.log(ret_phases)/np.log(10)
+    plt.imshow(ret_phases)
+    plt.colorbar()
+    plt.show()
+
+def visuallize_global_phases(obj, logarithmic=False):
+    fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
+    if logarithmic:
+        global_phases = np.log(global_phases)/np.log(10)
+    plt.imshow(global_phases)
+    plt.colorbar()
+    plt.show()
+
+def visualize_fast_axes(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
+    fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
+    A, B = obj.shape[-2], obj.shape[-1]
+    plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
+                 fast_ax_angles=fast_ax_angles, scale=scale, fast_axes=True)
+
+def visualize_attenuations(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
+    fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
+    A, B = obj.shape[-2], obj.shape[-1]
+    plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
+                 phases=ret_phases, atten_slow=atten_slow, atten_fast=atten_fast,
+                 fast_ax_angles=fast_ax_angles, scale=scale, attenuations=True)
+
+def visualize_probe(probe, scale=1, num_of_el_along_x=20, num_of_el_along_y=20):
+    a = np.abs(probe[..., 0, :, :])
+    b = np.abs(probe[..., 1, :, :])
+    phases = np.angle(probe[..., 1, :, :]) - np.angle(probe[..., 0, :, :])
+    A, B = np.asarray(probe.shape[-2]), np.asarray(probe.shape[-1])
+    plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
+                     phases=phases, atten_fast=a, atten_slow=b, scale=scale,
+                     probe=True)
+
+
+#
+# def object_from_components(a, b, c, d):
+#     ab = t.stack((a, b), dim=-3)
+#     cd = t.stack((c, d), dim=-3)
+#     return t.stack((ab, cd), dim=-4)
+#
+# def object_from_quarters(a, b, c, d):
+#     ab = t.cat((a, b), dim=-1)
+#     cd = t.cat((c, d), dim=-1)
+#     return t.cat((ab, cd), dim=-2)
+#
+# def generate_birefringent_obj(shape, func_axes=None, func_ret_phases=None, func_global_phases=None, fast_axes=[0, 0, 90, 90], phases=[0, 18, 40, 18]):
+#     ret = [polarization.generate_birefringent_obj(fast_axis=i, phase_ret=j) for i, j in zip(fast_axes, phases)]
+#     A = shape[0]
+#     B = shape[1]
+#     def to_rad(angle):
+#         angle = t.as_tensor(angle)
+#         return t.deg2rad(angle)
+#
+#     if func_axes is None:
+#         # 4 sets of components for each quarter
+#         components = [[ret[i][j][k].repeat(A//2, B//2) for j, k in zip(*iterator(2, 2))] for i in range(4)]
+#         # build the quarters from the components
+#         quarters = [object_from_components(*comp) for comp in components]
+#         obj = object_from_quarters(*quarters)
+#     else:
+#         X, Y = t.arange(A), t.arange(B)
+#         X, Y = t.meshgrid(X, Y)
+#         obj = t.empty(2, 2, A, B, dtype=t.cfloat)
+#         for i, j in zip(*iterator(A, B)):
+#             x, y = X[i, j], Y[i, j]
+#             axis = func_axes(x, y)
+#             ret = func_ret_phases(x, y)
+#             gl = func_global_phases(x, y)
+#             obj[:, :, i, j] = polarization.generate_birefringent_obj(fast_axis=axis, phase_ret=ret, global_phase=gl)
+#             w, v = t.linalg.eig(obj[:, :, i, j])
+#
+#     # t.tensor of shape (2, 2, A, B)
+#     return obj
+#
+# def axes(x, y):
+#     # return (x-25)**2 + (y-25)**2
+#     return  x * 20
+#
+# def phases(x, y):
+#     # return ((x - 5) ** 2 + (y - 5) ** 2) * 20
+#     return  (x) * 10
+#
+# def glob(x, y):
+#     return x * 10
+#
+# def amp(x, y):
+#     return 1 + 0.05 * (x + y)
+#
+# def build_probe(shape, phase_ret_func=None, amps_func=None):
+#     probe = t.empty(2, shape[0], shape[1], dtype=t.cfloat)
+#     for x, y in zip(*iterator(probe.shape[-2], probe.shape[-1])):
+#         phase = phase_ret_func(x, y)
+#         phase = t.deg2rad(phase).to(dtype=t.cfloat)
+#         abs = amps_func(x, y).to(dtype=t.cfloat)
+#         probe[0, x, y] = abs
+#         probe[1, x, y] = abs * t.exp(phase * 1j)
+#
+#     return probe
+#
+#
+# obj = generate_birefringent_obj((10, 10), func_axes=axes, func_ret_phases=phases, func_global_phases=glob)
+# probe = build_probe((10, 10), phase_ret_func=phases, amps_func=amp)
+# # visualize_probe(probe, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
+# visualize_fast_axes(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
+# # visualize_phase_ret(obj)
+# # visuallize_global_phases(obj)
+# # visualize_attenuations(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.3)
+# plt.show()
diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py
index f74d108..1374895 100644
--- a/CDTools/tools/polarization/polarization.py
+++ b/CDTools/tools/polarization/polarization.py
@@ -12,29 +12,32 @@ __all__ = ['apply_linear_polarizer',
            'apply_half_wave_plate',
            'apply_quarter_wave_plate',
            'apply_circular_polarizer',
-           'apply_jones_matrix']
+           'apply_jones_matrix',
+           'generate_linear_polarizer',
+           'generate_birefringent_obj']
 
 
-# Abe - split these into two functions
-
-# Note for the future: this function should
 def generate_linear_polarizer(pol_angle):
     single_angle = False
 
-    pol_angle = t.as_tensor(pol_angle)
+    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)
-    jones_matrices = t.stack([t.tensor([[(t.cos(p)) ** 2, t.sin(p) * t.cos(p)],
-                                        [t.sin(p) * t.cos(p), (t.sin(p)) ** 2]])
-                              for p in polarizer])
+    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) 
+        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):
     """
@@ -51,13 +54,17 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True
     Returns:
     --------
     linearly polarized probe: t.Tensor
-        (N)(P)x2x1xMxL 
+        (N)(P)x2x1xMxL
     """
-    jones_matrices = generate_linear_polarizer(polarization)
+    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
 
@@ -66,64 +73,40 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
     probe: t.Tensor
         A (N)(P)x2xMxL tensor representing the probe
     jones_matrix: t.tensor
-        (N)x2x2x(M)x(L) 
+        (N)x2x2x(M)x(L)
 
     Returns:
     --------
     a probe with the jones matrix applied: t.Tensor
-        (N)(P)x2xMxL 
+        (N)(P)x2xMxL
 
-    Assume that if the probe has a dimension (N), so does the jones matrix
     """
-    if multiple_modes:
-        if transpose:
-            if len(jones_matrix.shape) >= 4:
-                jones_matrix = jones_matrix[..., None, :, :, :, :]  
-
-            else:
-                jones_matrix = jones_matrix[..., None, :, :, None, None]
-            probe = probe[..., None, :, :]
-            # if jones matrices do not differ from pattern to pattern
-            if len(probe.shape) > len(jones_matrix.shape):
-                jones_matrix = jones_matrix[None, ...]
-            jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) 
-            # (N)1xMxLx2x2 or (N)1x1x1x2x2
-            probe = probe.transpose(-1, -3).transpose(-2, -4)
-            output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
-            # (N)Px2xMxL
-
-        else:
-            if len(jones_matrix.shape) < 4:
-                jones_matrix = jones_matrix[..., None, :, :, :, :]
-            probe = t.stack((probe, probe), dim=-4)
-            output = t.sum(jones_matrix * probe, dim=-3)
-            #(N)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:
-        if transpose:   
-            if len(jones_matrix.shape) < 4:
-                jones_matrix = jones_matrix[..., None, None]
-            probe = probe[..., None, :, :]
-            # if jones matrices do not differ from pattern to pattern
-            if len(probe.shape) > len(jones_matrix.shape):
-                jones_matrix = jones_matrix[None, ...]
-            probe = probe.transpose(-1, -3).transpose(-2, -4)
-            jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
-            output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
+        raise NotImplementedError
 
-
-        else:
-            if len(jones_matrix.shape) < 4:
-                jones_matrix = jones_matrix[..., None, None]
-            probe = t.stack((probe, probe), dim=-4)
-            output = t.sum(jones_matrix * probe, dim=-3)
-    
     return output
 
 
-def apply_phase_retardance(probe, phase_shift):
+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 
+    Shifts the y-component of the field wrt the x-component by a given phase shift
 
     Parameters:
     ----------
@@ -135,44 +118,41 @@ def apply_phase_retardance(probe, phase_shift):
     Returns:
     --------
     probe: t.Tensor
-        (...)x2x1xMxL 
+        (...)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, phase_shift]])
-    probe = probe.transpose(-1, -3).transpose(-2, -4)
-    polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
+    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)
 
-    # Transpose it back
-    return polarized_probe.transpose(-1, -3).transpose(-2, -4)
+    return polarized
 
-def apply_circular_polarizer(probe, left_polarized=True):
+def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
     """
     Applies a circular polarizer to the probe
 
     Parameters:
     ----------
     probe: t.Tensor
-        A (...)x2x1xMxL tensor representing the probe
+        A (...)x2xMxL tensor representing the probe
     left_polarizd: bool
         True for the left-polarization, False for the right
-    
+
     Returns:
     --------
     circularly polarized probe: t.Tensor
-        (...)x2x1xMxL 
+        (...)x2xMxL
     """
     probe = probe.to(dtype=t.cfloat)
     if left_polarized:
-        jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]]))
+        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]])
-    probe = probe.transpose(-1, -3).transpose(-2, -4)
-    polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
+        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
 
-    # Transpose it back
-    return polarized_probe.transpose(-1, -3).transpose(-2, -4)
-
-def apply_quarter_wave_plate(probe, fast_axis_angle):
+def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
     """
     Parameters:
     ----------
@@ -184,19 +164,17 @@ def apply_quarter_wave_plate(probe, fast_axis_angle):
     Returns:
     --------
     polarized probe: t.Tensor
-        (...)x2x1xMxL 
+        (...)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]])
-    probe = probe.transpose(-1, -3).transpose(-2, -4)
-    polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
-    # Transpose it back
-    return polarized_probe.transpose(-1, -3).transpose(-2, -4)
+    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):
+def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
     """
     Parameters:
     ----------
@@ -208,37 +186,32 @@ def apply_half_wave_plate(probe, fast_axis_angle):
     Returns:
     --------
     polarized probe: t.Tensor
-        (...)x2x1xMxL 
+        (...)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]])
-    probe = probe.transpose(-1, -3).transpose(-2, -4)
-    polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
-    # Transpose it back
-    return polarized_probe.transpose(-1, -3).transpose(-2, -4)
+    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)
 
-# probe = t.rand(17, 7, 2, 6, 4)
-# polarizer = t.rand(7)
-# out = apply_linear_polarizer(probe, polarizer)
-# out2 = apply_linear_polarizer(probe, polarizer, transpose=False)
+    return out
 
-# print(out.shape)
-# print(out2.shape)
+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
 
-# a = t.ones(17, 8, 2, 3, 4)
-# b = t.ones(2, 1, 1)
+    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)
 
-
-# probe = t.ones(5, 2, 3, 3)
-
-# polarizer = t.tensor([45])
-
-# exitw = apply_linear_polarizer(probe, polarizer)
-# print(exitw[:, 0, :, :])
-# print('y', exitw[:, 1, :, :])
-
-#a = t.ones(2, 4)
-#print(t.sum(a, dim=1).shape)
+    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))
diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py
index fc124b7..bca164e 100644
--- a/CDTools/tools/propagators/propagators.py
+++ b/CDTools/tools/propagators/propagators.py
@@ -382,6 +382,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
     # 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)
diff --git a/temp_tests/simulated_dataset.py b/temp_tests/simulated_dataset.py
new file mode 100644
index 0000000..817996b
--- /dev/null
+++ b/temp_tests/simulated_dataset.py
@@ -0,0 +1,132 @@
+import numpy as np
+import torch as t
+from scipy import misc
+from CDTools.models import PolarizedFancyPtycho
+from CDTools.datasets import PolarizedPtycho2DDataset
+import CDTools
+from CDTools.tools import polarization
+from CDTools import tools
+from matplotlib import pyplot as plt
+from PIL import Image
+
+# upolad 4 different images representing 4 components of the object
+# and 2 gaaussian functionas corresponding to the probe components
+
+f = misc.ascent()
+x , y = np.shape(f)
+aa = f[:x//2, :y//2]
+bb = f[:x//2, -y//2:]
+cc = f[-x//2:, :y//2]
+dd = f[-x//2:, -y//2:]
+
+print(1)
+
+def simulate_polarized_dataset(probe_size, obj_size, num_patt, a, b, c, d):
+    a, b, c, d = t.as_tensor(a, dtype=t.cfloat), t.as_tensor(b, dtype=t.cfloat), t.as_tensor(c, dtype=t.cfloat), t.as_tensor(d, dtype=t.cfloat)
+    # a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
+    # b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+    # c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+    # d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+
+    translations = []
+    xs, ys = np.mgrid[:num_patt, :num_patt]
+    xs, ys = np.ravel(xs), np.ravel(ys)
+    for x, y in zip(xs, ys):
+        for j in range(9):
+            translations.append((5*x, 5*y))
+    translations = t.as_tensor(translations, dtype=t.float32)
+    obj = t.stack((t.stack((a, b), dim=-3), t.stack((c, d), dim=-3)), dim=-4)
+
+    # for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1]):
+    #     plt.imshow(np.real(obj[i, j, ...]))
+    #     plt.show()
+    probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), np.array((2, 2)))
+    probe = t.stack((probe, probe), dim=-3)
+    probe = polarization.apply_circular_polarizer(probe, multiple_modes=False)
+    polarizers = polarization.generate_linear_polarizer(t.tensor([0, 45, 90]))
+    num_transl = len(translations)
+    pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i], multiple_modes=False) for i in range(3)]
+    analyzers = t.stack(([polarizers[i % 3] for i in range(num_transl)]), dim=0)
+    analyzer = t.tensor([(i % 3) * 45 for i in range(num_transl)])
+    polarizer = t.tensor([(i // 3) % 3 * 45 for i in range(num_transl)])
+    polarized_probes = t.stack(([pol_probes[(i // 3) % 3] for i in range(num_transl)]), dim=0)
+    polarized_wavefields = t.stack([tools.interactions.ptycho_2D_sinc(polarized_probes[i], obj, translations[i], polarized=True, multiple_modes=False) for i in range(num_transl)])
+    wavefields = tools.propagators.far_field(polarized_wavefields)
+    wavefields = polarization.apply_jones_matrix(wavefields, analyzers, multiple_modes=False)
+    patterns = t.abs(wavefields[:, 0, :, :])**2 + t.abs(wavefields[:, 1, :, :])**2
+    detector_basis = t.transpose(t.tensor([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]), 0, 1)
+    det_shape = t.Size((obj_size, obj_size))
+    wavelength = 532e-9
+    real_basis = tools.initializers.exit_wave_geometry(detector_basis, det_shape, wavelength, 2.5e-2)[0]
+    # print('f', type(real_basis), translations.shape)
+    real_translations = tools.interactions.pixel_to_translations(real_basis, translations)
+    # print(type(real_translations))
+    patterns = patterns.numpy()
+    detector_geometry = {
+                         'corner': np.array([probe_size*4.8e-6/2, probe_size*4.8e-6/2, 2.5e-2]),
+                         'basis': np.array([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]).transpose(),
+                         'distance': 2.5e-2
+                        }
+
+    return PolarizedPtycho2DDataset(real_translations, polarizer, analyzer, patterns,
+                                    axes=("x", "y"), detector_geometry=detector_geometry, wavelength=wavelength)
+
+# print(1)
+dataset = simulate_polarized_dataset(100, 500, 10, aa, bb, cc, dd)
+# print(2)
+dataset.inspect()
+# print(3)
+plt.show()
+model = PolarizedFancyPtycho.from_dataset(dataset, propagation_distance=1e-3)
+model.inspect()
+# plt.show()
+model.compare(dataset)
+plt.show()
+
+# for loss in model.Adam_optimize(400, dataset, batch_size=5, lr=0.002, schedule=True):
+#     # And we liveplot the updates to the model as they happen
+#     print(model.report())
+#     model.inspect(dataset)
+#     model.save_figures(prefix='simulated', extension='png')
+#     res = model.save_results(dataset)
+#     np.save('simulated_dataset.npy', res)
+
+dataset = simulate_polarized_dataset(50, 100, 10, aa, bb, cc, dd)
+
+res = np.load('simulated_dataset.npy', allow_pickle=True)
+res = res[()]
+print(type(res))
+Ws = t.ones(len(dataset))
+ewg = CDTools.tools.initializers.exit_wave_geometry
+probe_basis, probe_shape, det_slice = ewg(res['basis'],
+                                          dataset[0][1].shape,
+                                          dataset.wavelength,
+                                          dataset.detector_geometry['distance'],
+                                          center=None,
+                                          padding=0)
+print('object', res['obj'].shape)
+obj = res['obj'][..., 200:275, 200:275]
+# print('obj', obj)
+a = obj[0, 0, :, :]
+b = obj[0, 1, :, :]
+c = obj[1, 0, :, :]
+d = obj[1, 1, :, :]
+for i in [a, b, c, d]:
+    plt.imshow(np.real(i))
+    # plt.imshow(np.log(np))
+    plt.colorbar()
+    plt.show()
+print(np.allclose(np.real(a), np.real(b)))
+print('a', a[..., 10:20, 10:20])
+print('b', b[..., 10:20, 10:20])
+models = [CDTools.models.FancyPtycho(dataset.wavelength, dataset.detector_geometry, probe_basis,
+                                     res['probe'], component, surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
+                                     min_translation=t.tensor([0, 0], dtype=t.float32),
+                                     background=t.tensor(res['background']), translation_offsets=None, mask=None,
+                                     weights=Ws, translation_scale=1, saturation=None,
+                                     probe_support=None, oversampling=1,
+                                     loss='amplitude mse', units='um') for component in [a, b, c, d]]
+# print('model is created')
+# for model in models:
+#     model.inspect()
+#     plt.show()
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index b27bca5..86597cd 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -1,5 +1,6 @@
 import numpy as np
 from scipy import fftpack as ffts
+from scipy import linalg as la
 import torch as t
 from itertools import combinations
 
@@ -239,7 +240,6 @@ def test_calc_deconvolved_cross_correlation():
     assert np.allclose(test_cor_t.numpy(), np_cor)
 
     
-
 def test_calc_frc():
 
     obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
@@ -300,3 +300,141 @@ def test_calc_frc():
     assert np.allclose(bins, test_bins_t.numpy())
     assert np.allclose(frc, test_frc_t.numpy())
     assert np.allclose(threshold, test_threshold_t.numpy())
+
+    
+def test_calc_rms_error():
+    field_1 = t.rand(14,19, dtype=t.complex64)
+    field_2 = t.rand(14,19, dtype=t.complex64)
+
+    # Check that the calculation is insensitive to phase
+    assert t.allclose(analysis.calc_rms_error(field_1, field_2),
+                      analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
+
+    # And that it is sensitive to phase if we turn off the 
+    assert not t.allclose(
+        analysis.calc_rms_error(field_1, field_2, align_phases=False),
+        analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
+                                align_phases=False))
+
+    # Check that the result is positive
+    assert analysis.calc_rms_error(field_1, field_2) > 0
+
+    # And that it is a smaller number with align_phases on
+    assert (analysis.calc_rms_error(field_1, field_2) <=
+            analysis.calc_rms_error(field_1, field_2, align_phases=False))
+
+    # Now we check against an explicit implementation:
+    gamma = field_1 * t.conj(field_2)
+    gamma /= t.abs(gamma)
+
+    # This is an alternate way of doing the calculation. Actually, would this
+    # be a better implementation anyway? Probably no difference tbh.
+    rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
+                                t.mean(t.abs(field_2)**2) -
+                                2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
+    assert t.allclose(rms_error_nophase,
+                      analysis.calc_rms_error(field_1, field_2))
+
+    rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
+                              t.mean(t.abs(field_2)**2) -
+                              2 * t.real(t.mean(field_1 * t.conj(field_2)))))
+
+    assert t.allclose(rms_error_phase,
+                      analysis.calc_rms_error(field_1, field_2,
+                                              align_phases=False))
+
+    # Now let's test that it works along a dimension:
+    
+    field_1 = t.rand(3,14,19, dtype=t.complex64)
+    field_2 = t.rand(3,14,19, dtype=t.complex64)
+    result = analysis.calc_rms_error(field_1, field_2, normalize=True)
+    assert (result.shape == t.Size([3]))
+    
+    for i in range(3):
+        assert t.allclose(analysis.calc_rms_error(field_1[i],
+                                                  field_2[i],
+                                                  normalize=True),
+                          result[i])
+    
+
+def test_calc_fidelity():
+
+    fields_1 = t.rand(2,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+
+    dm_1 = t.reshape(fields_1, (2,-1))
+    dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
+    dm_2 = t.reshape(fields_2, (3,-1))
+    dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
+    
+    inner_mat = la.sqrtm(np.dot(np.dot(la.sqrtm(dm_1),dm_2),la.sqrtm(dm_1)))
+    fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
+    
+    assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
+
+    fields_1 = t.rand(1,30,17, dtype=t.complex128)
+    fields_2 = t.rand(1,30,17, dtype=t.complex128)
+
+    assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
+                     analysis.calc_fidelity(fields_1, fields_2))
+
+    # Checking that it works with extra dimensions
+    fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
+    field_3 = t.rand(1,30,17, dtype=t.complex128)
+
+    fidelities = analysis.calc_fidelity(fields_1, fields_2)
+    fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
+    for i in range(3):
+        assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
+                         fidelities[i])
+        assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
+                         fidelities_2[i])
+
+    # Check that the diensionality argument works
+    fields_1 = t.rand(3,2,12, dtype=t.complex128)
+    fields_2 = t.rand(3,2,12, dtype=t.complex128)
+
+    assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
+            == t.Size([3]))
+    
+    fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
+    fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
+
+    assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
+            == t.Size([3]))
+
+def test_calc_generalized_rms_error():
+
+    # Test that it matches the rms error for coherent fields
+    
+    fields_1 = t.rand(1,30,17, dtype=t.complex128)
+    fields_2 = t.rand(1,30,17, dtype=t.complex128)
+    
+    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
+                     analysis.calc_rms_error(fields_1[0], fields_2[0],
+                                             align_phases=True))
+
+    # Test that it is independent of field order
+    fields_1 = t.rand(5,30,17, dtype=t.complex128)
+    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+    fields_3 = fields_2.flip(0)
+    
+    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
+                     analysis.calc_generalized_rms_error(fields_1, fields_3))
+
+    # Test with leading dimensions
+    fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
+    fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
+    
+    assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
+            == t.Size([3,4]))
+
+    # And test with different number of dimensions dims
+    # Test that it is independent of field order
+    fields_1 = t.rand(3,6,17, dtype=t.complex128)
+    fields_2 = t.rand(3,1,17, dtype=t.complex128)
+    fields_3 = fields_2.flip(0)
+
+    assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
+
diff --git a/tests/tools/test_polarization.py b/tests/tools/test_polarization.py
index 04cbbfa..74b8f94 100644
--- a/tests/tools/test_polarization.py
+++ b/tests/tools/test_polarization.py
@@ -1,15 +1,17 @@
 import numpy as np
 import torch as t
+from CDTools.tools.polarization import apply_linear_polarizer, generate_linear_polarizer
 from CDTools.tools.polarization import apply_jones_matrix as jones
 
+
 # Abe - I removed all the imports that didn't need to be here.
 
 # Abe - A few issues. First, you could just write "from math import cos, sin"
 # Second, we already have numpy imported, so better to use np.cos and np.sin
-from math import cos as cos, sin as sin
+# from math import cos as cos, sin as sin
 # numpy also has np.pi and np.deg2rad 
-import math
 
+from numpy import cos, sin, deg2rad
 
 # Further comments:
 #
@@ -34,408 +36,589 @@ angle = 87
 angle_2 = angle - 45
 
 def polarizer(angle):
-	theta = math.radians(angle)
-	polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
-	return polarizer
+    theta = deg2rad(angle)
+    polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
+    return polarizer
 
-exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
-theta2 = math.radians(angle_2)
+exponent = t.exp(-1j * np.pi / 4 * t.ones(2, 2)).to(dtype=t.cfloat)
+theta2 = deg2rad(angle_2)
 quarter_plate = t.tensor([[(cos(theta2))**2 + 1j * (sin(theta2))**2, (1 - 1j) * sin(theta2) * cos(theta2)], 
-								   [(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]])
+                                   [(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]]).to(dtype=t.cfloat)
 
 def build_from_quarters(jones1, jones2, jones3, jones4):
-	x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
-	y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
-	x = t.stack((x, x), dim=-2)
-	y = t.stack((y, y), dim=-2)
-	return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
+    x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
+    y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
+    x = t.stack((x, x), dim=-2)
+    y = t.stack((y, y), dim=-2)
+    return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
 
 jones_plate = t.matmul(quarter_plate, polarizer(angle))
 jones0 = polarizer(0)
 jones90 = polarizer(90)
 jones45 = polarizer(45)
 
-
-'''	
-after applying the polarizer and the quarter_plate, the probe should get circularly polarized
-probe: no multiple modes, 1 diffr pattern
-	2xMxL
-jones_matrix: same jones matrix applied to all the pixels
-	2x2
-'''
 transpose = True
-def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():	
-	probe = t.rand(2, 3, 4, dtype=t.cfloat)
-	print(polarizer)
-	out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose), 
-		quarter_plate, multiple_modes=False, transpose=transpose)
-	print('expected shape:(2, 3, 4)')
-	print('actual:', out.shape)
-	print('simulated:', out)
 
-	assert out.shape == t.Size((2,3,4))
+def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr(): 
+    ''' 
+    after applying the polarizer and the quarter_plate, the probe should get circularly polarized
+    probe: no multiple modes, 1 diffr pattern
+        2xMxL
+    jones_matrix: same jones matrix applied to all the pixels
+        2x2
+    '''
+    probe = t.rand(2, 3, 4, dtype=t.cfloat)
+    print(polarizer)
+    out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose), 
+        quarter_plate, multiple_modes=False, transpose=transpose)
+    print('expected shape:(2, 3, 4)')
+    print('actual:', out.shape)
+    print('simulated:', out)
 
-	assert np.allclose(np.real(out[0]), np.imag(out[1]))
+    assert np.allclose(np.real(out[0]), np.imag(out[1]))
+    assert out.shape == t.Size((2, 3, 4))
 
 
-'''probe: no multiple modes, 1 diffr pattern
-	2xMxL = 2x4x4
-jones_matrix: jones matrices differ from pixel to pixel
-	2x2xMxL = 2x2x4x4
-4 quarters: 
-	1: [:, :, :-2, :-2] - circular_polarizer, 
-	2: [:, :, :-2, -2:] - 0
-	3: [:, :, -2:, :-2] - 90
-	4: [:, :, -2:, -2:] - 45
-'''
+def test_generate_linear_polarizer():
+    pol_angles = [0, 45, 90]
+    pol_angle1 = 45
+    pol_angle2 = t.tensor(90)
+    pol_angle3 = t.tensor([0])
+    pols = generate_linear_polarizer(pol_angles)
+    pol1 = generate_linear_polarizer(pol_angle1)
+    pol2 = generate_linear_polarizer(pol_angle2)
+    pol3 = generate_linear_polarizer(pol_angle3)
+    print('polarizers 0, 45, 90 (1D tensor) shape:', pols.shape)
+    print('shape of the polarizer generated from int:', pol1.shape)
+    print('shape of the polarizer generated from 0D tensor:', pol2.shape)
+    print('shape of the linear polarizer generated from t.Size(0) tensor:', pol3.shape)
+    print('90', pol2)
+    print(jones90)
+    probe = t.ones(4, 4)
+    jones_m = [jones0, jones45, jones90]
+    jones_m = t.stack([matr for matr in jones_m])
+    assert pols.shape == t.Size((3, 2, 2))
+    assert pol1.shape == t.Size((2, 2))
+    assert pol2.shape == t.Size((2, 2))
+    assert pol3.shape == t.Size((1, 2, 2))
+    assert t.allclose(pol1, jones45)
 
-def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():	
-	jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
-	probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
-	print('jones:', jones_matr)
-	# print('jones:', jones_matr)
-	out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
 
-	# jones -> (2,2,x,y), probe, output probe
-	# interaction -> 
-	print('expected shape:(2, 4, 4)')
-	print('simulated:', out.shape)
-	print('simulated:', out)
+def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():    
+    '''probe: no multiple modes, 1 diffr pattern
+        2xMxL = 2x4x4
+    jones_matrix: jones matrices differ from pixel to pixel
+        2x2xMxL = 2x2x4x4
+    4 quarters: 
+        1: [:, :, :-2, :-2] - circular_polarizer, 
+        2: [:, :, :-2, -2:] - 0
+        3: [:, :, -2:, :-2] - 90
+        4: [:, :, -2:, -2:] - 45
+    '''
+    jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
+    probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
+    print('jones:', jones_matr)
+    # print('jones:', jones_matr)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
 
-	# Example of using multiple asserts
-	assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
-	assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
-	assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-	assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-	assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
-	assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])
-	
+    # jones -> (2,2,x,y), probe, output probe
+    # interaction -> 
+    print('expected shape:(2, 4, 4)')
+    print('simulated:', out.shape)
+    print('simulated:', out)
+
+    # Example of using multiple asserts
+    assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
+    assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])
+    assert out.shape == t.Size((2, 4, 4))
 
-'''
-probe: no multiple modes, multiple diffr patterns
-	Nx2xMxL = 3x2x4x4
-jones_matrix: same jones matrix applied to all the pixels
-	Nx2x2 = 3x2x2
-	3 different matrices for each probe:
-		1: 0
-		2: 45
-		3: 90
-'''
 def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr():
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
-	print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+    '''
+    probe: no multiple modes, multiple diffr patterns
+        Nx2xMxL = 3x2x4x4
+    jones_matrix: same jones matrix applied to all the pixels
+        Nx2x2 = 3x2x2
+        3 different matrices for each probe:
+            1: 0
+            2: 45
+            3: 90
+    '''
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
+    print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
 
-	print('expected shape: (3, 2, 4, 4)')
-	print('actual:', out.shape)
-	print('simulated:', out)
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual:', out.shape)
+    print('simulated:', out)
 
-	assert (t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
-			and t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
-			and t.allclose(out[1, 0, :, :], out[1, 1])
-			and t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
-			and t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat)))
+    assert t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
+    assert t.allclose(out[1, 0, :, :], out[1, 1])
+    assert t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat))
+    assert out.shape == t.Size((3, 2, 4, 4))
 
-
-'''
-probe: no multiple modes, multiple diffr pattern
-	Nx2xMxL = 3x2x4x4 
-jones_matrix: jones matrices differ from pixel to pixel
-	Nx2x2xMxL = 3x2x2x4x4
-	jones matrix for the 1st pattern:
-		1: quat plate, 2: 90, 3: 0, 4: 45
-	jones matrix for the 2nd pattern:
-		1: 0, 2: 45, 3: 90, 4: quat plate
-	jones matrix for the 3rd pattern:
-	 	1: 90, 2: 0, 3: 45, 4: quat plate
-'''
 def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr():
-	jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
-								 build_from_quarters(jones0, jones45, jones90, jones_plate),
-								 build_from_quarters(jones90, jones0, jones45, jones_plate)]
-	jones_matr = t.stack(([i for i in jones_m]), dim=0)
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+    '''
+    probe: no multiple modes, multiple diffr pattern
+        Nx2xMxL = 3x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        Nx2x2xMxL = 3x2x2x4x4
+        jones matrix for the 1st pattern:
+            1: quat plate, 2: 90, 3: 0, 4: 45
+        jones matrix for the 2nd pattern:
+            1: 0, 2: 45, 3: 90, 4: quat plate
+        jones matrix for the 3rd pattern:
+            1: 90, 2: 0, 3: 45, 4: quat plate
+    '''
+    jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
+                                 build_from_quarters(jones0, jones45, jones90, jones_plate),
+                                 build_from_quarters(jones90, jones0, jones45, jones_plate)]
+    jones_matr = t.stack(([i for i in jones_m]), dim=0)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
 
-	print('expected shape: (3, 2, 4, 4)')
-	print('actual shape:', out.shape)
-	print('simulated:', out)
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated:', out)
 
 
-	assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
-			and t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
+    assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
+    assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
 
-			and t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
-			and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
-			and np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
+    assert t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
+    assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
+    assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
 
-			and t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
-			and np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:])))
+    assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
+    assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
+    assert out.shape == t.Size((3, 2, 4, 4))
 
-
-'''
-probe: multiple modes, 1 diffr pattern
-	Px2xMxL = 2x2x3x4
-jones_matrix: same jones matrix applied to all the pixels
-	2x2 - quarter plate
-'''
 def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr():
-	probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
-	out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, 1 diffr pattern
+        Px2xMxL = 2x2x3x4
+    jones_matrix: same jones matrix applied to all the pixels
+        2x2 - quarter plate
+    '''
+    probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
+    out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (2, 2, 3, 4)')
-	print('actual:', out.shape)
-	print('simulated:', out)
+    print('expected shape: (2, 2, 3, 4)')
+    print('actual:', out.shape)
+    print('simulated:', out)
 
-	assert (np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
-			and np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :])))
+    assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
+    assert np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :]))
+    assert out.shape == t.Size((2, 2, 3, 4))
 
-
-'''
-probe: multiple modes, 1 diffr pattern
-	Px2xMxL = 3x2x4x4 
-jones_matrix: jones matrices differ from pixel to pixel
-	2xMxL = 2x4x4
-4 quarters: 
-	1: [:, :, :-2, :-2] - circular_polarizer, 
-	2: [:, :, :-2, -2:] - 0
-	3: [:, :, -2:, :-2] - 90
-	4: [:, :, -2:, -2:] - 45
-'''
 def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr():
-	jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, 1 diffr pattern
+        Px2xMxL = 3x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        2xMxL = 2x4x4
+    4 quarters: 
+        1: [:, :, :-2, :-2] - circular_polarizer, 
+        2: [:, :, :-2, -2:] - 0
+        3: [:, :, -2:, :-2] - 90
+        4: [:, :, -2:, -2:] - 45
+    '''
+    jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (3, 2, 4, 4)')
-	print('actual shape:', out.shape)
-	print('simulated:', out)
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated:', out)
 
-	assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
-		and t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
+    assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
+    assert t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
 
-		and np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
-		and t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
+    assert np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
+    assert t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
 
-		and np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
-		and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
+    assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
+    assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
+    assert out.shape == t.Size((3, 2, 4, 4))
 
-
-'''
-probe: multiple modes, multiple diffr patterns
-	NxPx2xMxL = 3x7x2x3x4
-jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
-	Nx2x2 = 3x2x2
-		3 different matrices for each probe in one mode:
-		1: 0
-		2: 45
-		3: 90
-'''
 def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr():
-	probe = t.ones(2, 3, 4, dtype=t.cfloat)
-	# 1st mode
-	probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	# 2nd mode
-	probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)	
-	probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
-	jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
-	print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, multiple diffr patterns
+        NxPx2xMxL = 3x7x2x3x4
+    jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
+        Nx2x2 = 3x2x2
+            3 different matrices for each probe in one mode:
+            1: 0
+            2: 45
+            3: 90
+    '''
+    probe = t.ones(2, 3, 4, dtype=t.cfloat)
+    # 1st mode
+    probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    # 2nd mode
+    probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)    
+    probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
+    jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
+    print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (3, 7, 2, 3, 4)')
-	print('actual:', out.shape)
-	print('simulated (patterns in one mode):', out[:, 6, :, :, :])
-	
-	# we'll be checking only one mode
-	assert (t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
-			and t.allclose(out[0, 6,  1, :, :], t.zeros(3, 4, dtype=t.cfloat))
-			and t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
-			and t.allclose(out[2, 6, 0, :, :], 3 * (10**6) *  t.zeros(3, 4, dtype=t.cfloat))
-			and t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat)))
+    print('expected shape: (3, 7, 2, 3, 4)')
+    print('actual:', out.shape)
+    print('simulated (patterns in one mode):', out[:, 6, :, :, :])
+    
+    # we'll be checking only one mode
+    assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[0, 6,  1, :, :], t.zeros(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
+    assert t.allclose(out[2, 6, 0, :, :], 3 * (10**6) *  t.zeros(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat))
+    assert out.shape == t.Size((3, 7, 2, 3, 4))
 
-
-'''
-probe: multiple modes, multiple diffr pattern
-	NxPx2xMxL = 3x4x2x4x4 
-jones_matrix: jones matrices differ from pixel to pixel
-	Nx2x2xMxL = 3x2x2x4x4
-	(differs across the patterns in each mode)
-	jones matrix for the 1st pattern:
-		1: quat plate, 2: 90, 3: 0, 4: 45
-	jones matrix for the 2nd pattern:
-		1: 0, 2: 45, 3: 90, 4: quat plate
-	jones matrix for the 3rd pattern:
-	 	1: 90, 2: 0, 3: 45, 4: quat plate
-'''
 def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr():
-	jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
-						 build_from_quarters(jones0, jones45, jones90, jones_plate),
-						 build_from_quarters(jones90, jones0, jones45, jones_plate)]
-	jones_matr = t.stack(([i for i in jones_m]), dim=0)
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	# one mode
-	probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)		
-	probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
-	print('probe:', probe.shape)
-	print('jones:', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, multiple diffr pattern
+        NxPx2xMxL = 3x4x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        Nx2x2xMxL = 3x2x2x4x4
+        (differs across the patterns in each mode)
+        jones matrix for the 1st pattern:
+            1: quat plate, 2: 90, 3: 0, 4: 45
+        jones matrix for the 2nd pattern:
+            1: 0, 2: 45, 3: 90, 4: quat plate
+        jones matrix for the 3rd pattern:
+            1: 90, 2: 0, 3: 45, 4: quat plate
+    '''
+    jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
+                         build_from_quarters(jones0, jones45, jones90, jones_plate),
+                         build_from_quarters(jones90, jones0, jones45, jones_plate)]
+    jones_matr = t.stack(([i for i in jones_m]), dim=0)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    # one mode
+    probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)      
+    probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
+    print('probe:', probe.shape)
+    print('jones:', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (3, 2, 2, 4, 4)')
-	print('actual shape:', out.shape)
-	print('simulated patterns in one mode:', out[:, 3, :, :, :])
+    print('expected shape: (3, 4, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated patterns in one mode:', out[:, 3, :, :, :])
 
-	o = 10 ** 3
-	# we'll be checking only one mode (4th)
-	assert (np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
-		and t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
+    o = 10 ** 3
+    # we'll be checking only one mode (4th)
+    assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
+    assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
 
-		and t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
-		and t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
-		and np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
+    assert t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
+    assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
+    assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
 
-		and t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
-		and np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:])))
+    assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
+    assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
+    assert out.shape == t.Size((3, 4, 2, 4, 4))
 
-
-'''
-probe: no multiple modes, multiple diffr patterns
-	Nx2xMxL = 3x2x4x4
-jones_matrix: same jones matrix applied to all the pixels
-	2x2 = 2x2 - a quarter waveplate 
-'''
 def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1():
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	jones_matr = jones_plate
-	print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+    '''
+    probe: no multiple modes, multiple diffr patterns
+        Nx2xMxL = 3x2x4x4
+    jones_matrix: same jones matrix applied to all the pixels
+        2x2 = 2x2 - a quarter waveplate 
+    '''
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    jones_matr = jones_plate
+    print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
 
-	print('expected shape: (3, 2, 4, 4)')
-	print('actual:', out.shape)
-	print('simulated:', out)
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual:', out.shape)
+    print('simulated:', out)
 
-	assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
+    assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
+    assert out.shape == t.Size((3, 2, 4, 4))
 
-
-'''
-probe: no multiple modes, multiple diffr pattern
-	Nx2xMxL = 3x2x4x4 
-jones_matrix: jones matrices differ from pixel to pixel
-	2x2xMxL = 2x2x4x4
-	1: quat plate, 2: 90, 3: 0, 4: 45
-
-'''
 def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_1():
-	jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+    '''
+    probe: no multiple modes, multiple diffr pattern
+        Nx2xMxL = 3x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        2x2xMxL = 2x2x4x4
+        1: quat plate, 2: 90, 3: 0, 4: 45
 
-	print('expected shape: (3, 2, 4, 4)')
-	print('actual shape:', out.shape)
-	print('simulated:', out)
+    '''
+    jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated:', out)
 
 # check the 3rd pattern
-	assert (np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
-			and t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-			and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
+    assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
+    assert t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
+    assert out.shape == t.Size((3, 2, 4, 4))
 
-
-'''
-probe: multiple modes, multiple diffr patterns
-	NxPx2xMxL = 3x7x2x3x4
-jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
-	2x2 = 2x2
-	quarter wave plate
-'''
 def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1():
-	probe = t.ones(2, 3, 4, dtype=t.cfloat)
-	# 1st mode
-	probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
-	probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
-	jones_matr = jones_plate
-	print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, multiple diffr patterns
+        NxPx2xMxL = 3x7x2x3x4
+    jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
+        2x2 = 2x2
+        quarter wave plate
+    '''
+    probe = t.ones(2, 3, 4, dtype=t.cfloat)
+    # 1st mode
+    probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
+    probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
+    jones_matr = jones_plate
+    print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (3, 7, 2, 3, 4)')
-	print('actual:', out.shape)
-	print('simulated (patterns in one mode):', out[:, 6, :, :, :])
-	
-	# we'll be checking only one mode
-	assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
+    print('expected shape: (3, 7, 2, 3, 4)')
+    print('actual:', out.shape)
+    print('simulated (patterns in one mode):', out[:, 6, :, :, :])
+    
+    # we'll be checking only one mode
+    assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
+    assert out.shape == t.Size((3, 7, 2, 3, 4))
 
-
-'''
-probe: multiple modes, multiple diffr pattern
-	NxPx2xMxL = 3x4x2x4x4 
-jones_matrix: jones matrices differ from pixel to pixel
-	2x2xMxL = 2x2x4x4
-	jones matrix:
-	1: quat plate, 2: 90, 3: 0, 4: 45
-'''
 def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_1():
-	jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
-	probe = t.ones(2, 4, 4, dtype=t.cfloat)
-	# one mode
-	probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)		
-	probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
-	print('probe:', probe.shape)
-	print('jones:', jones_matr.shape)
-	out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+    '''
+    probe: multiple modes, multiple diffr pattern
+        NxPx2xMxL = 3x4x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        2x2xMxL = 2x2x4x4
+        jones matrix:
+        1: quat plate, 2: 90, 3: 0, 4: 45
+    '''
+    jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    # one mode
+    probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)      
+    probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
+    print('probe:', probe.shape)
+    print('jones:', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
 
-	print('expected shape: (3, 2, 2, 4, 4)')
-	print('actual shape:', out.shape)
-	print('simulated patterns in one mode:', out[:, 3, :, :, :])
+    print('expected shape: (3, 4, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated patterns in one mode:', out[:, 3, :, :, :])
 
-	o = 10 ** 2
-	# we'll be checking only one mode (3th)
-	assert (np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
-		and t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
-		and t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:]))
+    o = 10 ** 2
+    # we'll be checking only one mode (3th)
+    assert np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
+    assert t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:])
+    assert out.shape == t.Size((3, 4, 2, 4, 4))
 
-	return None
+    return None
 
+
+def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_1():
+    '''
+    probe: 
+        2xMxL = 2x3x4 
+    jones_matrix: 
+        Nx2x2 = 3x2x2
+        jones matrices:
+        1: quat plate, 2: 90, 3: 0
+    '''
+    jones_matr = t.stack((jones_plate, jones90, jones0))
+    probe = t.ones(2, 3, 4, dtype=t.cfloat)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+
+    print('expected shape: (3, 2, 3, 4)')
+    print('actual shape:', out.shape)
+    print('simulated:', out)
+
+    # we'll be checking only one mode (3th)
+    assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
+    assert t.allclose(out[1, 0, :, :], t.zeros(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, :, :], t.ones(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, :, :], t.ones(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
+    assert out.shape == t.Size((3, 2, 3, 4))
+
+    return None
+
+
+def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_2():
+    '''
+    probe: 
+        2xMxL = 2x4x4
+    jones_matrix: jones matrices differ from pixel to pixel
+        Nx2x2xMxL = 3x2x2x4x4
+        jones matrix for the 1st pattern:
+            1: quat plate, 2: 90, 3: 0, 4: 45
+        jones matrix for the 2nd pattern:
+            1: 0, 2: 45, 3: 90, 4: quat plate
+        jones matrix for the 3rd pattern:
+            1: 90, 2: 0, 3: 45, 4: quat plate
+    '''
+    jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
+                         build_from_quarters(jones0, jones45, jones90, jones_plate),
+                         build_from_quarters(jones90, jones0, jones45, jones_plate)]
+    jones_matr = t.stack(([i for i in jones_m]), dim=0)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)
+    out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
+
+    print('expected shape: (3, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated:', out)
+
+    assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
+    assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
+
+    assert t.allclose(out[1, 0, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
+    assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
+
+    assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
+    assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
+    assert out.shape == t.Size((3, 2, 4, 4))
+
+
+def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_1():
+    '''
+    probe: multiple modes, multiple diffr patterns
+        Px2xMxL = 7x2x3x4
+    jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
+        Nx2x2 = 3x2x2
+            3 different matrices for each pattern:
+            1: 0
+            2: 45
+            3: 90
+    '''
+    probe = t.ones(2, 3, 4, dtype=t.cfloat)  
+    probe = t.stack(([probe * 10 ** i for i in range(7)]), dim=0)
+    jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
+    print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+
+    print('expected shape: (3, 7, 2, 3, 4)')
+    print('actual:', out.shape)
+    print('simulated (patterns in one mode):', out[:, 6, :, :, :])
+    
+    # we'll be checking only one mode
+    assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[0, 6,  1, :, :], t.zeros(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
+    assert t.allclose(out[2, 6, 0, :, :], (10**6) *  t.zeros(3, 4, dtype=t.cfloat))
+    assert t.allclose(out[2, 6, 1, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
+    assert out.shape == t.Size((3, 7, 2, 3, 4))
+
+
+def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_2():
+    '''
+    probe: multiple modes, multiple diffr pattern
+        Px2xMxL = 4x2x4x4 
+    jones_matrix: jones matrices differ from pixel to pixel
+        Nx2x2xMxL = 3x2x2x4x4
+        (differs across the patterns in each mode)
+        jones matrix for the 1st pattern:
+            1: quat plate, 2: 90, 3: 0, 4: 45
+        jones matrix for the 2nd pattern:
+            1: 0, 2: 45, 3: 90, 4: quat plate
+        jones matrix for the 3rd pattern:
+            1: 90, 2: 0, 3: 45, 4: quat plate
+    '''
+    jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45), 
+                         build_from_quarters(jones0, jones45, jones90, jones_plate),
+                         build_from_quarters(jones90, jones0, jones45, jones_plate)]
+    jones_matr = t.stack(([i for i in jones_m]), dim=0)
+    probe = t.ones(2, 4, 4, dtype=t.cfloat)    
+    probe = t.stack(([probe * (10 ** i) for i in range(4)]), dim=0)
+    print('probe:', probe.shape)
+    print('jones:', jones_matr.shape)
+    out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
+
+    print('expected shape: (3, 4, 2, 4, 4)')
+    print('actual shape:', out.shape)
+    print('simulated patterns in one mode:', out[:, 3, :, :, :])
+
+    o = 10 ** 3
+    # we'll be checking only one mode (4th)
+    assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
+    assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
+
+    assert t.allclose(out[1, 3, 0, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
+    assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[1, 3, 1, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
+
+    assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 1, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 0, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
+    assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
+    assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
+    assert out.shape == t.Size((3, 4, 2, 4, 4))
\ No newline at end of file