Maybe last time it didn't actually commit?

CDTools/datasets/ptycho_2d_dataset.py

@@ -337,7 +337,11 @@ class Ptycho2DDataset(CDataset):
                 idx = slider.val - 1
             elif event.key == 'down' or event.button == 'down' or event.key == 'left':
                 idx = slider.val + 1
-
+            else:
+                # This prevents errors from being thrown on irrelevant key
+                # or mouse input
+                return
+            
             # Handle the wraparound and trigger the update
             idx = int(idx) % len(self)
             slider.set_val(idx)

CDTools/models/__init__.py

@@ -24,7 +24,7 @@ from __future__ import division, print_function, absolute_import
 
 # 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', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI']
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI', 'UnifiedModePtycho']
 
 from CDTools.models.base import CDIModel
 from CDTools.models.simple_ptycho import SimplePtycho
@@ -34,3 +34,5 @@ from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
 from CDTools.models.s_matrix_ptycho import SMatrixPtycho
 from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
 from CDTools.models.rpi import RPI
+from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
+from CDTools.models.unified_mode_ptycho2 import UnifiedModePtycho2

CDTools/models/base.py

@@ -36,6 +36,9 @@ from matplotlib import pyplot as plt
 from matplotlib.widgets import Slider
 from matplotlib import ticker
 import numpy as np
+import threading
+import queue
+import time
 
 __all__ = ['CDIModel']
 
@@ -97,7 +100,7 @@ class CDIModel(t.nn.Module):
         raise NotImplementedError()
 
     def AD_optimize(self, iterations, data_loader,  optimizer,\
-                    scheduler=None, regularization_factor=None):
+                    scheduler=None, regularization_factor=None, thread=True):
         """Runs a round of reconstruction using the provided optimizer
         
         This is the basic automatic differentiation reconstruction tool
@@ -118,44 +121,75 @@ class CDIModel(t.nn.Module):
             Optional, a learning rate scheduler to use
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         """
         # First, calculate the normalization
         normalization = 0
         for inputs, patterns in data_loader:
             normalization += t.sum(patterns).cpu().numpy()
-            
-        for it in range(iterations):
+
+        def run_iteration(stop_event=None):
             loss = 0
             N = 0
             for inputs, patterns in data_loader:
                 N += 1
                 def closure():
+                    # This is just used to allow graceful exit when threading
+                    if stop_event is not None and stop_event.is_set():
+                        exit()
                     optimizer.zero_grad()
                     sim_patterns = self.forward(*inputs)
                     if hasattr(self, 'mask'):
                         loss = self.loss(patterns,sim_patterns, mask=self.mask)
                     else:
                         loss = self.loss(patterns,sim_patterns)
-
+                        
                     if regularization_factor is not None \
                        and hasattr(self, 'regularizer'):
                         loss += self.regularizer(regularization_factor)
-                        
+                    #print(loss)
                     loss.backward()
                     return loss
-
+                
                 loss += optimizer.step(closure).detach().cpu().numpy()
-
+                
             loss /= normalization
             if scheduler is not None:
                 scheduler.step(loss)
-                
-            yield loss
+
+            return loss
+
+        if thread:
+            result_queue = queue.Queue()
+            stop_event = threading.Event()
+            def target():
+                result_queue.put(run_iteration(stop_event))
+            
+        for it in range(iterations):
+            if thread:
+                calc = threading.Thread(target=target, name='calculator', daemon=True)
+                try:
+                    calc.start()
+                    while calc.is_alive():
+                        if hasattr(self, 'figs'):
+                            self.figs[0].canvas.start_event_loop(0.01) 
+                        else:
+                            calc.join()
+                except KeyboardInterrupt as e:
+                    stop_event.set()
+                    print('\nAsking execution thread to stop cleanly - please be patient.')
+                    calc.join()
+                    raise e
+                    
+                yield result_queue.get()
+            else:
+                yield run_iteration()
 
 
     def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005,
                       schedule=False, amsgrad=False, subset=None,
-                      regularization_factor=None):
+                      regularization_factor=None, thread=True):
         """Runs a round of reconstruction using the Adam optimizer
         
         This is generally accepted to be the most robust algorithm for use
@@ -179,6 +213,8 @@ class CDIModel(t.nn.Module):
             Optional, a pattern index or list of pattern indices to use
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         """
 
         if subset is not None:
@@ -203,12 +239,13 @@ class CDIModel(t.nn.Module):
 
         return self.AD_optimize(iterations, data_loader, optimizer,
                                 scheduler=scheduler,
-                                regularization_factor=regularization_factor)
+                                regularization_factor=regularization_factor,
+                                thread=thread)
 
 
     def LBFGS_optimize(self, iterations, dataset, batch_size=None,
                        lr=0.1,history_size=2, subset=None,
-                       regularization_factor=None):
+                       regularization_factor=None, thread=True):
         """Runs a round of reconstruction using the L-BFGS optimizer
         
         This algorithm is often less stable that Adam, however in certain
@@ -232,6 +269,8 @@ class CDIModel(t.nn.Module):
             Optional, a pattern index or list of pattern indices to ues
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         """
         if subset is not None:
             # if just one pattern, turn into a list for convenience
@@ -252,12 +291,14 @@ class CDIModel(t.nn.Module):
                                   lr = lr, history_size=history_size)
 
         return self.AD_optimize(iterations, data_loader, optimizer,
-                                regularization_factor=regularization_factor)
+                                regularization_factor=regularization_factor,
+                                thread=thread)
 
 
     def SGD_optimize(self, iterations, dataset, batch_size=None,
                      lr=0.01, momentum=0, dampening=0, weight_decay=0,
-                     nesterov=False, subset=None, regularization_factor=None):
+                     nesterov=False, subset=None, regularization_factor=None,
+                     thread=True):
         """Runs a round of reconstruction using the SGDoptimizer
         
         This algorithm is often less stable that Adam, but it is simpler
@@ -279,6 +320,8 @@ class CDIModel(t.nn.Module):
             Optional, a pattern index or list of pattern indices to use
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        thread : bool
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         """
 
         if subset is not None:
@@ -303,7 +346,8 @@ class CDIModel(t.nn.Module):
                                 nesterov=nesterov)
 
         return self.AD_optimize(iterations, data_loader, optimizer,
-                                regularization_factor=regularization_factor)
+                                regularization_factor=regularization_factor,
+                                thread=thread)
 
 
     # By default, the plot_list is empty
