Start documenting the models, and change the default FOV on fancy_ptycho reconstruction inspections

examples/gold_ball_ptycho.py

@@ -15,7 +15,7 @@ model.to(device='cuda')
 dataset.get_as(device='cuda')
 
 # Now, we run a short reconstruction from the dataset
-for loss in model.Adam_optimize(100, dataset, batch_size=50, schedule=True):
+for loss in model.Adam_optimize(10, dataset, batch_size=50, schedule=True):
     # And we liveplot the updates to the model as they happen
     print(model.report())
     model.inspect(dataset)

examples/multi_gpu_simple_ptycho.py

@@ -16,11 +16,6 @@ def demo():
     model = cdtools.models.SimplePtycho.from_dataset(dataset)
     
     
-    #class MyDataParallel(t.nn.DataParallel):
-    #    def __getattr__(self, name):
-    #        return getattr(self.module, name)
-    
-    
     print(dist.is_torchelastic_launched())
     dist.init_process_group(backend="nccl", init_method='env://')
     rank = dist.get_rank()
@@ -29,7 +24,9 @@ def demo():
     print('device id', device_id)
     model = model.to(device_id)
     ddp_model = DDP(model, device_ids=[device_id])
-    
+    print(model.loss_train)
+    print(model.Adam_optimize)
+    print(ddp_model.Adam_optimize)
     
     sampler = torchdata.distributed.DistributedSampler(dataset)
     # Make a dataloader
@@ -78,6 +75,8 @@ def demo():
         #loss = dist.all_reduce(loss)
         #print(it, 'time', time.time()-t0)
         #print(loss / normalization)
+        if rank==0:
+            model.inspect()
     return loss#.cpu().numpy()
 
 if __name__=='__main__':

examples/simple_ptycho.py

@@ -14,6 +14,8 @@ dataset.patterns = t.cat([dataset.patterns]*5)
 # Next, we create a ptychography model from the dataset
 model = cdtools.models.SimplePtycho.from_dataset(dataset)
 
+print(model.mask)
+exit()
 #class MyDataParallel(t.nn.DataParallel):
 #    def __getattr__(self, name):
 #        return getattr(self.module, name)

src/cdtools/models/base.py

@@ -37,7 +37,6 @@ import numpy as np
 import threading
 import queue
 import time
-from .complex_adam import MyAdam
 from .complex_lbfgs import MyLBFGS
 
 __all__ = ['CDIModel']
@@ -95,18 +94,52 @@ class CDIModel(t.nn.Module):
         raise NotImplementedError()
 
     
-    def store_detector_geometry(self, detector_geometry):
-        if 'distance' in detector_geometry:
+    def simulate_to_dataset(self, args_list):
+        raise NotImplementedError()
+
+    
+    def store_detector_geometry(self, detector_geometry, dtype=t.float32):
+        """Registers the information in a detector geometry dictionary
+
+        Information about the detector geometry is passed in as a dictionary,
+        but we want the various properties to be registered as buffers in
+        the model. This has nice effects, for example automatically updating
+        with model.to, and making it possible to automatically save them out.
+
+        Parameters
+        ----------
+        detector_geometry : dict
+            A dictionary containing at least the two entries 'distance' and 'basis'
+        dtype : torch.dtype, default: torch.float32
+            The datatype to convert the values to before registering
+        """
+        self.register_buffer('det_basis',
+                             t.tensor(detector_geometry['basis'],
+                                      dtype=dtype))
+        
+        if 'distance' in detector_geometry \
+           and detector_geometry['distance'] is not None:                
             self.register_buffer('det_distance',
-                                 t.as_tensor(detector_geometry['distance']))
-        if 'basis' in detector_geometry:
-            self.register_buffer('det_basis',
-                                 t.as_tensor(detector_geometry['basis']))
-        if 'corner' in detector_geometry:
+                                 t.tensor(detector_geometry['distance'],
+                                          dtype=dtype))
+        if 'corner' in detector_geometry \
+           and detector_geometry['corner'] is not None:
             self.register_buffer('det_corner',
-                                 t.as_tensor(detector_geometry['corner']))
+                                 t.tensor(detector_geometry['corner'],
+                                          dtype=dtype))
 
     def get_detector_geometry(self):
+        """Makes a detector geometry dictionary from the registered buffers
+
+        This extracts a dictionary with the detector geometry data from
+        the registered buffers, helpful for functions which expect the
+        geometry data to be in this format.
+
+        Returns
+        -------
+        detector_geometry : dict
+            A dictionary containing at least the two entries 'distance' and 'basis', pulled from the model's buffers
+        """
         detector_geometry = {}
         if hasattr(self, 'det_distance'):
             detector_geometry['distance'] = self.det_distance
@@ -116,9 +149,6 @@ class CDIModel(t.nn.Module):
             detector_geometry['corner'] = self.det_corner
         return detector_geometry    
     