@@ -337,6 +381,7 @@ class CDIModel(t.nn.Module):
             Whether to update existing plots or plot new ones
         
         """
+                
         first_update = False
         if update and hasattr(self, 'figs') and self.figs:
             figs = self.figs
@@ -375,10 +420,10 @@ class CDIModel(t.nn.Module):
                     try:
                         plotter(self, fig, dataset)
                         plt.title(name)
-                    except (IndexError, KeyError, AttributeError) as e:
+                    except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
                         pass
 
-            except (IndexError, KeyError, AttributeError) as e:
+            except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
                 pass
 
             idx += 1
@@ -386,7 +431,7 @@ class CDIModel(t.nn.Module):
             if update:
                 plt.draw()
                 fig.canvas.start_event_loop(0.001)
-
+                
         if first_update:
             plt.pause(0.05 * len(self.figs))
 

CDTools/models/bragg_2d_ptycho.py

@@ -170,12 +170,13 @@ class Bragg2DPtycho(CDIModel):
             # recall that here we always want the shape of the detector
             # before it's cut down by the detector slice to match the
             # physical detector region
-            det_shape = self.probe[0].shape[:-1]
-
+            probe_shape = self.probe[0].shape[:-1]
+            
             self.k_map, self.intensity_map = \
                 tools.propagators.generate_high_NA_k_intensity_map(
-                    self.probe_basis, self.detector_geometry['basis'],
-                    det_shape,
+                    self.probe_basis,
+                    self.detector_geometry['basis']/ oversampling,
+                    probe_shape,
                     self.detector_geometry['distance'],
                     self.wavelength,dtype=t.float32,
                     lens=lens)

CDTools/models/fancy_ptycho.py

@@ -144,7 +144,7 @@ class FancyPtycho(CDIModel):
             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)
+            surface_normal /= -np.linalg.norm(surface_normal)
 
 
         # Next generate the object geometry from the probe geometry and

CDTools/models/multislice_2d_ptycho.py

@@ -31,7 +31,7 @@ class Multislice2DPtycho(CDIModel):
         super(Multislice2DPtycho,self).__init__()
         self.wavelength = t.Tensor([wavelength])
         self.detector_geometry = copy(detector_geometry)
-        self.dz = dz
+        self.dz = -dz
         self.nz = nz
         det_geo = self.detector_geometry
         if hasattr(det_geo, 'distance'):
@@ -95,22 +95,28 @@ class Multislice2DPtycho(CDIModel):
 
         if obj_support is not None:
             self.obj_support = obj_support
-            self.obj.data = self.obj * obj_support
+            if self.obj.dim() == 3:
+                self.obj.data = self.obj * obj_support
+            elif self.obj.dim() == 4:
+                self.obj.data = self.obj * obj_support[None,...]
         else:
-            self.obj_support = t.ones_like(self.obj)
-
+            if self.obj.dim() == 3:
+                self.obj_support = t.ones_like(self.obj)
+            elif self.obj.dim() == 4:
+                self.obj_support = t.ones_like(self.obj[0])
         self.oversampling = oversampling
 
         spacing = np.linalg.norm(self.probe_basis,axis=0)
         shape = np.array(self.probe.shape[1:-1])
 
         self.bandlimit = bandlimit
-        
+
+        # Big question: Should there be a minus sign before self.dz, or not?
         self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit)
 
         
     @classmethod
-    def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5):
+    def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5, replicate_slice=False):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -185,8 +191,14 @@ class Multislice2DPtycho(CDIModel):
         probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
         probe = t.stack([probe,] + probe_stack)
 
-        obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
-                              
+        # Consider a different start
+        obj = t.zeros(obj_size+(2,))
+        #obj = tools.cmath.expi(t.zeros(obj_size))
+        # If we will use a separate object per slice
+        if not replicate_slice:
+            obj = t.stack([obj]*nz)
+            
+        
         det_geo = dataset.detector_geometry
 
         translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
@@ -212,7 +224,10 @@ class Multislice2DPtycho(CDIModel):
             ro = restrict_obj
             os = np.array(obj_size)
             ps = np.array(probe_shape)
-            obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            if replicate_slice:
+                obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            else:
+                obj_support = t.zeros_like(obj[0].to(dtype=t.float32))
             obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
                         ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
         else:
@@ -251,7 +266,13 @@ class Multislice2DPtycho(CDIModel):
 
         if self.translation_offsets is not None:
             pix_trans += self.translation_scale * self.translation_offsets[index]
-            
+        if len(pix_trans.shape) == 1:
+            pix_trans = [pix_trans]
+            index = [index]
+            strip_first_index = True
+        else:
+            strip_first_index = False
+
         all_exit_waves = []
         for i in range(self.probe.shape[0]):
             # For a Fourier-space probe
@@ -263,45 +284,32 @@ class Multislice2DPtycho(CDIModel):
             #print(self.probe_norm)
             #for i in range(self.nz):
             exit_waves = []
-            if len(pix_trans.shape) == 1:
-                pix_trans = [pix_trans]
+                
             for trans in pix_trans:
                 exit_wave = self.probe_norm * pr
-                for i in range(self.nz-1):
-                    
-                    exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
-                                            self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj.data,
-                                            trans,
-                                            shift_probe=True)
-                    #exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
-                    #                        self.obj_support * cmath.cexpi(self.obj.data/self.nz),
-                    #                        trans)
-                    exit_wave = tools.propagators.near_field(exit_wave,self.as_prop)
+                for i in range(self.nz):
+                    # If only one object slice
+                    if self.obj.dim() == 3:
+                        exit_wave = tools.interactions.ptycho_2D_sinc(
+                            exit_wave,
+                            self.obj_support*cmath.cexpi(self.obj/self.nz),
+                            trans, shift_probe=True)
+                    # If separate slices
+                    elif self.obj.dim() == 4:
+                        exit_wave = tools.interactions.ptycho_2D_sinc(
+                            exit_wave,
+                            self.obj_support*cmath.cexpi(self.obj[i]/self.nz),
+                            trans, shift_probe=True)
+                        
+                    exit_wave = tools.propagators.near_field(
+                        exit_wave,self.as_prop)
                     
                     #tools.plotting.plot_amplitude(exit_wave)
                     #plt.show()
-
-                # only final layer gets a derivative
-                exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
-                                            self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj,
-                                            trans,
-                                            shift_probe=True)
-                #exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
-                #                        self.obj_support * cmath.cexpi(self.obj/self.nz),
-                #                        trans)
-                # One final propagation to enforce the bandlimit
-                exit_wave = tools.propagators.near_field(exit_wave,self.as_prop)
-                exit_waves.append(exit_wave)
                 
+                exit_waves.append(exit_wave)
             exit_waves = t.stack(exit_waves)
             
-
-            if np.array(index).size == 1:
-                index = [index]
-                strip_first_index = True
-            else:
-                strip_first_index = False
-            
             if exit_waves.dim() == 4:
                 exit_waves =  self.weights[index][:,None,None,None] * exit_waves
             else:
@@ -312,7 +320,6 @@ class Multislice2DPtycho(CDIModel):
                 
             all_exit_waves.append(exit_waves)
 
-        
         return t.stack(all_exit_waves)
     
         
@@ -422,14 +429,10 @@ class Multislice2DPtycho(CDIModel):
         ('Subdominant Probe Phase',
          lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
          lambda self: len(self.probe) >=2),
-        #('Object Amplitude', 
-        # lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
-        #('Object Phase',
-        # lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
-        ('Real Part of T', 
-         lambda self, fig: p.plot_amplitude(self.obj[:,:,0].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)),
-        ('Imaginary Part of T',
-         lambda self, fig: p.plot_amplitude(self.obj[:,:,1].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)),
+        ('Integrated Real Part of T', 
+         lambda self, fig: p.plot_real(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)),
+        ('Integrated Imaginary Part of T',
+         lambda self, fig: p.plot_imag(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)),
         ('Corrected Translations',
          lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
         ('Background',

CDTools/models/unified_mode_ptycho.py

@@ -0,0 +1,449 @@
+from __future__ import division, print_function, absolute_import
+
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import cmath
+from CDTools.tools import analysis
+from CDTools.tools import plotting as p
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+__all__ = ['UnifiedModePtycho']
+
+class UnifiedModePtycho(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_guess, rhos_guess,
+                 detector_slice=None,
+                 surface_normal=np.array([0.,0.,1.]),
+                 min_translation = t.Tensor([0,0]),
+                 background = None, translation_offsets=None, mask=None,
+                 translation_scale = 1, saturation=None,
+                 probe_support = None, obj_support=None, oversampling=1):
+        
+        super(UnifiedModePtycho,self).__init__()
+        self.wavelength = t.Tensor([wavelength])
+        self.detector_geometry = copy(detector_geometry)
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = t.Tensor(det_geo['distance'])
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = t.Tensor(det_geo['basis'])
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = t.Tensor(det_geo['corner'])
+
+        self.min_translation = t.Tensor(min_translation)
+
+        self.probe_basis = t.Tensor(probe_basis)
+        self.detector_slice = detector_slice
+        self.surface_normal = t.Tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.BoolTensor(mask)
+        
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 3:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
+        else:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))         
+            
+        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+                                    / self.probe_norm)
+        
+        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+        
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape[:-1])
+
+                
+        self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
+
+        self.rhos = t.nn.Parameter(t.Tensor(rhos_guess).to(t.float32))
+        
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) 
+
+        self.translation_scale = translation_scale
+
+        if probe_support is not None:
+            self.probe_support = probe_support
+        else:
+            self.probe_support = t.ones_like(self.probe[0])
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support
+        else:
+            self.obj_support = t.ones_like(self.obj)
+
+        self.oversampling = oversampling
+
+        
+    @classmethod
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, mixing_mode='unified'):
+        
+        wavelength = dataset.wavelength
+        det_basis = dataset.detector_geometry['basis']
+        det_shape = dataset[0][1].shape
+        distance = dataset.detector_geometry['distance']
+
+        # always do this on the cpu
+        get_as_args = dataset.get_as_args
+        dataset.get_as(device='cpu')
+        (indices, translations), patterns = dataset[:]
+        dataset.get_as(*get_as_args[0],**get_as_args[1])
+
+        # Set to none to avoid issues with things outside the detector
+        if auto_center:
+            center = tools.image_processing.centroid(t.sum(patterns,dim=0))
+        else:
+            center = None
+            
+        # Then, generate the probe geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center,
+                                                   padding=padding,
+                                                   opt_for_fft=opt_for_fft,
+                                                   oversampling=oversampling)
+
+
+        if hasattr(dataset, 'sample_info') and \
+           dataset.sample_info is not None and \
+           'orientation' in dataset.sample_info:
+            surface_normal = dataset.sample_info['orientation'][2]
+        else:
+            surface_normal = np.array([0.,0.,1.])
+
+
+        # If this information is supplied when the function is called,
+        # then we override the information in the .cxi file
+        if scattering_mode in {'t', 'transmission'}:
+            surface_normal = np.array([0.,0.,1.])
+        elif scattering_mode in {'r', 'reflection'}:
+            outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
+            outgoing_dir /= np.linalg.norm(outgoing_dir)
+            surface_normal = outgoing_dir + np.array([0.,0.,1.])
+            surface_normal /= -np.linalg.norm(surface_normal)
+
+
+        # Next generate the object geometry from the probe geometry and
+        # the translations
+        pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
+
+        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
+
+        if hasattr(dataset, 'background') and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        else:
+            background = None
+
+        # Finally, initialize the probe and  object using this information
+        if probe_size is None:
+            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        else:
+            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(cmath.cabs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
+        probe = t.stack([probe,] + probe_stack)
+        #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
+
+        obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
+                              
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
+
+        # 
+        if mixing_mode.lower().strip() == 'unified':
+            rhos = t.zeros(len(dataset),n_modes,n_modes)
+            rhos[:,0,0] = 1
+            for i in range(1,n_modes):
+                rhos[:,i,i] = 1/n_modes
+        
+        if hasattr(dataset, 'mask') and dataset.mask is not None:
+            mask = dataset.mask.to(t.bool)
+        else:
+            mask = None
+
+        if probe_support_radius is not None:
+            probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
+            p_cent = np.array(probe.shape[1:3]).astype(int) // 2
+            psr = int(probe_support_radius)
+            probe_support[p_cent[0]-psr:p_cent[0]+psr,
+                          p_cent[1]-psr:p_cent[1]+psr] = 1
+            probe = probe * probe_support[None,:,:]
+        else:
+            probe_support = None;
+
+        if restrict_obj != -1:
+            ro = restrict_obj
+            os = np.array(obj_size)
+            ps = np.array(probe_shape)
+            obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
+                        ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
+        else:
+            obj_support = None
+
+        return cls(wavelength, det_geo, probe_basis, probe, obj, rhos,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets = translation_offsets,
+                   mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   obj_support=obj_support,
+                   oversampling=oversampling)
+                   
+    
+    def interaction(self, index, translations):
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations,
+                                                             surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+
+        probes = []
+        all_exit_waves = []
+        for i in range(self.probe.shape[0]):
+            # from storing the probe in Fourier space
+            #pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support
+            pr = self.probe[i] * self.probe_support
+            exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
+                                                           self.obj_support * self.obj,
+                                                           pix_trans,
+                                                           shift_probe=True)
+            exit_waves = exit_waves * self.probe_support[...,:,:]
+                
+            all_exit_waves.append(exit_waves)
+
+        
+        return t.stack(all_exit_waves)
+    
+        
+    def forward_propagator(self, wavefields):
+        return tools.propagators.far_field(wavefields)
+
+
+    def backward_propagator(self, wavefields):
+        return tools.propagators.inverse_far_field(wavefields)
+
+    
+    def measurement(self, wavefields, indices):
+        #return tools.measurements.density_matrix(wavefields,self.rhos[indices],
+        #                    detector_slice=self.detector_slice,
+        #                    saturation=self.saturation,
+        #                    oversampling=self.oversampling)
+        return tools.measurements.quadratic_background(wavefields,
+                            self.background, self.rhos[indices], 
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.density_matrix,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+
+
+    def forward(self, *args):
+        """The complete forward model
+        
+        We need to override this to enable the wavefield mixing at the
+        level of the measurement function
+        """
+        indices = args[0]
+        return self.measurement(self.forward_propagator(self.interaction(*args)),indices)
+
+    def loss(self, sim_data, real_data, mask=None):
+        return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
+
+    
+    def to(self, *args, **kwargs):
+        super(UnifiedModePtycho, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args,**kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
+
+        if self.mask is not None:
+            self.mask = self.mask.to(*args, **kwargs)
+
+            
+        self.min_translation = self.min_translation.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.probe_norm = self.probe_norm.to(*args,**kwargs)
+        self.probe_support = self.probe_support.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+        
+    def sim_to_dataset(self, args_list):
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        surface_normal = self.surface_normal.detach().cpu().numpy()
+        xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
+        xsurfacevec /= np.linalg.norm(xsurfacevec)
+        ysurfacevec = np.cross(surface_normal, xsurfacevec)
+        ysurfacevec /= np.linalg.norm(ysurfacevec)
+        orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
+        
+        sample_info = {'description': 'A simulated sample',
+                       'orientation': orientation}
+
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho2DDataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
+
+    def get_rhos(self):
+
+        rhos_out = np.zeros([self.rhos.shape[0],
+                             self.rhos.shape[1],self.rhos.shape[2]],
+                            dtype=np.complex64)
+        for (i,j) in ((i,j) for i in range(rhos_out.shape[1])
+                      for j in range(rhos_out.shape[2])):
+            if i == j:
+                rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
+            if i < j: # upper triangle, real part
+                rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
+                rhos_out[:,j,i] += self.rhos.data[:,i,j].cpu().detach().numpy()
+            if i > j: # upper triangle, real part
+                rhos_out[:,j,i] += 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
+                rhos_out[:,i,j] -= 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
+        return rhos_out
+    
+    def tidy_probes(self, normalization=1):
+        """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
+        
+        """
+        rhos = self.get_rhos()
+        overall_rho = np.mean(rhos,axis=0)
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        ortho_probes, A = analysis.orthogonalize_probes(probe,
+                                        density_matrix=overall_rho,
+                                        keep_transform=True,
+                                        normalize=True)
+        Aconj = A.conj()
+        Atrans = np.transpose(A)
+        new_rhos = np.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1)
+
+        new_rhos /= normalization
+        ortho_probes *= np.sqrt(normalization)
+        #print(np.dot(Aconjinv,rhos).shape)
+        #print(np.dot(rhos,Aconj).shape)
+        new_rhos = cmath.complex_to_torch(new_rhos).to(
+            dtype=self.rhos.dtype,device=self.rhos.device)
+
+        # This repacks the data into the format used internally
+        for (i,j) in ((i,j) for i in range(new_rhos.shape[1])
+                      for j in range(new_rhos.shape[2])):
+            if i == j:
+                self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
+            if i < j: # upper triangle, real part
+                self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
+                self.rhos.data[:,j,i] = new_rhos[:,i,j,1]
+
+        
+        self.probe.data = cmath.complex_to_torch(ortho_probes).to(
+            device=self.probe.device,dtype=self.probe.dtype)
+        
+    
+    def corrected_translations(self,dataset):
+        translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
+        t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
+        return translations + t_offset
+
+
+    # Needs to be updated to allow for plotting to an existing figure
+    plot_list = [
+        ('Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Subdominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Subdominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
+         lambda self: len(self.probe) >=2),
+        ('Von Neumann Entropy (only accurate after tidy_probes)',
+          lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_vn_entropy(self.get_rhos()), fig=fig)),
+        ('Object Amplitude', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = cmath.torch_to_complex(self.obj.detach().cpu())
+        background = self.background.detach().cpu().numpy()**2
+        weights = self.weights.detach().cpu().numpy()
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'obj':obj,
+                'background':background,
+                'weights':weights}

CDTools/models/unified_mode_ptycho2.py

@@ -0,0 +1,490 @@
+from __future__ import division, print_function, absolute_import
+
+import torch as t
+from CDTools.models import CDIModel
+from CDTools.datasets import Ptycho2DDataset
+from CDTools import tools
+from CDTools.tools import cmath
+from CDTools.tools import analysis
+from CDTools.tools import plotting as p
+from matplotlib import pyplot as plt
+from datetime import datetime
+import numpy as np
+from copy import copy
+
+__all__ = ['UnifiedModePtycho2']
+
+class UnifiedModePtycho2(CDIModel):
+
+    def __init__(self, wavelength, detector_geometry,
+                 probe_basis,
+                 probe_guess, obj_guess, Ws_guess,
+                 detector_slice=None,
+                 surface_normal=np.array([0.,0.,1.]),
+                 min_translation = t.Tensor([0,0]),
+                 background = None, translation_offsets=None, mask=None,
+                 translation_scale = 1, saturation=None,
+                 probe_support = None, obj_support=None, oversampling=1):
+        
+        super(UnifiedModePtycho2,self).__init__()
+        self.wavelength = t.Tensor([wavelength])
+        self.detector_geometry = copy(detector_geometry)
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = t.Tensor(det_geo['distance'])
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = t.Tensor(det_geo['basis'])
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = t.Tensor(det_geo['corner'])
+
+        self.min_translation = t.Tensor(min_translation)
+
+        self.probe_basis = t.Tensor(probe_basis)
+        self.detector_slice = detector_slice
+        self.surface_normal = t.Tensor(surface_normal)
+        
+        self.saturation = saturation
+        
+        if mask is None:
+            self.mask = mask
+        else:
+            self.mask = t.BoolTensor(mask)
+        
+        # We rescale the probe here so it learns at the same rate as the
+        # object
+        if probe_guess.dim() > 3:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
+        else:
+            self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))         
+            
+        self.probe = t.nn.Parameter(probe_guess.to(t.float32)
+                                    / self.probe_norm)
+        
+        self.obj = t.nn.Parameter(obj_guess.to(t.float32))
+        
+        if background is None:
+            if detector_slice is not None:
+                background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
+            else:
+                background = 1e-6 * t.ones(self.probe[0].shape[:-1])
+
+                
+        self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
+
+        if type(Ws_guess) == type(t.zeros(1)):
+            self.Ws = t.nn.Parameter(Ws_guess.to(t.float32))        
+        else:
+            self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).to(t.float32))
+        
+        if translation_offsets is None:
+            self.translation_offsets = None
+        else:
+            self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) 
+
+        self.translation_scale = translation_scale
+
+        if probe_support is not None:
+            self.probe_support = probe_support
+        else:
+            self.probe_support = t.ones_like(self.probe[0])
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support
+        else:
+            self.obj_support = t.ones_like(self.obj)
+
+        self.oversampling = oversampling
+
+        
+    @classmethod
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=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')
+        (indices, translations), patterns = dataset[:]
+        dataset.get_as(*get_as_args[0],**get_as_args[1])
+
+        # Set to none to avoid issues with things outside the detector
+        if auto_center:
+            center = tools.image_processing.centroid(t.sum(patterns,dim=0))
+        else:
+            center = None
+            
+        # Then, generate the probe geometry from the dataset
+        ewg = tools.initializers.exit_wave_geometry
+        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+                                                   det_shape,
+                                                   wavelength,
+                                                   distance,
+                                                   center=center,
+                                                   padding=padding,
+                                                   opt_for_fft=opt_for_fft,
+                                                   oversampling=oversampling)
+
+
+        if hasattr(dataset, 'sample_info') and \
+           dataset.sample_info is not None and \
+           'orientation' in dataset.sample_info:
+            surface_normal = dataset.sample_info['orientation'][2]
+        else:
+            surface_normal = np.array([0.,0.,1.])
+
+
+        # If this information is supplied when the function is called,
+        # then we override the information in the .cxi file
+        if scattering_mode in {'t', 'transmission'}:
+            surface_normal = np.array([0.,0.,1.])
+        elif scattering_mode in {'r', 'reflection'}:
+            outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
+            outgoing_dir /= np.linalg.norm(outgoing_dir)
+            surface_normal = outgoing_dir + np.array([0.,0.,1.])
+            surface_normal /= -np.linalg.norm(surface_normal)
+
+
+        # Next generate the object geometry from the probe geometry and
+        # the translations
+        pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
+
+        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
+
+        if hasattr(dataset, 'background') and dataset.background is not None:
+            background = t.sqrt(dataset.background)
+        else:
+            background = None
+
+        # Finally, initialize the probe and  object using this information
+        if probe_size is None:
+            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        else:
+            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+
+
+        # Now we initialize all the subdominant probe modes
+        probe_max = t.max(cmath.cabs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
+        probe = t.stack([probe,] + probe_stack)
+        #probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
+
+        obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
+                              
+        det_geo = dataset.detector_geometry
+
+        translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
+
+        # dm_rank defines the rank of the shot-by-shot density matrices
+        if dm_rank > n_modes:
+            raise KeyError('Density matrix rank cannot be greater than the number of modes')
+        elif dm_rank != 0:
+            if dm_rank == -1:
+                dm_rank = n_modes
+            Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
+            Ws[:,0,0,0] = 1
+            for i in range(1,dm_rank):
+                Ws[:,i,i,0] = 1/np.sqrt(n_modes)
+        else:
+            # dm_rank=0 is a special case defining a purely stable, incoherent
+            # mode mixing model. This is passed on by defining 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_like(probe[0].to(dtype=t.float32))
+            p_cent = np.array(probe.shape[1:3]).astype(int) // 2
+            psr = int(probe_support_radius)
+            probe_support[p_cent[0]-psr:p_cent[0]+psr,
+                          p_cent[1]-psr:p_cent[1]+psr] = 1
+            probe = probe * probe_support[None,:,:]
+        else:
+            probe_support = None;
+
+        if restrict_obj != -1:
+            ro = restrict_obj
+            os = np.array(obj_size)
+            ps = np.array(probe_shape)
+            obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
+                        ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
+        else:
+            obj_support = None
+
+        return cls(wavelength, det_geo, probe_basis, probe, obj, Ws,
+                   detector_slice=det_slice,
+                   surface_normal=surface_normal,
+                   min_translation=min_translation,
+                   translation_offsets = translation_offsets,
+                   mask=mask, background=background,
+                   translation_scale=translation_scale,
+                   saturation=saturation,
+                   probe_support=probe_support,
+                   obj_support=obj_support,
+                   oversampling=oversampling)
+                   
+    
+    def interaction(self, index, translations):
+        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                             translations,
+                                                             surface_normal=self.surface_normal)
+        pix_trans -= self.min_translation
+
+        if self.translation_offsets is not None:
+            pix_trans += self.translation_scale * self.translation_offsets[index]
+
+        Ws = self.Ws[index]
+        
+        if type(index) == type(0):
+            index = [index]
+            Ws = [Ws]
+            pix_trans = [pix_trans]
+            single_frame = True
+        else:
+            single_frame = False
+
+        probes = []
+        all_exit_waves = []
+
+        # This is the case if a purely stable, incoherent model is defined.
+        if len(Ws[0].shape) == 0:
+            # What we do here is generate an identity matrix, and multiply
+            # that identity matrix by the per-frame weight.
+            Ws = [W * t.stack([t.eye(self.probe.shape[0]),
+                               t.zeros([self.probe.shape[0]]*2)],dim=-1).to(
+                                   dtype=W.dtype, device=W.device)
+                  for W in Ws]
+            
+        
+        # Outer iteration is the mode index iteration
+        for i in range(Ws[0].shape[0]):
+            # Now we need to separately treat each mode
+            exit_waves = []
+            for W, pix_tran in zip(Ws, pix_trans):
+                # from storing the probe in Fourier space
+                pr = [cmath.cmult(W[i,j,:], self.probe[j] * self.probe_support)
+                      for j in range(self.probe.shape[0])]
+                
+                pr = t.sum(t.stack(pr), axis=0)
+                
+                exit_waves.append(self.probe_norm *
+                                  tools.interactions.ptycho_2D_sinc(
+                                      pr, self.obj_support * self.obj,
+                                      pix_tran, shift_probe=True))
+                
+            exit_waves = t.stack(exit_waves)
+
+            if single_frame:
+                exit_waves = exit_waves[0]
+            # Multiply again by probe support to suppress the fringes from the
+            # sinc-interpolated shift
+            exit_waves = exit_waves * self.probe_support[...,:,:]
+
+            all_exit_waves.append(exit_waves)
+
+        
+        return t.stack(all_exit_waves)
+    
+        
+    def forward_propagator(self, wavefields):
+        return tools.propagators.far_field(wavefields)
+
+
+    def backward_propagator(self, wavefields):
+        return tools.propagators.inverse_far_field(wavefields)
+
+    
+    def measurement(self, wavefields):
+        return tools.measurements.quadratic_background(wavefields,
+                            self.background, 
+                            detector_slice=self.detector_slice,
+                            measurement=tools.measurements.incoherent_sum,
+                            saturation=self.saturation,
+                            oversampling=self.oversampling)
+
+    def loss(self, sim_data, real_data, mask=None):
+        return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
+        #return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
+
+    
+    def to(self, *args, **kwargs):
+        super(UnifiedModePtycho2, self).to(*args, **kwargs)
+        self.wavelength = self.wavelength.to(*args,**kwargs)
+        # move the detector geometry too
+        det_geo = self.detector_geometry
+        if hasattr(det_geo, 'distance'):
+            det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
+        if hasattr(det_geo, 'basis'):
+            det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
+        if hasattr(det_geo, 'corner'):
+            det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
+
+        if self.mask is not None:
+            self.mask = self.mask.to(*args, **kwargs)
+
+            
+        self.min_translation = self.min_translation.to(*args,**kwargs)
+        self.probe_basis = self.probe_basis.to(*args,**kwargs)
+        self.probe_norm = self.probe_norm.to(*args,**kwargs)
+        self.probe_support = self.probe_support.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
+        self.surface_normal = self.surface_normal.to(*args, **kwargs)
+
+        
+    def sim_to_dataset(self, args_list):
+        # In the future, potentially add more control
+        # over what metadata is saved (names, etc.)
+        
+        # First, I need to gather all the relevant data
+        # that needs to be added to the dataset
+        entry_info = {'program_name': 'CDTools',
+                      'instrument_n': 'Simulated Data',
+                      'start_time': datetime.now()}
+
+        surface_normal = self.surface_normal.detach().cpu().numpy()
+        xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
+        xsurfacevec /= np.linalg.norm(xsurfacevec)
+        ysurfacevec = np.cross(surface_normal, xsurfacevec)
+        ysurfacevec /= np.linalg.norm(ysurfacevec)
+        orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
+        
+        sample_info = {'description': 'A simulated sample',
+                       'orientation': orientation}
+
+        
+        detector_geometry = self.detector_geometry
+        mask = self.mask
+        wavelength = self.wavelength
+        indices, translations = args_list
+        
+        # Then we simulate the results
+        data = self.forward(indices, translations)
+
+        # And finally, we make the dataset
+        return Ptycho2DDataset(translations, data,
+                                 entry_info = entry_info,
+                                 sample_info = sample_info,
+                                 wavelength=wavelength,
+                                 detector_geometry=detector_geometry,
+                                 mask=mask)
+
+    def get_rhos(self):
+        # If this is not a purely stable model
+        if len(self.Ws.shape) >= 2:
+            Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
+            rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
+            return rhos_out
+        else:
+            return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0],
+                            dtype=np.complex64)
+        
+    def tidy_probes(self, normalization=1):
+        """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
+        
+        """
+
+        # Must also implement a version that works appropriately with
+        # a purely incoherent model
+        
+        #
+        # 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. This could avoid potential stability issues
+        # due to the existence of zero eigenvalues in the full rho matrix
+        # when dm_rank < n_modes
+        # 
+        
+        rhos = self.get_rhos()
+        overall_rho = np.mean(rhos,axis=0)
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        ortho_probes, A = analysis.orthogonalize_probes(probe,
+                                        density_matrix=overall_rho,
+                                        keep_transform=True,
+                                        normalize=True)
+        Aconj = A.conj()
+        Atrans = np.transpose(A)
+        new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
+
+        new_rhos /= normalization
+        ortho_probes *= np.sqrt(normalization)
+
+        dm_rank = self.Ws.shape[1]
+        
+        new_Ws = []
+        for rho in new_rhos:
+            # These are returned from smalles to largest - we want to keep
+            # the largest ones
+            w,v = np.linalg.eigh(rho)
+            w = w[::-1][:dm_rank]
+            v = v[:,::-1][:,:dm_rank]
+            new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
+
+        new_Ws = np.array(new_Ws)
+
+        self.Ws.data = cmath.complex_to_torch(new_Ws).to(
+            dtype=self.Ws.dtype,device=self.Ws.device)
+        
+        self.probe.data = cmath.complex_to_torch(ortho_probes).to(
+            device=self.probe.device,dtype=self.probe.dtype)
+        
+    
+    def corrected_translations(self,dataset):
+        translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
+        t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
+        return translations + t_offset
+
+
+    # Needs to be updated to allow for plotting to an existing figure
+    plot_list = [
+        ('Dominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Dominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
+        ('Subdominant Probe Amplitude',
+         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Subdominant Probe Phase',
+         lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
+         lambda self: len(self.probe) >=2),
+        ('Average Density Matrix Amplitudes',
+         lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
+         lambda self: len(self.Ws.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),
+         lambda self: len(self.Ws.shape) >=2),
+        ('Object Amplitude', 
+         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
+        ('Corrected Translations',
+         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
+        ('Background',
+         lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
+    
+    def save_results(self, dataset):
+        basis = self.probe_basis.detach().cpu().numpy()
+        translations = self.corrected_translations(dataset).detach().cpu().numpy()
+        probe = cmath.torch_to_complex(self.probe.detach().cpu())
+        probe = probe * self.probe_norm.detach().cpu().numpy()
+        obj = cmath.torch_to_complex(self.obj.detach().cpu())
+        background = self.background.detach().cpu().numpy()**2
+        Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
+        
+        return {'basis':basis, 'translation':translations,
+                'probe':probe,'obj':obj,
+                'background':background,
+                'Ws':Ws}

CDTools/tools/analysis/analysis.py

@@ -12,14 +12,14 @@ import numpy as np
 from CDTools.tools import cmath
 from CDTools.tools import image_processing as ip
 from scipy import fftpack
+from scipy import linalg as sla
 
-__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions',
+__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
            'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
-           'calc_frc']
+           'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction']
 
 
-from matplotlib import pyplot as plt
-def orthogonalize_probes(probes):
+def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
     """Orthogonalizes a set of incoherently mixing probes
 
     The strategy is to define a reduced orthogonal basis that spans
@@ -27,55 +27,104 @@ def orthogonalize_probes(probes):
     defined by the probes in that basis. After diagonalization, the
     eigenvectors can be recast into the original basis and returned
 
+    By default, it assumes that the set of probes are defined just as
+    standard incoherently mixing probe modes, and orthogonalizes them.
+    However, if a density matrix is explicitly given, it will instead
+    consider the problem of extracting the eigenbasis of the matrix
+    probes * denstity_matrix * probes^dagger, where probes is the
+    column matrix of the given probe functions. This latter problem arises
+    in the generalization of the probe mixing model, and reduces to the
+    simpler case when the density matrix is equal to the identity matrix
+    
+    If the parameter "keep_transform" is set, the function will additionally
+    return the matrix A such that A * ortho_probes^dagger = probes^dagger
+
+    If the parameter "normalize" is False (as is the default), the variation
+    in intensities in the probe modes will be kept in the probe modes, as is
+    natural for a purely incoherent model. If it is set to "True", the
+    returned probe modes will all be normalized instead.
+
     Parameters
     ----------
     probes : array
         An l x n x m complex array representing  a stack of probes
+    density_matrix : np.array
+        An optional l x l density matrix further elaborating on the state
+    keep_transform : bool
+        Default False, whether to return the map from probes to ortho_probes
+    normalize : bool
+        Default False, whether to normalize the probe modes
 
     Returns
     -------
     ortho_probes: array
         An l x n x m complex array representing a stack of probes
     """
-
+    
     try:
         probes = cmath.torch_to_complex(probes.detach().cpu())
         send_to_torch = True
     except:
         send_to_torch = False
 
-    bases = []
-    coefficients = np.zeros((probes.shape[0],probes.shape[0]), dtype=np.complex64)
-    for i, probe in enumerate(probes):
-        ortho_probe = np.copy(probe)
-        for j, basis in enumerate(bases):
-            coefficients[j,i] = np.sum(basis.conj()*ortho_probe)
-            ortho_probe -= basis * coefficients[j,i]
+    from matplotlib import pyplot as plt
+    # We can do the orthogonalization with an SVD, so first we have to
+    # reshape the final two dimensions (the image shape) into a single
+    # vectorized dimension. This matrix is probes^dagger, hence the
+    # conjugation
+    probes_mat = probes.reshape(probes.shape[0],
+                                probes.shape[1]*probes.shape[2])
 
+    if density_matrix is None:
+        density_matrix = np.eye(probes.shape[0])
+    
+        
+    # next we want to extract the eigendecomposition of the density matrix
+    # itself
+    w,v = sla.eigh(density_matrix)
+    w,v = np.linalg.eigh(density_matrix)
+    w = w[::-1]
+    v = v[:,::-1]
 
-        coefficients[i,i] = np.sqrt(np.sum(np.abs(ortho_probe)**2))
-        bases.append(ortho_probe / coefficients[i,i])
+    # We do this just to avoid total failure when the density
+    # matrix is not positive definite.
+    w = np.abs(w)
+    # Note: this will fail if any w are less than zero
+    # Probably should figure out a good way to deal with that
+    B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose())
 
+    #u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    u,s,vh = sla.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
+    
 
-    density_mat = coefficients.dot(np.conj(coefficients).transpose())
-    eigvals, eigvecs = np.linalg.eigh(density_mat)
+    if normalize:
+        ortho_probes = vh.reshape(probes.shape[0],
+                                  probes.shape[1],
+                                  probes.shape[2])
 
-    ortho_probes = []
-    for i in range(len(eigvals)):
-        coefficients = np.sqrt(eigvals[i]) * eigvecs[:,i]
-        probe = np.zeros(bases[0].shape, dtype=np.complex64)
-        for coefficient, basis in zip(coefficients, bases):
-            probe += basis * coefficient
-        ortho_probes.append(probe)
-
-
-    if send_to_torch:
-        return cmath.complex_to_torch(np.stack(ortho_probes[::-1]))
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,np.dot(u,np.diag(s)))
+        #A_dagger = np.dot(np.linalg.pinv(np.diag(s)),
+        #                  np.dot(np.transpose(u).conj(),B_dagger))
     else:
-        return np.stack(ortho_probes[::-1])
-
+        ortho_probes = np.dot(np.diag(s),vh).reshape(probes.shape[0],
+                                                     probes.shape[1],
+                                                     probes.shape[2])
+        B_dagger_inv = np.linalg.pinv(B_dagger)
+        A = np.dot(B_dagger_inv,u)
+        #A_dagger = np.dot(np.transpose(u).conj(),B_dagger)
+    
+    if send_to_torch:
+        ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes))
+        A = cmath.complex_to_torch(A)
+        #A_dagger = cmath.complex_to_torch(A_dagger)
 
+    if keep_transform:
+        return ortho_probes, A#_dagger
+    else:
+        return ortho_probes 
 
+    
 def standardize(probe, obj, obj_slice=None, correct_ramp=False):
     """Standardizes a probe and object to prepare them for comparison
 