-    def simulate_to_dataset(self, args_list):
-        raise NotImplementedError()
-
     def save_results(self):
         """A convenience function to get the state dict as numpy arrays
 
@@ -130,29 +160,40 @@ class CDIModel(t.nn.Module):
         results of the reconstruction
 
         Second, because display, further processing, long-term storage,
-        etc. are often done with dictionaries of numpy files, it's useful
+        etc. are often done with dictionaries of numpy arrays. So, it's useful
         to have a convenience function which does that conversion
         automatically.
+        
+        Returns
+        -------
+        results : dict
+            A dictionary containing all the parameters and buffers of the model, i.e. the result of self.state_dict(), converted to numpy.
         """
         return {k: v.cpu().numpy() for k, v in self.state_dict().items()}
 
+    
     def AD_optimize(self, iterations, data_loader,  optimizer,\
                     scheduler=None, regularization_factor=None, thread=True,
                     calculation_width=10):
         """Runs a round of reconstruction using the provided optimizer
 
         This is the basic automatic differentiation reconstruction tool
-        which all the other, algorithm-specific tools, use.
+        which all the other, algorithm-specific tools, use. It is a
+        generator which yields the average loss each epoch, ending after
+        the specified number of iterations.
 
-        Like all the other optimization routines, it is defined as a
-        generator function which yields the average loss each epoch.
+        By default, the computation will be run in a separate thread. This
+        is done to enable live plotting with matplotlib during a reconstruction.
+        If the computation was done in the main thread, this would freeze
+        the plots. This behavior can be turned off by setting the keyword
+        argument 'thread' to False.        
 
         Parameters
         ----------
         iterations : int
             How many epochs of the algorithm to run
-        dataset : CDataset
-            The dataset to reconstruct against
+        data_loader : torch.utils.data.DataLoader
+            A data loader loading the CDataset to reconstruct
         optimizer : torch.optim.Optimizer
             The optimizer to run the reconstruction with
         scheduler : torch.optim.lr_scheduler._LRScheduler
@@ -162,22 +203,34 @@ class CDIModel(t.nn.Module):
         thread : bool
             Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         calculation_width : int
-            Default 10, how many translations to pass through at once for each round of gradient accumulation
-        """
-        # First, calculate the normalization
-        normalization = 0
-        for inputs, patterns in data_loader:
-            normalization += t.sum(patterns).cpu().numpy()
+            Default 10, how many translations to pass through at once for each round of gradient accumulation. This does not affect the result, but may affect the calculation speed.
 
-        def run_iteration(stop_event=None):
+        Yields
+        ------
+        loss : float
+            The summed loss over the latest epoch, divided by the total diffraction pattern intensity
+        """
+
+        def run_epoch(stop_event=None):
+            """Runs one full epoch of the reconstruction."""
+            # First, initialize some tracking variables
+            normalization = 0
             loss = 0
             N = 0
             t0 = time.time()
+
+            # The data loader is responsible for setting the minibatch
+            # size, so each set is a minibatch
             for inputs, patterns in data_loader:
+                normalization += t.sum(patterns).cpu().numpy()
                 N += 1
                 def closure():
                     optimizer.zero_grad()
 
+                    # We further break up the minibatch into a set of chunks.
+                    # This lets us use larger minibatches than can fit
+                    # on the GPU at once, while still doing batch processing
+                    # for efficiency
                     input_chunks = [[inp[i:i + calculation_width]
                                      for inp in inputs]
                                     for i in range(0, len(inputs[0]),
@@ -188,31 +241,39 @@ class CDIModel(t.nn.Module):
 
                     total_loss = 0
                     for inp, pats in zip(input_chunks, pattern_chunks):
-                        # This is just used to allow graceful exit when
-                        # threading
+                        # This check allows for graceful exit when threading
                         if stop_event is not None and stop_event.is_set():
                             exit()
 
+                        # Run the simulation
                         sim_patterns = self.forward(*inp)
 
+                        # Calculate the loss
                         if hasattr(self, 'mask'):
                             loss = self.loss(pats,sim_patterns, mask=self.mask)
                         else:
                             loss = self.loss(pats,sim_patterns)
 
+                        # And accumulate the gradients
                         loss.backward()
-
                         total_loss += loss.detach()
 
+                    # If we have a regularizer, we can calculate it separately,
+                    # and the gradients will add to the minibatch gradient
                     if regularization_factor is not None \
                        and hasattr(self, 'regularizer'):
                         loss = self.regularizer(regularization_factor)
                         loss.backward()
+                    
                     return total_loss
 
+                # This takes the step for this minibatch
                 loss += optimizer.step(closure).detach().cpu().numpy()
 
+            
             loss /= normalization
+
+            # We step the scheduler after the full epoch
             if scheduler is not None:
                 scheduler.step(loss)
 
@@ -220,19 +281,26 @@ class CDIModel(t.nn.Module):
             self.latest_iteration_time = time.time() - t0
             return loss
 
-        if thread:
+        # If we don't want to run in a different thread, this is easy
+        if not thread:
+            for it in range(iterations):
+                yield run_iteration()
+
+        # But if we do want to thread, it's annoying:
+        else:
+            # Here we set up the communication with the computation thread
             result_queue = queue.Queue()
             stop_event = threading.Event()
             def target():
                 try:
-                    result_queue.put(run_iteration(stop_event))
+                    result_queue.put(run_epoch(stop_event))
                 except Exception as e:
                     # If something bad happens, put the exception into the
                     # result queue
                     result_queue.put(e)
 
-        for it in range(iterations):
-            if thread:
+            # And this actually starts and monitors the thread
+            for it in range(iterations):
                 calc = threading.Thread(target=target, name='calculator', daemon=True)
                 try:
                     calc.start()
@@ -256,9 +324,6 @@ class CDIModel(t.nn.Module):
 
                 yield res
 
-            else:
-                yield run_iteration()
-
 
     def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005,
                       schedule=False, amsgrad=False, subset=None,
@@ -280,7 +345,7 @@ class CDIModel(t.nn.Module):
         batch_size : int
             Optional, the size of the minibatches to use
         lr : float
-            Optional, The learning rate (alpha) to use
+            Optional, The learning rate (alpha) to use. 0.05 is typically the highest value with any chance of being stable
         schedule : float
             Optional, whether to use the ReduceLROnPlateau scheduler
         subset : list(int) or int
@@ -290,24 +355,23 @@ class CDIModel(t.nn.Module):
         thread : bool
             Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
         calculation_width : int
-            Default 1, how many translations to pass through at once for each round of gradient accumulation
+            Default 10, how many translations to pass through at once for each round of gradient accumulation. Does not affect the result, only the calculation speed
 
         """
 
         if subset is not None:
-            # if just one pattern, turn into a list for convenience
+            # if subset is just one pattern, turn into a list for convenience
             if type(subset) == type(1):
                 subset = [subset]
             dataset = torchdata.Subset(dataset, subset)
 
         # Make a dataloader
-        data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
+        data_loader = torchdata.DataLoader(dataset,
+                                           batch_size=batch_size,
                                            shuffle=True)
 
         # Define the optimizer
-        #optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
-        optimizer = MyAdam(self.parameters(), lr = lr, amsgrad=amsgrad)
-
+        optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
 
         # Define the scheduler
         if schedule:
@@ -351,7 +415,7 @@ class CDIModel(t.nn.Module):
         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
+            Default True, whether to run the computation in a separate thread to allow interaction with plots during computation.
 
         """
         if subset is not None:
@@ -382,7 +446,7 @@ class CDIModel(t.nn.Module):
                      lr=0.01, momentum=0, dampening=0, weight_decay=0,
                      nesterov=False, subset=None, regularization_factor=None,
                      thread=True, calculation_width=10):
-        """Runs a round of reconstruction using the SGDoptimizer
+        """Runs a round of reconstruction using the SGD optimizer
 
         This algorithm is often less stable that Adam, but it is simpler
         and is the basic workhorse of gradience descent.
@@ -438,7 +502,7 @@ class CDIModel(t.nn.Module):
 
 
     def report(self):
-        """Returns a string informing on the latest reconstruction iteration
+        """Returns a string with info about the latest reconstruction iteration
 
         Returns
         -------
@@ -463,7 +527,7 @@ class CDIModel(t.nn.Module):
         of plots, if one exists, and then redraw them. Otherwise, it will
         plot a new set, and any subsequent updates will update the new set
 
-        Optionally, a dataset can be passed, which then will plot any
+        Optionally, a dataset can be passed, which will allow plotting of any
         registered plots which need to incorporate some information from
         the dataset (such as geometry or a comparison with measured data).
 
@@ -479,22 +543,11 @@ class CDIModel(t.nn.Module):
         ----------
         dataset : CDataset
             Optional, a dataset matched to the model type
-        update : bool
-            Default True, whether to update existing plots or plot new ones
+        update : bool, default: True
+            Whether to update existing plots or plot new ones
 
         """
-
-        #print('base models inspect: checking the object')
-        #a = self.obj.detach()
-        #def saveobj(a, filename):
-        #    a = np.abs(a)
-        #    plt.imshow(a)
-        #    plt.savefig(filename)
-        #f = ['base_a.png', 'base_b.png', 'base_c.png', 'base_d.png']
-        #comp = [a[i, j, :, :] for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1])]
-        #for i in range(4):
-        #    saveobj(comp[i], f[i])
-        
+        # We find or create all the figures
         first_update = False
         if update and hasattr(self, 'figs') and self.figs:
             figs = self.figs
@@ -508,7 +561,8 @@ class CDIModel(t.nn.Module):
 
         idx = 0
         for plots in self.plot_list:
-            # If a conditional is included in the plot
+            # If a conditional is included in the plot, we check whether
+            # it is True
             try:
                 if len(plots) >=3 and not plots[2](self):
                     continue
@@ -525,33 +579,36 @@ class CDIModel(t.nn.Module):
             else:
                 fig = figs[idx]
 
-            try:
+                
+            try: # We try just plotting using the simplest allowed signature
                 plotter(self,fig)
                 plt.title(name)
-
             except TypeError as e:
+                # TypeError implies it wanted another argument, i.e. a dataset
                 if dataset is not None:
                     try:
                         plotter(self, fig, dataset)
                         plt.title(name)
-
-                    except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
+                    except Exception as e: # Don't raise errors: it's just plots
                         pass
 
-            except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
+            except Exception as e: # Don't raise errors, it's just a plot
                 pass
 
             idx += 1
 
             if update:
+                # This seems to update the figure without blocking.
                 plt.draw()
                 fig.canvas.start_event_loop(0.001)
 
         if first_update:
+            # But this is needed the first time the figures update, or
+            # they won't get drawn at all
             plt.pause(0.05 * len(self.figs))
 
 
-    def save_figures(self, prefix='', extension='.eps'):
+    def save_figures(self, prefix='', extension='.pdf'):
         """Saves all currently open inspection figures.
 
         Note that this function is not very intelligent - so, for example,
@@ -586,10 +643,18 @@ class CDIModel(t.nn.Module):
     def compare(self, dataset, logarithmic=False):
         """Opens a tool for comparing simulated and measured diffraction patterns
 
+        This does what it says on the tin.
+        
+        Also, I am very sorry, the implementation was done while I was
+        possessed by Beezlebub - do not try to fix this, if it breaks just
+        kill it and start from scratch.
+
         Parameters
         ----------
         dataset : CDataset
             A dataset containing the simulated diffraction patterns to compare against
+        logarithmic : bool, default: False
+            Whether to plot the diffraction on a logarithmic scale
         """
 
         fig, axes = plt.subplots(1,3,figsize=(12,5.3))
@@ -598,13 +663,6 @@ class CDIModel(t.nn.Module):
 
 
         def update_colorbar(im):
-            # If the update brought the colorbar out of whack
-            # (say, from clicking back in the navbar)
-            # Holy fuck this was annoying. Sorry future for how
-            # crappy this solution is.
-            #if not np.allclose(im.colorbar.ax.get_xlim(),
-            #                   (np.min(im.get_array()),
-            #                    np.max(im.get_array()))):
             if hasattr(im, 'norecurse') and im.norecurse:
                 im.norecurse=False
                 return
@@ -617,9 +675,9 @@ class CDIModel(t.nn.Module):
             fig.pattern_idx = idx
             updating = True if len(axes[0].images) >= 1 else False
 
-            inputs, output = dataset[idx]
-            sim_data = self.forward(*inputs).detach().cpu().numpy()
-            meas_data = output.detach().cpu().numpy()
+            inputs, output = dataset[idx:idx+1]
+            sim_data = self.forward(*inputs).detach().cpu().numpy()[0]
+            meas_data = output.detach().cpu().numpy()[0]
             if hasattr(self, 'mask') and self.mask is not None:
                 mask = self.mask.detach().cpu().numpy()
             else:

src/cdtools/models/bragg_2d_ptycho.py

@@ -72,7 +72,7 @@ class Bragg2DPtycho(CDIModel):
                  background=None, translation_offsets=None, mask=None,
                  weights=None, translation_scale=1, saturation=None,
                  probe_support=None, oversampling=1,
-                 propagate_probe=True, correct_tilt=True, lens=False):
+                 propagate_probe=True, correct_tilt=True, lens=False, units='um'):
 
         # We need the detector geometry
         # We need the probe basis (but in this case, we don't need the surface
@@ -84,6 +84,7 @@ class Bragg2DPtycho(CDIModel):
         # propagate_probe and correct_tilt are important!
 
         super(Bragg2DPtycho, self).__init__()
+        self.units = units
         self.wavelength = t.tensor(wavelength)
         self.detector_geometry = copy(detector_geometry)
         det_geo = self.detector_geometry
@@ -375,14 +376,15 @@ class Bragg2DPtycho(CDIModel):
         prs = Ws[...,None,None,None] * self.probe * self.probe_support[...,:,:]
         # Now we need to propagate each of the probes
 
-        
-        for j in range(prs.shape[0]):
-            # I believe this -1 sign is in error, but I need a dataset with
-            # well understood geometry to figure it out
-            propagator = t.exp(
-                1j*(props[j]*(2*np.pi)/self.wavelength)
-                * self.universal_propagator)
-            prs[j] = tools.propagators.near_field(prs[j], propagator)
+
+        if self.propagate_probe:
+            for j in range(prs.shape[0]):
+                # I believe this -1 sign is in error, but I need a dataset with
+                # well understood geometry to figure it out
+                propagator = t.exp(
+                    1j*(props[j]*(2*np.pi)/self.wavelength)
+                    * self.universal_propagator)
+                prs[j] = tools.propagators.near_field(prs[j], propagator)
 
         exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
             prs, self.obj,pix_trans,
@@ -495,20 +497,18 @@ class Bragg2DPtycho(CDIModel):
 
 
     plot_list = [
-        ('Dominant Probe Amplitude',
-         lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)),
-        ('Dominant Probe Phase',
-         lambda self, fig: p.plot_phase(self.probe[0], fig=fig)),
-        ('Subdominant Probe Amplitude',
-         lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig),
-         lambda self: len(self.probe) >=2),
-        ('Subdominant Probe Phase',
-         lambda self, fig: p.plot_phase(self.probe[1], fig=fig),
-         lambda self: len(self.probe) >=2),
+        ('Basis Probe Fourier Space Amplitudes',
+         lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Basis Probe Fourier Space Phases',
+         lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)),
+        ('Basis Probe Real Space Amplitudes',
+         lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
+        ('Basis Probe Real Space Phases',
+         lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
         ('Object Amplitude', 
-         lambda self, fig: p.plot_amplitude(self.obj, fig=fig)),
+         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)),
+         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',

src/cdtools/models/complex_adam.py

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

src/cdtools/models/fancy_ptycho.py

@@ -38,33 +38,40 @@ class FancyPtycho(CDIModel):
                  units='um',
                  simulate_probe_translation=False,
                  simulate_finite_pixels=False,
+                 dtype=t.float32,
+                 obj_view_crop=0
                  ):
 
         super(FancyPtycho, self).__init__()
-        self.wavelength = t.tensor(wavelength)
-        self.detector_geometry = copy(detector_geometry)
-        det_geo = self.detector_geometry
-        if 'distance' in det_geo:
-            det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
-        if 'basis' in det_geo:
-            det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
-        if 'corner' in det_geo and det_geo['corner'] is not None:
-            det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
+        self.register_buffer('wavelength',
+                             t.tensor(wavelength, dtype=dtype))
+        self.store_detector_geometry(detector_geometry,
+                                     dtype=dtype)
 
-        self.min_translation = t.tensor(min_translation)
+        self.register_buffer('min_translation',
+                             t.tensor(min_translation, dtype=dtype))
 
-        self.probe_basis = t.tensor(probe_basis)
+        self.register_buffer('probe_basis',
+                             t.tensor(probe_basis, dtype=dtype))
+        
         self.detector_slice = copy(detector_slice)
-        self.surface_normal = t.tensor(surface_normal)
-
-        self.saturation = saturation
+        self.register_buffer('surface_normal',
+                             t.tensor(surface_normal, dtype=dtype))
+        if saturation is None:
+            self.saturation = None
+        else:
+            self.register_buffer('saturation',
+                                 t.tensor(saturation, dtype=dtype))
+        # Not sure how to make this a buffer...
         self.units = units
+
         self.fourier_probe = fourier_probe
 
         if mask is None:
-            self.mask = mask
+            self.mask = None
         else:
-            self.mask = t.tensor(mask, dtype=t.bool)
+            self.register_buffer('mask',
+                                 t.tensor(mask, dtype=t.bool))
         
         probe_guess = t.tensor(probe_guess, dtype=t.complex64)
         obj_guess = t.tensor(obj_guess, dtype=t.complex64)
@@ -72,13 +79,18 @@ class FancyPtycho(CDIModel):
         # We rescale the probe here so it learns at the same rate as the
         # object
         if probe_guess.dim() > 2:
-            self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
+            probe_norm = 1 * t.max(t.abs(probe_guess[0]))
         else:
-            self.probe_norm = 1 * t.max(t.abs(probe_guess))        
+            probe_norm = 1 * t.max(t.abs(probe_guess))
+        self.register_buffer('probe_norm', probe_norm.to(dtype))
         
         self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
         self.obj = t.nn.Parameter(obj_guess)
 
+
+        self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop,
+                                    obj_view_crop:-obj_view_crop]
+        
         if background is None:
             if detector_slice is not None:
                 dummy_det = t.empty([s//oversampling
@@ -112,22 +124,28 @@ class FancyPtycho(CDIModel):
             t_o = t_o / translation_scale
             self.translation_offsets = t.nn.Parameter(t_o)
 
-        self.translation_scale = translation_scale
-
-        if probe_support is not None:
-            self.probe_support = probe_support
-        else:
-            self.probe_support = t.ones_like(self.probe[0], dtype=t.bool)
+        self.register_buffer('translation_scale',
+                             t.tensor(translation_scale, dtype=dtype))
 
+        if probe_support is None:
+            probe_support = t.ones_like(self.probe[0], dtype=t.bool)
+        self.register_buffer('probe_support',
+                             t.tensor(probe_support, dtype=t.bool))
+            
         self.oversampling = oversampling
 
         self.simulate_probe_translation = simulate_probe_translation
+
         if simulate_probe_translation:
-            Is = t.arange(self.probe.shape[-2], dtype=t.float32)
-            Js = t.arange(self.probe.shape[-1], dtype=t.float32)
+            Is = t.arange(self.probe.shape[-2], dtype=dtype)
+            Js = t.arange(self.probe.shape[-1], dtype=dtype)
             Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
-            self.I_phase = 2 * np.pi* Is * self.oversampling
-            self.J_phase = 2 * np.pi* Js * self.oversampling
+            
+            I_phase = 2 * np.pi* Is * self.oversampling
+            J_phase = 2 * np.pi* Js * self.oversampling
+            self.register_buffer('I_phase', I_phase)
+            self.register_buffer('J_phase', J_phase)
+            
 
         self.simulate_finite_pixels = simulate_finite_pixels
             
@@ -165,6 +183,8 @@ class FancyPtycho(CDIModel):
                      units='um',
                      simulate_probe_translation=False,
                      simulate_finite_pixels=False,
+                     obj_view_crop=None,
+                     obj_padding=200
                      ):
 
         wavelength = dataset.wavelength
@@ -221,7 +241,7 @@ class FancyPtycho(CDIModel):
 
         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)
+        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=obj_padding)
 
         if hasattr(dataset, 'background') and dataset.background is not None:
             background = t.sqrt(dataset.background)
@@ -254,6 +274,12 @@ class FancyPtycho(CDIModel):
         if n_obj_modes != 1:
             obj = t.stack([obj,] + [0.05*t.ones_like(obj),]*(n_obj_modes-1))
 
+        if obj_view_crop is None:
+            obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2
+        if obj_view_crop < 0:
+            obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2
+        obj_view_crop += obj_padding
+            
         det_geo = dataset.detector_geometry
 
         translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
@@ -312,7 +338,8 @@ class FancyPtycho(CDIModel):
                    oversampling=oversampling,
                    loss=loss, units=units,
                    simulate_probe_translation=simulate_probe_translation,
-                   simulate_finite_pixels=simulate_finite_pixels)
+                   simulate_finite_pixels=simulate_finite_pixels,
+                   obj_view_crop=obj_view_crop)
 
 
     def interaction(self, index, translations, *args):
@@ -338,7 +365,7 @@ class FancyPtycho(CDIModel):
         # For a Fourier-space probe, we take an IFT
         if self.fourier_probe:
             basis_prs = tools.propagators.inverse_far_field(basis_prs)
-
+            
         # Now we construct the probes for each shot from the basis probes
         if self.weights is not None:
             Ws = self.weights[index]
@@ -358,10 +385,10 @@ class FancyPtycho(CDIModel):
             # coherent mode index, then x,y, and then complex index
             # Maybe this can be done with a matmul now?
             prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3)
-
+        
         if self.simulate_probe_translation:
             det_pix_trans = tools.interactions.translations_to_pixel(
-                    self.detector_geometry['basis'],
+                    self.det_basis,
                     translations,
                     surface_normal=self.surface_normal)
             
@@ -375,8 +402,8 @@ class FancyPtycho(CDIModel):
         # which allows us to do stuff like let the object be super-resolution,
         # while restricting the probe to the detector resolution but still
         # doing an explicit real-space limitation of the probe
-        padding = [self.background.shape[-2] - prs.shape[-2],
-                   self.background.shape[-1] - prs.shape[-1]]
+        padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2],
+                   self.oversampling * self.background.shape[-1] - prs.shape[-1]]
 
         if any([p != 0 for p in padding]): # For probe_fourier_crop != 0.
             padding = [padding[-1]//2, padding[-1]-padding[-1]//2,
@@ -415,32 +442,6 @@ class FancyPtycho(CDIModel):
     # Note: No "loss" function is defined here, because it is added
     # dynamically during object creation in __init__
 
-    def to(self, *args, **kwargs):
-        super(FancyPtycho, self).to(*args, **kwargs)
-        self.wavelength = self.wavelength.to(*args, **kwargs)
-        # move the detector geometry too
-        det_geo = self.detector_geometry
-        if 'distance' in det_geo:
-            det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
-        if 'basis' in det_geo:
-            det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
-        if 'corner' in det_geo and det_geo['corner'] is not None:
-            det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
-
-        if self.mask is not None:
-            self.mask = self.mask.to(*args, **kwargs)
-
-        if self.simulate_probe_translation:
-            self.I_phase = self.I_phase.to(*args, **kwargs)
-            self.J_phase = self.J_phase.to(*args, **kwargs)
-
-        self.min_translation = self.min_translation.to(*args, **kwargs)
-        self.probe_basis = self.probe_basis.to(*args, **kwargs)
-        self.probe_norm = self.probe_norm.to(*args, **kwargs)
-        self.probe_support = self.probe_support.to(*args, **kwargs)
-        self.surface_normal = self.surface_normal.to(*args, **kwargs)
-
-
     def sim_to_dataset(self, args_list, calculation_width=None):
         # In the future, potentially add more control
         # over what metadata is saved (names, etc.)
@@ -462,7 +463,6 @@ class FancyPtycho(CDIModel):
                        'orientation': orientation}
 
 
-        detector_geometry = self.detector_geometry
         mask = self.mask
         wavelength = self.wavelength
         indices, translations = args_list
@@ -490,7 +490,7 @@ class FancyPtycho(CDIModel):
             entry_info=entry_info,
             sample_info=sample_info,
             wavelength=wavelength,
-            detector_geometry=detector_geometry,
+            detector_geometry=self.get_detector_geometry(),
             mask=mask)
 
 
@@ -629,32 +629,77 @@ class FancyPtycho(CDIModel):
         
     plot_list = [
         ('',
-         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'),
+         lambda self, fig, dataset: self.plot_wavefront_variation(
+             dataset,
+             fig=fig,
+             mode='root_sum_intensity',
+             image_title='Root Summed Probe Intensities',
+             image_colorbar_title='Square Root of Intensity'),
          lambda self: len(self.weights.shape) >= 2),
         ('',
-         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'),
+         lambda self, fig, dataset: self.plot_wavefront_variation(
+             dataset,
+             fig=fig,
+             mode='amplitude',
+             image_title='Probe Amplitudes (scroll to view modes)',
+             image_colorbar_title='Probe Amplitude'),
          lambda self: len(self.weights.shape) >= 2),
         ('',
-         lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'),
+         lambda self, fig, dataset: self.plot_wavefront_variation(
+             dataset,
+             fig=fig,
+             mode='phase',
+             image_title='Probe Phases (scroll to view modes)',
+             image_colorbar_title='Probe Phase'),
          lambda self: len(self.weights.shape) >= 2),
         ('Basis Probe Fourier Space Amplitudes',
-         lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+         lambda self, fig: p.plot_amplitude(
+             (self.probe if self.fourier_probe
+              else tools.propagators.inverse_far_field(self.probe)),
+              fig=fig)),
         ('Basis Probe Fourier Space Phases',
-         lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
+         lambda self, fig: p.plot_phase(
+             (self.probe if self.fourier_probe
+              else tools.propagators.inverse_far_field(self.probe))
+             , fig=fig)),
         ('Basis Probe Real Space Amplitudes',
-         lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+         lambda self, fig: p.plot_amplitude(
+             (self.probe if not self.fourier_probe
+              else tools.propagators.inverse_far_field(self.probe)),
+             fig=fig,
+             basis=self.probe_basis,
+             units=self.units)),
         ('Basis Probe Real Space Phases',
-         lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
+         lambda self, fig: p.plot_phase(
+             (self.probe if not self.fourier_probe
+              else tools.propagators.inverse_far_field(self.probe)),
+             fig=fig,
+             basis=self.probe_basis,
+             units=self.units)),
         ('Average Density Matrix Amplitudes',
-         lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
+         lambda self, fig: p.plot_amplitude(
+             np.nanmean(np.abs(self.get_rhos()), axis=0),
+             fig=fig),
          lambda self: len(self.weights.shape) >= 2),
         ('% Power in Top Mode (only accurate after tidy_probes)',
-         lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
+         lambda self, fig, dataset: p.plot_nanomap(
+             self.corrected_translations(dataset),
+             analysis.calc_top_mode_fraction(self.get_rhos()),
+             fig=fig,
+             units=self.units),
          lambda self: len(self.weights.shape) >= 2),
         ('Object Amplitude',
-         lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+         lambda self, fig: p.plot_amplitude(
+             self.obj[self.obj_view_slice],
+             fig=fig,
+             basis=self.probe_basis,
+             units=self.units)),
         ('Object Phase',
-         lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
+         lambda self, fig: p.plot_phase(
+             self.obj[self.obj_view_slice],
+             fig=fig,
+             basis=self.probe_basis,
+             units=self.units)),
         ('Corrected Translations',
          lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
         ('Background',

src/cdtools/models/simple_ptycho.py

@@ -45,22 +45,9 @@ class SimplePtycho(CDIModel):
         # object
         self.register_buffer('probe_norm', t.max(t.abs(probe_guess)))
 
-        #self.probe_data = complexWrapper(probe_guess/self.probe_norm)
-        self.probe_data = t.nn.Parameter(t.view_as_real(probe_guess / self.probe_norm))
-        self.obj_data = t.nn.Parameter(t.view_as_real(obj_guess))
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+        self.obj = t.nn.Parameter(obj_guess)
 
-    # Do this for all the complex-valued parameters which are stored as real
-    # valued for compatibility reasons. Note that updating model.probe.data
-    # won't update the underlying storage, but something like:
-    # model.probe.data[:] = ...
-    # will update the underlying storage
-    @property
-    def probe(self):
-        return t.view_as_complex(self.probe_data)
-
-    @property
-    def obj(self):
-        return t.view_as_complex(self.obj_data)
 
     @classmethod
     def from_dataset(cls, dataset):
@@ -94,8 +81,13 @@ class SimplePtycho(CDIModel):
 
         # 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)
+        pix_translations = tools.interactions.translations_to_pixel(
+            probe_basis,
+            translations,
+            surface_normal=surface_normal)
+        obj_size, min_translation = tools.initializers.calc_object_setup(
+            probe_shape,
+            pix_translations)
 
         # Finally, initialize the probe and  object using this information
         probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
@@ -109,17 +101,28 @@ class SimplePtycho(CDIModel):
         else:
             mask = None
 
-        return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask, surface_normal=surface_normal)
+        return cls(wavelength,
+                   det_geo,
+                   probe_basis,
+                   det_slice,
+                   probe,
+                   obj,
+                   min_translation=min_translation,
+                   mask=mask,
+                   surface_normal=surface_normal)
 
 
     def interaction(self, index, translations):
-        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
-                                                             translations,
-                                            surface_normal=self.surface_normal)
+        pix_trans = tools.interactions.translations_to_pixel(
+            self.probe_basis,
+            translations,
+            surface_normal=self.surface_normal)
+        
         pix_trans -= self.min_translation
-        return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
-                                                  self.obj,
-                                                  pix_trans)
+        return tools.interactions.ptycho_2D_round(
+            self.probe_norm * self.probe,
+            self.obj,
+            pix_trans)
 
 
     def forward_propagator(self, wavefields):
@@ -133,50 +136,12 @@ class SimplePtycho(CDIModel):
     def measurement(self, wavefields):
         return tools.measurements.intensity(wavefields,
                                             detector_slice=self.detector_slice)
-
+    
 
     def loss(self, real_data, sim_data, mask=None):
         return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
 
 
-    def sim_to_dataset(self, args_list):
-        # In the future, potentially add more control
-        # over what metadata is saved (names, etc.)
-
-        # First, I need to gather all the relevant data
-        # that needs to be added to the dataset
-        entry_info = {'program_name': 'cdtools',
-                      'instrument_n': 'Simulated Data',
-                      'start_time': datetime.now()}
-
-        surface_normal = self.surface_normal.detach().cpu().numpy()
-        xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
-        xsurfacevec /= np.linalg.norm(xsurfacevec)
-        ysurfacevec = np.cross(surface_normal, xsurfacevec)
-        ysurfacevec /= np.linalg.norm(ysurfacevec)
-        orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
-
-        sample_info = {'description': 'A simulated sample',
-                       'orientation': orientation}
-
-        detector_geometry = self.detector_geometry
-        mask = self.mask
-        wavelength = self.wavelength
-        indices, translations = args_list
-
-        # Then we simulate the results
-        data = self.forward(indices, translations)
-
-        # And finally, we make the dataset
-        return Ptycho2DDataset(translations, data,
-                                 entry_info = entry_info,
-                                 sample_info = sample_info,
-                                 wavelength=wavelength,
-                                 detector_geometry=detector_geometry,
-                                 mask=mask)
-
-
-
     plot_list = [
         ('Probe Amplitude',
          lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
@@ -188,65 +153,3 @@ class SimplePtycho(CDIModel):
          lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
     ]
     
-
-    def ePIE(self, iterations, dataset, beta = 1.0):
-        """Runs an ePIE reconstruction as described in `Maiden et al. (2017) <https://www.osapublishing.org/optica/abstract.cfm?uri=optica-4-7-736>`_.
-        Optional parameters are:
-
-        :arg ``iterations``: Controls the number of iterations run, defaults to 1.
-        :arg ``beta``: Algorithmic parameter described in Maiden's implementation of rPIE. Defaults to 0.15.
-        :arg ``probe``: Initial probe wavefunction.
-        :arg ``object``: Initial object wavefunction.
-        """
-        probe_shape = self.probe.shape
-
-        if self.mask is not None:
-            mask = self.mask[...,None]
-        else:
-            mask=None
-
-        def probe_update(exit_wave, exit_wave_corrected, probe, object, translation):
-            new_probe = probe + tools.cmath.cmult(beta * tools.cmath.cconj(object[translation])/(self.probe_norm*t.max(tools.cmath.cabssq(object))), exit_wave_corrected-exit_wave)
-            return new_probe
-
-        def object_update(exit_wave, exit_wave_corrected, probe, object, translation):
-            new_object = object.clone()
-            new_object[translation] = object[translation] + tools.cmath.cmult(beta * tools.cmath.cconj(probe)/(self.probe_norm*t.max(tools.cmath.cabssq(probe))), exit_wave_corrected-exit_wave)
-            return new_object
-
-        with t.no_grad():
-            data_loader = torchdata.DataLoader(dataset, shuffle=True)
-
-            for it in range(iterations):
-                loss = []
-                for (i, [translations]), [patterns] in data_loader:
-                    probe = self.probe.data.clone()
-                    object = self.obj.data.clone()
-
-                    exit_wave = self.interaction(i, translations).clone()
-                    # Apply modulus constraint
-                    exit_wave_corrected = exit_wave.clone()
-                    exit_wave_corrected = self.forward_propagator(exit_wave_corrected.clone())
-                    exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(exit_wave_corrected.clone()[self.detector_slice], patterns, mask = mask)
-                    exit_wave_corrected = self.backward_propagator(exit_wave_corrected.clone())
-
-                    # Calculate the section of the object wavefunction to be modified
-                    pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
-                                                                         translations)
-                    pix_trans -= self.min_translation
-
-                    pix_trans = t.round(pix_trans).to(dtype=t.int32).detach().cpu().numpy()
-
-                    object_slice = np.s_[pix_trans[0]:
-                                      pix_trans[0]+probe_shape[0],
-                                      pix_trans[1]:
-                                      pix_trans[1]+probe_shape[1]]
-
-                    # Apply probe and object updates
-                    self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
-                    self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
-
-                    # Calculate loss
-                    loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
-
-                yield t.mean(t.tensor(loss)).cpu().numpy()

src/cdtools/tools/interactions/interactions.py

@@ -95,8 +95,8 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
         The real space basis the wavefields are defined in
     translations : torch.Tensor
         A Jx2 stack of pixel-space translations, or a single translation
-    surface_normal : torch.Tensor
-        Optional, the sample's surface normal
+    surface_normal : torch.Tensor, default: torch.Tensor([0,0,1])
+        The sample's surface normal
 
     Returns
     -------

src/cdtools/tools/measurements/measurements.py

@@ -148,7 +148,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
             output = output[detector_slice]
         else:
             output = output[(np.s_[:],) + detector_slice]
-
+        
     if saturation is None:
         return output + epsilon
     else:
@@ -193,7 +193,7 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
         output = measurement(wavefield, *args, detector_slice=detector_slice,
                              epsilon=epsilon, oversampling=oversampling,
                              simulate_finite_pixels=simulate_finite_pixels) \
-                             + background**2
+                            + background**2
 
     # This has to be done after the background is added, hence we replicate
     # it here