@@ -512,3 +561,66 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
         threshold = t.tensor(threshold)
 
     return bins[:-1], frc, threshold
+
+
+def calc_vn_entropy(matrix):
+    """Calculates the Von Neumann entropy of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to calculate the entropy of
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The entropy or entropies of the arrays
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        eigs = [eig / np.sum(eig) for eig in eigs]
+        # And calculate the VN entropy!
+        entropies = [-np.sum(eig*np.log(eig)) for eig in eigs]
+        return np.array(entropies)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        entropy = -np.sum(eig*np.log(eig))/np.sum(eig)
+        return -np.trace(np.dot(matrix,sla.logm(matrix))) 
+
+def calc_top_mode_fraction(matrix):
+    """Calculates the fraction of total power in the top mode of a density matrix
+    
+    Will either accept a single matrix, or a stack of matrices. Matrices
+    are assumed to be Hermetian and positive definite, to be well-formed
+    density matrices
+    
+    Parameters
+    ----------
+    matrix : np.array
+        The nxn matrix or lxnxn stack of matrices to work from
+
+    
+    Returns
+    -------
+    entropy: float or np.array
+        The fraction of power in the top mode of each matrix
+    """
+
+    if len(matrix.shape) == 3:
+        # Get the eigenvalues
+        eigs = [np.linalg.eigh(mat)[0] for mat in matrix]
+        # Normalize them to match standard density matrix form
+        fractions = [np.max(eig) / np.sum(eig) for eig in eigs]
+        return np.array(fractions)
+    else:
+        eig = np.linalg.eigh(matrix)[0]
+        fraction = np.max(eig) / np.sum(eig)
+        return fraction

CDTools/tools/cmath/cmath.py

@@ -13,7 +13,7 @@ import torch as t
 
 
 __all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj',
-           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi']
+           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi']
 
 
 #

CDTools/tools/interactions/interactions.py

@@ -140,7 +140,15 @@ def project_translations_to_sample(sample_basis, translations):
     describing the location of the probe's intersection with the sample
     plane, assuming the basis found in sample_basis is used. Second, an
     amount that the probe must be propagated along the z-axis to reach
-    that location on the sample plane.
+    the sample plane at the given location. This includes both the effect of
+    the tilted sample plane and any explicitly defined motion along the z axis
+    of the probe-forming optic as included in the input translations.
+
+
+    Note that because of the sign convention (that this function returns th
+    relative amount the probe needs to be propagated to reach any given
+    location), a positive motion along the z-axis of the probe forming optics
+    will lead to a negative propagation distance.
     
     The assumed geometry is incoming radiation with a wavevector parallel
     to the +z axis, [0,0,1].
@@ -177,6 +185,8 @@ def project_translations_to_sample(sample_basis, translations):
         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)))
     
@@ -205,8 +215,9 @@ def project_translations_to_sample(sample_basis, translations):
         single_translation = True
 
     pixel_translations = t.mm(translations, sample_projection)
-    propagations = t.mm(translations, prop_projection)
-    
+    propagations = t.mm(translations, prop_projection) \
+        - t.mm(translations,propagation_dir[:,None])
+
     if single_translation:
         return pixel_translations[0], propagations[0]
     else:

CDTools/tools/measurements/measurements.py

@@ -17,7 +17,8 @@ from torch.nn.functional import avg_pool2d
 # intensity pattern on a detector
 #
 
-__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
+__all__ = ['intensity', 'incoherent_sum', 'density_matrix',
+           'quadratic_background']
 
 
 def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -66,6 +67,101 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
         return t.clamp(output + epsilon,0,saturation)
 
 
+def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
+    """Returns the intensities associated with a given density matrix state
+
+    The essential idea is that the most general description of a light field
+    at the detector plane will consist of a density matrix state. Here, that
+    low rank density matrix state is encoded as a set of basis wavefields
+    and a density matrix in that basis.
+
+    For computational efficiency, the density matrix is coded in an unusual
+    format. The density matrix formally is a complex Hermetian matrix,
+    which also happens to be positive definite. Here, we store it as a
+    real-valued matrix, where the upper triangle corresponds to the real
+    part of the elements in the upper triangle, and the lower triangle
+    corresponds to the imaginary parts. The elements on the diagonal are
+    purely real, and are stored as they are.
+
+    As with other multi-mode measurement functions, the modes are stored in
+    the first index, and the index of the diffraction pattern in the stack
+    of diffraction patterns is the second index. The stack-direction index
+    can be omitted if only a single pattern needs to be simulated
+    
+    It is important to note that this method does not inforce the positive
+    definiteness of the density matrix, this it is possible for negative
+    values of intensity to appear if the underlying density matrix passed
+    to this method is not positive definite
+
+    Parameters
+    ----------
+    wavefields : torch.Tensor
+        An Lx(Jx)MxNx2 stack of complex wavefields
+    density_matrix : torch.Tensor
+        A (Jx)LxL stack of real-valued representations of density matrices, as per above
+    saturation : float
+        Optional, a maximum saturation value to clamp the resulting intensities to
+    oversampling : int
+        Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
+
+    Returns
+    -------
+    sim_patterns : torch.Tensor 
+        A real Lx(Jx)MxN array storing the incoherently summed intensities
+
+    """
+    
+    #if wavefields.dim() == 4:
+    #    wavefields.unsqueeze(1)
+    #    single_frame = True
+    #elif wavefields.dim() == 5:
+    #    single_frame=False
+        
+    output = t.zeros(wavefields.shape[1:-1],
+                     dtype=wavefields.dtype,
+                     device=wavefields.device)
+    
+    # flat is better than nested, but simple is better than complex...
+    for (i,j) in ((i,j) for i in range(density_matrix.shape[-2])
+                  for j in range(density_matrix.shape[-1])):
+        if i == j: # diagonal
+            output += density_matrix[...,i,j,None,None] \
+                * cmath.cabssq(wavefields[i])
+        if i < j: # upper triangle, real part 
+            output += 2 * density_matrix[...,i,j,None,None] \
+                * (wavefields[i,...,0] * wavefields[j,...,0]
+                   +  wavefields[i,...,1] * wavefields[j,...,1])
+        if i > j: # lower triangle, imaginary part
+            # We pull the i,jth element from the density matrix,
+            # but this correponds to wavefield j and wavefield i,
+            # unlike above where it was wavefield i and j (swapped order).
+            # We also get the one negative sign because this is the imaginary
+            # part
+            output += 2 * density_matrix[...,i,j,None,None] \
+                * (wavefields[j,...,0] * wavefields[i,...,1]
+                   -  wavefields[j,...,1] * wavefields[i,...,0])
+        
+    # Now we apply oversampling
+    if oversampling != 1:
+        if wavefields.dim() == 4:
+            output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
+        else:
+            output = avg_pool2d(output, oversampling)
+            
+    # Then we grab the detector slice
+    if detector_slice is not None:
+        if wavefields.dim() == 4:
+            output = output[detector_slice]
+        else:
+            output = output[(np.s_[:],) + detector_slice]
+
+    if saturation is None:
+        return t.clamp(output,min=0) + epsilon
+    else:
+        return t.clamp(output + epsilon,0,saturation)
+    
+            
+            
 def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
     """Returns the incoherent sum of the intensities of the wavefields
     
@@ -94,7 +190,6 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
     sim_patterns : torch.Tensor 
         A real JXMxN array storing the incoherently summed intensities
     """
-    # This syntax just adds an axis to the slice to preserve the J direction
 
     output = t.sum(cmath.cabssq(wavefields),dim=0) 
 
@@ -118,7 +213,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
         return t.clamp(output + epsilon,0,saturation)
 
 
-def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
+def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
     """Returns the intensity of a wavefield plus a background
     
     The intensity is calculated via the given measurment function 
@@ -148,10 +243,10 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
     """
     
     if detector_slice is None:
-        output = measurement(wavefield, epsilon=epsilon,
+        output = measurement(wavefield, *args, epsilon=epsilon,
                              oversampling=oversampling) + background**2
     else:
-        output = measurement(wavefield, detector_slice,
+        output = measurement(wavefield, *args, detector_slice=detector_slice,
                              epsilon=epsilon, oversampling=oversampling) \
                              + background**2
 

CDTools/tools/plotting/plotting.py

@@ -17,7 +17,7 @@ from matplotlib.colors import hsv_to_rgb
 
 __all__ = ['colorize', 'plot_amplitude', 'plot_phase',
            'plot_colorized', 'plot_translations', 'get_units_factor',
-           'plot_nanomap']
+           'plot_nanomap', 'plot_real', 'plot_imag']
 
 
 def colorize(z):
@@ -82,6 +82,140 @@ def get_units_factor(units):
     return factor
 
 
+def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
+    """Plots the real part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    if isinstance(im, t.Tensor):
+        real = im[...,0].detach().cpu().numpy()
+    else:
+        real = np.real(im)
+
+    #Plot in a basis if it exists, otherwise dont
+    if basis is not None:
+        if isinstance(basis,t.Tensor):
+            basis = basis.detach().cpu().numpy()
+        # This fails if the 
+        basis_norm = np.linalg.norm(basis, axis = 0)
+        basis_norm = basis_norm * get_units_factor(units)
+
+        extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]]
+    else:
+        extent=None
+
+    plt.imshow(real, cmap = cmap, extent = extent)
+    cbar = plt.colorbar()
+    cbar.set_label('Real Part (a.u.)')
+
+    if basis is not None:
+        plt.xlabel('X (' + units + ')')
+        plt.ylabel('Y (' + units + ')')
+    else:
+        plt.xlabel('j (pixels)')
+        plt.ylabel('i (pixels)')
+
+    return fig
+
+
+def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
+    """Plots the imaginary part of a complex array with dimensions NxM
+
+    If a figure is given explicitly, it will clear that existing figure and
+    plot over it. Otherwise, it will generate a new figure.
+
+    If a basis is explicitly passed, the image will be plotted in real-space
+    coordinates
+
+    Parameters
+    ----------
+    im : array
+        An complex array with dimensions NxM
+    fig : matplotlib.figure.Figure
+        Default is a new figure, a matplotlib figure to use to plot
+    basis : np.array
+        Optional, the 3x2 probe basis
+    units : str
+        The length units to mark on the plot, default is um
+    cmap : str
+        Default is 'viridis', the colormap to plot with
+    \\**kwargs
+        All other args are passed to fig.add_subplot(111, \\**kwargs)
+
+    Returns
+    -------
+    used_fig : matplotlib.figure.Figure
+        The figure object that was actually plotted to.
+    """
+    if fig is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, **kwargs)
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    if isinstance(im, t.Tensor):
+        imag = im[...,1].detach().cpu().numpy()
+    else:
+        imag = np.imag(im)
+
+    #Plot in a basis if it exists, otherwise dont
+    if basis is not None:
+        if isinstance(basis,t.Tensor):
+            basis = basis.detach().cpu().numpy()
+        # This fails if the 
+        basis_norm = np.linalg.norm(basis, axis = 0)
+        basis_norm = basis_norm * get_units_factor(units)
+
+        extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]]
+    else:
+        extent=None
+
+    plt.imshow(imag, cmap = cmap, extent = extent)
+    cbar = plt.colorbar()
+    cbar.set_label('Imaginary Part (a.u.)')
+
+    if basis is not None:
+        plt.xlabel('X (' + units + ')')
+        plt.ylabel('Y (' + units + ')')
+    else:
+        plt.xlabel('j (pixels)')
+        plt.ylabel('i (pixels)')
+
+    return fig
+
+
 def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
     """Plots the amplitude of a complex array with dimensions NxM
 

CDTools/tools/propagators/propagators.py

@@ -205,7 +205,7 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     
     intensity_map = t.Tensor(intensity_map).to(*args, **kwargs)
 
-
+    
     # This accounts for the implied phase ramp along the exit wave direction
     # In other words, it prevents the diffraction pattern from sliding off the
     # detector when the sample is tilted but represented by an object with

examples/gold_ball_ptycho.py

@@ -3,6 +3,7 @@ from __future__ import division, print_function, absolute_import
 import CDTools
 from matplotlib import pyplot as plt
 import pickle
+import time
 
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@@ -18,8 +19,8 @@ dataset.get_as(device='cuda')
 
 for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
     # And we liveplot the updates to the model as they happen
-    model.inspect(dataset)
     print(i,loss)
+    model.inspect(dataset)
 
 # And we save the reconstruction out to a file
 with open('example_reconstructions/gold_balls.pickle', 'wb') as f:

examples/simple_ptycho.py

@@ -6,7 +6,7 @@ from matplotlib import pyplot as plt
 # First, we load an example dataset from a .cxi file
 filename = 'example_data/lab_ptycho_data.cxi'
 dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
-
+plt.ion()
 # Next, we create a ptychography model from the dataset
 model = CDTools.models.SimplePtycho.from_dataset(dataset)
 

examples/unified_modes.py

@@ -0,0 +1,39 @@
+from __future__ import division, print_function, absolute_import
+
+import CDTools
+from matplotlib import pyplot as plt
+import pickle
+
+filename = 'example_data/lab_ptycho_data.cxi'
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
+
+# dataset.inspect()
+# plt.show()
+
+#model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=2,n_modes=3)#, probe_support_radius=90)
+model = CDTools.models.UnifiedModePtycho2.from_dataset(dataset, oversampling=1,n_modes=3)#, probe_support_radius=90)
+#model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1)
+
+model.to(device='cuda')
+dataset.get_as(device='cuda')
+
+
+model.translation_offsets.requires_grad = False
+for i, loss in enumerate(model.Adam_optimize(100, dataset)):
+    model.inspect(dataset)
+    print(i,loss)
+
+model.tidy_probes()
+
+for i, loss in enumerate(model.Adam_optimize(20, dataset, lr=0.0001)):
+    model.inspect(dataset)
+    print(i,loss)
+
+model.tidy_probes()
+model.inspect(dataset)
+                         
+with open('example_reconstructions/unified_modes.pickle', 'wb') as f:
+    pickle.dump(model.save_results(dataset),f)
+
+model.compare(dataset)
+plt.show()

tests/tools/test_analysis.py

@@ -22,8 +22,8 @@ def test_orthogonalize_probes():
     probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
     probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
 
-    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2)),
-                       3*np.exp(-probe_Rs**2 / (2 * 12**2)),
+    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
+                       3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
                        1*np.exp(-probe_Rs**2 / (2 * 15**2))]).astype(np.complex64)
 
     # test that it works on numpy arrays
@@ -31,13 +31,15 @@ def test_orthogonalize_probes():
 
     # test that it also works on torch tensors
     ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes)))
-        
+
+    # This tests for orthogonality
     for p1,p2 in combinations(ortho_probes,2):
         assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
     for p1,p2 in combinations(ortho_probes_t,2):
         assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
+        
     probe_intensity = np.sum(np.abs(probes)**2,axis=0)
     ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0)
     ortho_probe_t_intensity = np.sum(np.abs(ortho_probes_t)**2,axis=0)
@@ -45,8 +47,35 @@ def test_orthogonalize_probes():
     assert np.allclose(probe_intensity,ortho_probe_intensity)
     assert np.allclose(probe_intensity,ortho_probe_t_intensity)
 
+    # Check that it returns normalized probes if we ask
+    ortho_probes = analysis.orthogonalize_probes(probes, normalize=True)
+    assert np.allclose([1,1,1],np.sum(np.abs(ortho_probes)**2,axis=(1,2)))
 
+    # And now we check that the A matrices actually work
+    ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
+                                                    normalize=False)
+    assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
 
+    # And now we check that the A matrices actually work
+    ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True,
+                                                    normalize=False)
+    assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes)
+
+    # The big problem here is that we haven't tested if it returns the
+    # probes or their complex conjugate.. we can test that by sending in
+    # one probe and checking that we get the right thing out
+    
+    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 20**2 + 100j))])
+    ortho_probes = analysis.orthogonalize_probes(probes)
+    # We need to correct for the arbitrary phase offset to be able to
+    # compare them
+    correction = ortho_probes[0]/probes[0]
+    correction = correction / np.abs(correction)
+    ortho_probes *= correction
+
+    assert np.allclose(ortho_probes,probes)
+
+        
 def test_standardize():
 
     # Start by making a probe and object that should meet the standardization

tests/tools/test_initializers.py

@@ -192,6 +192,32 @@ def test_SHARP_style_probe(ptycho_cxi_1):
 
 
 def test_RPI_spectral_init():
-    # Figure out a good way to test this
-    assert 0
-    RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None)
+    # I think we can only really meaningfully test that it doesn't throw errors,
+    # since the original implementation is in numpy and there aren't any clear
+    # cases that can be calculated analytically.
+
+    pattern = np.random.rand(230,253).astype(np.float32)
+    probe = np.random.rand(230,253).astype(np.complex64)
+    obj_shape = [37,53]
+    mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
+    background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05)
+
+    probe = cmath.complex_to_torch(probe)
+    pattern = t.Tensor(pattern)
+    
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
+    assert list(obj.shape) == [1]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, mask=mask)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, background=background)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+
+    obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
+                                         n_modes=2, mask=mask,
+                                         background=background)
+    assert list(obj.shape) == [2]+obj_shape+[2]
+

tests/tools/test_interactions.py

@@ -4,6 +4,8 @@ from CDTools.tools import cmath
 from CDTools.tools import interactions
 import numpy as np
 import torch as t
+from numpy import fft
+from scipy.fftpack import fftshift, ifftshift
 import pytest
 
 
@@ -98,8 +100,35 @@ def test_pixel_to_translations():
 
 
 def test_project_translations_to_sample():
-    # Needs to be tested
-    assert 0
+    # First, try the case where everything is ones and simple
+    basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
+    translations = t.rand((10,3))
+    pixels, props = interactions.project_translations_to_sample(basis, translations)
+
+    assert np.allclose(pixels[:,0].numpy(),-translations[:,1])
+    assert np.allclose(pixels[:,1].numpy(),-translations[:,0])
+    assert np.allclose(props.numpy(),-translations[:,2:].numpy())
+
+    # Next, a simple tilt along one axis. This is a 45 degree rotation
+    # around the positive y-axis
+    # Thus, y-axis translations are unaffected, but x-axis translations
+    # induce a motion of 1/sqrt(2) in the j- pixel space, as well as
+    # creating a propagation (negative propagation for positive x)
+    basis = t.Tensor([[0,-1e-3,0],[-np.sqrt(2)*1e-3,0,np.sqrt(2)*1e-3]]).t()
+    translations = t.rand((10,3))
+    pixels, props = interactions.project_translations_to_sample(basis, translations)
+
+    print(props.numpy())
+    print(-translations[:,2:].numpy() - translations[:,:1].numpy())
+    assert np.allclose(pixels[:,0].numpy(),-translations[:,1]*1e3)
+    assert np.allclose(pixels[:,1].numpy(),-translations[:,0]*1e3/np.sqrt(2))
+    assert np.allclose(props.numpy(),-translations[:,2:].numpy() - translations[:,:1].numpy())
+
+    # Finally, we check a non-orthogonal case
+
+
+    
+    
 
 def test_ptycho_2D_round(random_probe, random_obj):
     # Test a stack of images
@@ -227,22 +256,39 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
 
 def test_RPI_interaction(random_probe, random_obj):
 
-    random_obj1 = cmath.complex_to_torch(random_obj[:79,:68])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
 
+    random_obj1 = random_obj[:79,:68]
+    random_probe1 = random_probe
+    t_random_obj1 = cmath.complex_to_torch(random_obj1)
+    t_random_probe1 = cmath.complex_to_torch(random_probe1)
+    t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1)
 
-    random_obj1 = cmath.complex_to_torch(random_obj[:256,:256])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
+    obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
+    obj1_ups = np.zeros(random_probe1.shape[:2]).astype(np.complex128)
+    obj1_ups[(random_probe1.shape[0]-79)//2:
+             (random_probe1.shape[0]-79)//2 + 79,
+             (random_probe1.shape[1]-68)//2:
+             (random_probe1.shape[1]-68)//2 + 68] = obj1_fourier
+    output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
+                                                norm='ortho'))
+    
+    assert np.allclose(cmath.torch_to_complex(t_output1), output1)
 
     
-    random_obj1 = cmath.complex_to_torch(random_obj[:42,:103])
-    random_probe1 = cmath.complex_to_torch(random_probe)
-    output = interactions.RPI_interaction(random_probe1, random_obj1)
+    random_obj2 = np.stack([random_obj[:64,:89]]*3)
+    random_probe2 = random_probe[3:,5:]
+    t_random_obj2 = cmath.complex_to_torch(random_obj2)
+    t_random_probe2 = cmath.complex_to_torch(random_probe2)
+    t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
+
+    obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
+    obj2_ups = np.zeros((3,)+random_probe2.shape[:2]).astype(np.complex128)
+    obj2_ups[:,(random_probe2.shape[0]-64)//2:
+             (random_probe2.shape[0]-64)//2 + 64,
+             (random_probe2.shape[1]-89)//2:
+             (random_probe2.shape[1]-89)//2 + 89] = obj2_fourier
+    output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
+                                                norm='ortho'))
+
 
-    # Need to actually test against a numpy implementation
-    print(random_probe.shape)
-    print(random_obj.shape)
-    assert 0