From cb5942c39acec7c61f89c9c9ab66d3f888bcfacc Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Wed, 22 Jan 2025 16:41:18 +0100
Subject: [PATCH 01/55] Fix an issue for datasets with colinear translations

---
 src/cdtools/datasets/ptycho_2d_dataset.py |  5 ++++-
 src/cdtools/tools/analysis/analysis.py    | 24 ++++++++++++++++++-----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 8b6f6d8..971baf1 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -269,9 +269,12 @@ class Ptycho2DDataset(CDataset):
         """Plots the mean diffraction pattern across the dataset
 
         The output is normalized so that the summed intensity on the
-        detector is equal to the total intensity of light that passed
+        detector is roughly equal to the total intensity of light that passed
         through the sample within each detector conjugate field of view.
 
+        If the scan points are colinear (which causes issues for this
+        estimation), the mean pattern is displayed unscaled.
+
         The plot is plotted as log base 10 of the output plus log_offset.
         By default, log_offset is set equal to 1, which is a good level for
         shot-noise limited data captured in units of photons. More
diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py
index 0fd913e..9844edb 100644
--- a/src/cdtools/tools/analysis/analysis.py
+++ b/src/cdtools/tools/analysis/analysis.py
@@ -14,6 +14,7 @@ from scipy import linalg as sla
 from scipy import special
 from scipy import optimize as opt
 from scipy import spatial
+import warnings
 
 __all__ = [
     'product_svd',
@@ -1443,6 +1444,13 @@ def calc_spectral_info(dataset, nbins=50):
     the scan pattern whose area matches one detector conjugate field of
     view.
 
+    This estimation will start to deviate from the truth if the scan area
+    is not significantly larger than the illumination function, because
+    the nonzero size of the illumination function is not taken into account.
+    Furthermore, in the edge case where all the scan points are colinear,
+    the estimate will fail, and the mean diffraction pattern will be returned
+    instead
+
     Parameters
     ----------
     dataset : Ptycho2DDataset
@@ -1461,10 +1469,12 @@ def calc_spectral_info(dataset, nbins=50):
     
     """
 
-    scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
-    
-    scan_area = scan_hull.volume
-
+    try:
+        scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
+        scan_area = scan_hull.volume
+    except spatial._qhull.QhullError as e:
+        scan_area = None
+        
     ewg = cdtools.tools.initializers.exit_wave_geometry
     obj_basis = ewg(
         dataset.detector_geometry['basis'],
@@ -1477,8 +1487,12 @@ def calc_spectral_info(dataset, nbins=50):
         np.cross(obj_basis[:,0]*dataset.patterns.shape[-2],
                  obj_basis[:,1]*dataset.patterns.shape[-1])
     )
-    scale_factor = det_conj_fov_area / scan_area
 
+    if scan_area is not None:
+        scale_factor = det_conj_fov_area / scan_area
+    else:
+        warnings.warn("The scan points in this dataset are all colinear. The mean pattern will be calculated rather than a scaled mean based on the scanned area.")
+        scale_factor = 1/len(dataset)
 
     mask = dataset.mask.cpu().numpy().astype(int)
     sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor

From 91717b65d6ab78b3175b5ec6512c24a3c4aeea3b Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 7 Apr 2025 09:35:59 +0200
Subject: [PATCH 02/55] Change the behavior of FancyPtycho.center_probes to
 center the probe in real space even when fourier_probe is set to True

---
 src/cdtools/models/fancy_ptycho.py | 35 +++++++++++++++++++++++----
 tests/models/test_fancy_ptycho.py  | 39 ++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 1be40c6..8854d41 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -597,15 +597,40 @@ class FancyPtycho(CDIModel):
 
         
     def center_probes(self, iterations=4):
-        """Centers the probes
+        """Centers the probes in real space
+
+        Takes the current guess of the illumination function and centers it
+        using a shift with periodic boundary conditions. It uses
+        cdtools.tools.image_processing.center internally to do the centering.
+        Multiple iterations of an algorithm are run, which is helpful if the
+        illumination is reconstructed near the corners and "wraps around" the
+        probe field of view.
+
+        Note that the centering is always performed in real space, even if
+        the probe array is defined in Fourier space.
         
-        Note that this does not compensate for the centering by adjusting
+        Note also that this does not compensate for the centering by adjusting
         the object, so it's a good idea to reset the object after centering
         the probes
+
+        Parameters
+        ----------
+        iterations : int
+            Default 4, how many iterations of the centering algorithm to run
         """
-        centered_probe = tools.image_processing.center(
-            self.probe.data.cpu(), iterations=iterations)
-        self.probe.data = centered_probe.to(device=self.probe.data.device)
+        if self.fourier_probe:
+            prs = tools.propagators.inverse_far_field(self.probe.detach()).cpu()
+        else:
+            prs = self.probe.detach().cpu()
+        
+        centered_prs = tools.image_processing.center(prs, iterations=iterations)
+
+        if self.fourier_probe:
+            self.probe.data = tools.propagators.far_field(
+                centered_prs.to(device=self.probe.data.device))
+        else:
+            self.probe.data = centered_prs.to(device=self.probe.data.device)
+
 
 
     def tidy_probes(self):
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index c9eaa1b..4bcbac1 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -5,6 +5,45 @@ import torch as t
 import cdtools
 from matplotlib import pyplot as plt
 
+
+def test_center_probe(lab_ptycho_cxi):
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
+    model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=False
+    )
+    base_probe = model.probe.detach().clone()
+    model.center_probes()
+    centered_probe = model.probe.detach().clone()
+
+    fourier_model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=True,
+    )
+
+    fourier_model.probe.data = cdtools.tools.propagators.far_field(
+        base_probe
+    )
+    
+    fourier_base_probe = fourier_model.probe.detach().clone()
+    fourier_model.center_probes()
+    fourier_centered_probe = fourier_model.probe.detach().clone()
+    ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
+        fourier_centered_probe)
+
+    # So we know the code had to do something
+    assert not t.allclose(base_probe, centered_probe)
+    # And checking that they both do the same thing, whether or not
+    # fourier_probe was set to True
+    assert t.allclose(
+        centered_probe,
+        ifft_fourier_centered_probe,
+        atol=1e-4,
+        rtol=1e-3
+    )
+    
 @pytest.mark.slow
 def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 

From ac1499f511ee494a2a32d2673df0f82630819c82 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 7 Apr 2025 11:28:15 +0200
Subject: [PATCH 03/55] get it running in principle, but it is not good

---
 src/cdtools/models/fancy_ptycho.py            | 68 +++++++++++++------
 .../tools/measurements/measurements.py        | 47 +++++++++----
 2 files changed, 83 insertions(+), 32 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 1be40c6..9459855 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -28,6 +28,7 @@ class FancyPtycho(CDIModel):
                  probe_fourier_shifts=None,
                  mask=None,
                  weights=None,
+                 qe_mask=None,
                  translation_scale=1,
                  saturation=None,
                  probe_support=None,
@@ -87,7 +88,14 @@ class FancyPtycho(CDIModel):
         else:
             self.register_buffer('mask',
                                  t.as_tensor(mask, dtype=t.bool))
-        
+
+
+        if qe_mask is None:
+            self.qe_mask = None
+        else:
+            self.qe_mask = t.nn.Parameter(
+                t.as_tensor(qe_mask, dtype=dtype))
+            
         probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
         obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
 
@@ -202,6 +210,7 @@ class FancyPtycho(CDIModel):
                      dm_rank=None,
                      translation_scale=1,
                      saturation=None,
+                     use_qe_mask=False,
                      probe_support_radius=None,
                      probe_fourier_crop=None,
                      propagation_distance=None,
@@ -376,6 +385,11 @@ class FancyPtycho(CDIModel):
         else:
             mask = None
 
+        if use_qe_mask:
+            qe_mask = t.ones(mask.shape, dtype=t.float32)
+        else:
+            qe_mask = None
+            
         if probe_support_radius is not None:
             probe_support = t.zeros(probe[0].shape, dtype=t.bool)
             xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
@@ -389,24 +403,34 @@ class FancyPtycho(CDIModel):
         else:
             probe_support = None
 
-        return cls(wavelength, det_geo, obj_basis, probe, obj,
-                   surface_normal=surface_normal,
-                   min_translation=min_translation,
-                   translation_offsets=translation_offsets,
-                   weights=Ws, mask=mask, background=background,
-                   translation_scale=translation_scale,
-                   saturation=saturation,
-                   probe_basis=probe_basis,
-                   probe_support=probe_support,
-                   fourier_probe=fourier_probe,
-                   oversampling=oversampling,
-                   loss=loss, units=units,
-                   probe_fourier_shifts=probe_fourier_shifts,
-                   simulate_probe_translation=simulate_probe_translation,
-                   simulate_finite_pixels=simulate_finite_pixels,
-                   phase_only=phase_only,
-                   exponentiate_obj=exponentiate_obj,
-                   obj_view_crop=obj_view_crop)
+        return cls(
+            wavelength,
+            det_geo,
+            obj_basis,
+            probe,
+            obj,
+            surface_normal=surface_normal,
+            min_translation=min_translation,
+            translation_offsets=translation_offsets,
+            weights=Ws,
+            mask=mask,
+            background=background,
+            qe_mask=qe_mask,
+            translation_scale=translation_scale,
+            saturation=saturation,
+            probe_basis=probe_basis,
+            probe_support=probe_support,
+            fourier_probe=fourier_probe,
+            oversampling=oversampling,
+            loss=loss,
+            units=units,
+            probe_fourier_shifts=probe_fourier_shifts,
+            simulate_probe_translation=simulate_probe_translation,
+            simulate_finite_pixels=simulate_finite_pixels,
+            phase_only=phase_only,
+            exponentiate_obj=exponentiate_obj,
+            obj_view_crop=obj_view_crop
+        )
 
 
     def interaction(self, index, translations, *args):
@@ -521,6 +545,7 @@ class FancyPtycho(CDIModel):
             wavefields,
             self.background,
             measurement=tools.measurements.incoherent_sum,
+            qe_mask=self.qe_mask,
             saturation=self.saturation,
             oversampling=self.oversampling,
             simulate_finite_pixels=self.simulate_finite_pixels,
@@ -840,7 +865,10 @@ class FancyPtycho(CDIModel):
         ('Corrected Translations',
          lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
         ('Background',
-         lambda self, fig: p.plot_amplitude(self.background**2, fig=fig))
+         lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
+        ('Quantum Efficiency Mask',
+         lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
+         lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None))
     ]
     
     
diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py
index e726149..8ce41ae 100644
--- a/src/cdtools/tools/measurements/measurements.py
+++ b/src/cdtools/tools/measurements/measurements.py
@@ -155,7 +155,18 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
         return t.clamp(output + epsilon,0,saturation)
 
 
-def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False):
+def quadratic_background(
+        wavefield,
+        background,
+        *args,
+        detector_slice=None,
+        measurement=intensity,
+        epsilon=1e-7,
+        qe_mask=None,
+        saturation=None,
+        oversampling=1,
+        simulate_finite_pixels=False
+):
     """Returns the intensity of a wavefield plus a background
     
     The intensity is calculated via the given measurment function 
@@ -173,6 +184,8 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
         Optional, a slice or tuple of slices defining a section of the simulation to return
     measurement : function
         Default is measurements.intensity, the measurement function to use.
+    qe_mask : torch.Tensor
+        A tensor storing the per-pixel quantum efficiency (up to an unknown global scaling factor)
     saturation : float
         Optional, a maximum saturation value to clamp the resulting intensities to
     oversampling : int
@@ -183,18 +196,28 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
     sim_patterns : torch.Tensor
         A real MxN array storing the wavefield's intensities
     """
-    
-    if detector_slice is None:
-        output = measurement(wavefield, *args, epsilon=epsilon,
-                             oversampling=oversampling,
-                             simulate_finite_pixels=simulate_finite_pixels) \
-                             + background**2
-    else:
-        output = measurement(wavefield, *args, detector_slice=detector_slice,
-                             epsilon=epsilon, oversampling=oversampling,
-                             simulate_finite_pixels=simulate_finite_pixels) \
-                            + background**2
 
+    if detector_slice is None:
+        raw_intensity = measurement(
+            wavefield,
+            *args,
+            epsilon=epsilon,
+            oversampling=oversampling,
+            simulate_finite_pixels=simulate_finite_pixels)
+    else:
+        raw_intensity = measurement(
+            wavefield,
+            *args,
+            detector_slice=detector_slice,
+            epsilon=epsilon,
+            oversampling=oversampling,
+            simulate_finite_pixels=simulate_finite_pixels)
+
+    if qe_mask is None:
+        output = raw_intensity + background**2
+    else:
+        output = (qe_mask * raw_intensity) + background**2
+        
     # This has to be done after the background is added, hence we replicate
     # it here
     if saturation is None:

From c86e505ae5e5871e8da88e2aa1339d872eae126c Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Tue, 13 May 2025 14:06:02 +0200
Subject: [PATCH 04/55] Make it possible to add a quantum efficiency mask to a
 dataset, allow fancy_ptycho to load that and use it in the forward model

---
 src/cdtools/datasets/base.py              | 20 ++++++--
 src/cdtools/datasets/ptycho_2d_dataset.py | 56 ++++++++++++++++++-----
 src/cdtools/models/fancy_ptycho.py        |  9 +++-
 3 files changed, 69 insertions(+), 16 deletions(-)

diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py
index 9dc4801..85ec854 100644
--- a/src/cdtools/datasets/base.py
+++ b/src/cdtools/datasets/base.py
@@ -34,10 +34,16 @@ class CDataset(torchdata.Dataset):
     needed to allow for easy mixing of data on the CPU and GPU.
     """
 
-    def __init__(self, entry_info=None, sample_info=None,
-                 wavelength=None,
-                 detector_geometry=None, mask=None,
-                 background=None):
+    def __init__(
+            self,
+            entry_info=None,
+            sample_info=None,
+            wavelength=None,
+            detector_geometry=None,
+            mask=None,
+            qe_mask=None,
+            background=None,
+    ):
 
         """The __init__ function allows construction from python objects.
 
@@ -73,6 +79,12 @@ class CDataset(torchdata.Dataset):
             self.mask = t.tensor(mask, dtype=t.bool)
         else:
             self.mask = None
+
+        if qe_mask is not None:
+            self.qe_mask = t.as_tensor(qe_mask, dtype=t.float32)
+        else:
+            self.qe_mask = None
+
         if background is not None:
             self.background = t.tensor(background, dtype=t.float32)
         else:
diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 8b6f6d8..f5b92f9 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -77,6 +77,7 @@ class Ptycho2DDataset(CDataset):
             self.intensities = t.as_tensor(intensities, dtype=t.float32)
         else:
             self.intensities = None
+
             
     def __len__(self):
         return self.patterns.shape[0]
@@ -372,10 +373,19 @@ class Ptycho2DDataset(CDataset):
         equal to the sum of a <factor> x <factor> region of pixels in the
         input pattern. This summation is done by pytorch.functional.avg_pool2d.
         
-        Any mask and background data which is stored with the dataset is
-        downsampled with the data. The background is downsampled using the same
-        method as the data. The mask is expanded so that any output pixel
-        containing a masked pixel will be masked.
+        Any mask, quantum efficiency, and background data which is stored with
+        the dataset is downsampled with the data. The background is downsampled
+        using the same method as the data.
+
+        If there is no quantum efficiency mask, then the mask is downsapled so
+        that any output pixel containing a masked pixel will be masked. If there
+        is a quantum efficiency mask, then the quantum efficiency mask is
+        downsampled using the same method as the data, and the mask is
+        downsampled to include any pixels for which there is at least one valid
+        pixel.
+
+        To avoid leakage of data from masked pixels, the data is first
+        multiplied by the mask before downsampling.
         
         Parameters
         ----------
@@ -383,17 +393,41 @@ class Ptycho2DDataset(CDataset):
             Default 2, the factor to downsample by
 
         """
-        self.patterns = t.nn.functional.avg_pool2d(
-            self.patterns.unsqueeze(0), factor, divisor_override=1)[0]
-        self.mask = t.logical_not(t.nn.functional.max_pool2d(
-            (1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
-            factor
-        )[0,0].to(dtype=t.bool))
+        if hasattr(self, 'mask') and self.mask is not None:
+            self.patterns = t.nn.functional.avg_pool2d(
+                (self.mask * self.patterns).unsqueeze(0),
+                factor, divisor_override=1)[0]
+        else:
+            self.patterns = t.nn.functional.avg_pool2d(
+                self.patterns.unsqueeze(0),
+                factor, divisor_override=1)[0]
+            
+
+        # If we have a QE mask, we want to include all pixels for which at
+        # least one of the input pixels was unmasked, because we can account
+        # for the masked pixels through quantum efficiency
+        if hasattr(self, 'qe_mask') and self.qe_mask is not None:
+            self.qe_mask = t.nn.functional.avg_pool2d(
+                (self.mask * self.qe_mask).unsqueeze(0).unsqueeze(0),
+                factor)[0,0]
+            self.mask = t.nn.functional.max_pool2d(
+                self.mask.to(dtype=t.uint8).unsqueeze(0).unsqueeze(0),
+                factor)[0,0].to(dtype=t.bool)
+            
+        # But if there is no QE mask, we need to only preserve pixels for
+        # which all input pixels were unmasked
+        elif hasattr(self, 'mask') and self.mask is not None:
+            self.mask = t.logical_not(t.nn.functional.max_pool2d(
+                (1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
+                factor
+            )[0,0].to(dtype=t.bool))
         
         self.detector_geometry['basis'] = \
             self.detector_geometry['basis'] * factor
 
-        if self.background is not None:
+
+        
+        if hasattr(self, 'background') and self.background is not None:
             self.background = t.nn.functional.avg_pool2d(
                 self.background.unsqueeze(0).unsqueeze(0),
                 factor,
diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 9459855..0f0323e 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -95,6 +95,10 @@ class FancyPtycho(CDIModel):
         else:
             self.qe_mask = t.nn.Parameter(
                 t.as_tensor(qe_mask, dtype=dtype))
+            # I want the ability to optimize over this, but experience shows
+            # that it is wildly unstable, so I think it's best to keep
+            # gradients turned off by default
+            self.qe_mask.requires_grad=False
             
         probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
         obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
@@ -386,7 +390,10 @@ class FancyPtycho(CDIModel):
             mask = None
 
         if use_qe_mask:
-            qe_mask = t.ones(mask.shape, dtype=t.float32)
+            if hasattr(dataset, 'qe_mask') and dataset.qe_mask is not None:
+                qe_mask = t.as_tensor(dataset.qe_mask, dtype=t.float32)
+            else:
+                qe_mask = t.ones(dataset.patterns.shape[-2:], dtype=t.float32)
         else:
             qe_mask = None
             

From 2c2f2d2d935708a5b203ad4ffb434d159df68089 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Tue, 13 May 2025 16:05:16 +0200
Subject: [PATCH 05/55] Add a way to save and load the quantum efficiency
 masks, and add test coverage

---
 src/cdtools/datasets/base.py      | 19 ++++++--
 src/cdtools/tools/data/data.py    | 75 ++++++++++++++++++++++++++++
 tests/conftest.py                 | 14 ++++++
 tests/models/test_fancy_ptycho.py |  3 +-
 tests/test_datasets.py            | 81 ++++++++++++++++++++++++++-----
 tests/tools/test_data.py          | 22 ++++++++-
 6 files changed, 195 insertions(+), 19 deletions(-)

diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py
index 85ec854..3f8ec8c 100644
--- a/src/cdtools/datasets/base.py
+++ b/src/cdtools/datasets/base.py
@@ -110,6 +110,8 @@ class CDataset(torchdata.Dataset):
 
         if self.mask is not None:
             self.mask = self.mask.to(*args,**mask_kwargs)
+        if self.qe_mask is not None:
+            self.qe_mask = self.qe_mask.to(*args,**kwargs)
         if self.background is not None:
             self.background = self.background.to(*args,**kwargs)
 
@@ -205,12 +207,17 @@ class CDataset(torchdata.Dataset):
                              'basis'    : basis,
                              'corner'   : corner}
         mask = cdtdata.get_mask(cxi_file)
+        qe_mask = cdtdata.get_qe_mask(cxi_file)
         dark = cdtdata.get_dark(cxi_file)
-        return cls(entry_info = entry_info,
-                   sample_info = sample_info,
-                   wavelength=wavelength,
-                   detector_geometry=detector_geometry,
-                   mask=mask, background=dark)
+        return cls(
+            entry_info=entry_info,
+            sample_info=sample_info,
+            wavelength=wavelength,
+            detector_geometry=detector_geometry,
+            mask=mask,
+            qe_mask=qe_mask,
+            background=dark,
+        )
     
     
     def to_cxi(self, cxi_file):
@@ -248,6 +255,8 @@ class CDataset(torchdata.Dataset):
                                corner = corner)
         if self.mask is not None:
             cdtdata.add_mask(cxi_file, self.mask)
+        if self.qe_mask is not None:
+            cdtdata.add_qe_mask(cxi_file, self.qe_mask)
         if self.background is not None:
             cdtdata.add_dark(cxi_file, self.background)
         
diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py
index b69d0c0..89879dd 100644
--- a/src/cdtools/tools/data/data.py
+++ b/src/cdtools/tools/data/data.py
@@ -22,6 +22,7 @@ __all__ = ['get_entry_info',
            'get_wavelength',
            'get_detector_geometry',
            'get_mask',
+           'get_qe_mask',
            'get_dark',
            'get_data',
            'get_shot_to_shot_info',
@@ -32,6 +33,7 @@ __all__ = ['get_entry_info',
            'add_source',
            'add_detector',
            'add_mask',
+           'add_qe_mask',
            'add_dark',
            'add_data',
            'add_shot_to_shot_info',
@@ -300,6 +302,42 @@ def get_mask(cxi_file):
         return None
 
 
+def get_qe_mask(cxi_file):
+    """Returns the quantum efficiency mask defined in the cxi file object
+
+    There is no way to store a quantum efficiency mask (a.k.a. a flat-field
+    image) in the .cxi file specification, but experience has indicated that
+    this is often a valuable thing to store, because just correcting for a
+    flatfield with e.g. a division will mess up the photon counting statistics.
+
+    Because there is no specification, I have simply chosen to store the
+    quantum efficiency mask as a float32 array in the same location as the
+    mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
+    
+    The stored quantum efficiency mask should be defined as the mask that
+    a simulated intensity pattern needs to be multiplied by to realize the
+    measured image. In other words, it should be a flat-field image, not the
+    inverse of a flat-field image.
+    
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        A file object to be read
+
+    Returns
+    -------
+    qe_mask : np.array
+        A float32 array storing the quantum efficiency mask from the cxi file
+    """
+
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1/qe_mask' in i1:
+        qe_mask = i1['detector_1/qe_mask'][()].astype(np.float32)
+        return qe_mask
+    else:
+        return None
+
+
 def get_dark(cxi_file):
     """Returns an array with a dark image to use for initialization of a background model
 
@@ -635,6 +673,43 @@ def add_mask(cxi_file, mask):
     d1.create_dataset('mask',data=mask_to_save)
 
 
+def add_qe_mask(cxi_file, qe_mask):
+    """Adds the specified quantum efficiency mask to the cxi file
+
+    There is no way to store a quantum efficiency mask (a.k.a. a flat-field
+    image) in the .cxi file specification, but experience has indicated that
+    this is often a valuable thing to store, because just correcting for a
+    flatfield with e.g. a division will mess up the photon counting statistics.
+
+    Because there is no specification, I have simply chosen to store the
+    quantum efficiency mask as an array in the same location as the
+    mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
+    
+    The stored quantum efficiency mask should be defined as the mask that
+    a simulated intensity pattern needs to be multiplied by to realize the
+    measured image. In other words, it should be a flat-field image, not the
+    inverse of a flat-field image.
+
+    Parameters
+    ----------
+    cxi_file : h5py.File
+        The file to add the mask to
+    qe_mask : array
+        The quantum efficiency mask to save out to the file
+    """
+
+    if 'entry_1/instrument_1' not in cxi_file:
+        cxi_file['entry_1'].create_group('instrument_1')
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1' not in i1:
+        i1.create_group('detector_1')
+    d1 = i1['detector_1']
+    if isinstance(qe_mask, t.Tensor):
+        qe_mask = qe_mask.detach().cpu().numpy()
+
+    d1.create_dataset('qe_mask',data=qe_mask)
+
+
 def add_dark(cxi_file, dark):
     """Adds the specified dark image to a cxi file
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 51a58dc..781dd9e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -137,9 +137,16 @@ def ptycho_cxi_1():
     # Remember the format for the CXI file differs from the format used
     # internally
     mask = np.zeros((256,256)).astype(np.int32)
+    mask[5,8] = 1
     expected['mask'] = np.ones((256,256)).astype(bool)
+    expected['mask'][5,8] = 0
     d1f.create_dataset('mask',data=mask)
 
+    # There is no specification for this in the CXI file format :(
+    qe_mask = np.ones((256,256)).astype(np.float32)
+    expected['qe_mask'] = qe_mask
+    d1f.create_dataset('qe_mask',data=qe_mask)
+    
     # Create an initial background
     dark = np.ones((256,256)) * 0.01
     expected['dark'] = dark
@@ -228,6 +235,8 @@ def ptycho_cxi_2():
     # internally
     expected['mask'] = None
 
+    expected['qe_mask'] = None
+    
     # Test with a set of dark images
     dark = np.ones((10,256,256)) * 0.01
     expected['dark'] = np.nanmean(dark,axis=0)
@@ -305,8 +314,13 @@ def ptycho_cxi_3():
     # Remember the format for the CXI file differs from the format used
     # internally
     mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
+    mask[15,47] = 38
     expected['mask'] = np.ones((256,256)).astype(bool)
+    expected['mask'][15,47] = 0
     d1f.create_dataset('mask',data=mask)
+
+    expected['qe_mask'] = None
+    
     expected['dark'] = None
     
     data1f = e1f.create_group('data_1')
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index c9eaa1b..1591bf9 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -21,6 +21,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         propagation_distance=5e-3, 
         units='mm', 
         obj_view_crop=-50,
+        use_qe_mask=True, # test this in the case where no qe mask is defined
     )
     
     print('Running reconstruction on provided reconstruction_device,',
@@ -28,7 +29,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
     model.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
-    for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
+    for loss in model.Adam_optimize(70, dataset, lr=0.02, batch_size=10):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 6897991..3bb91be 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -62,6 +62,9 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
         if expected['mask'] is not None:
             assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
 
+        if expected['qe_mask'] is not None:
+            assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
+
         if expected['dark'] is not None:
             assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
                               dataset.background))
@@ -101,6 +104,9 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
         if dataset.mask is not None:
             assert t.all(t.eq(dataset.mask,read_dataset.mask))
 
+        if dataset.qe_mask is not None:
+            assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
+
         if dataset.background is not None:
             assert t.all(t.eq(dataset.background, read_dataset.background))
 
@@ -115,6 +121,7 @@ def test_CDataset_to(ptycho_cxi_1):
     if t.cuda.is_available():
         dataset.to(device='cuda:0')
         assert dataset.mask.device == t.device('cuda:0')
+        assert dataset.qe_mask.device == t.device('cuda:0')
         assert dataset.background.device == t.device('cuda:0')
 
 
@@ -135,6 +142,7 @@ def test_Ptycho2DDataset_init():
                                             [-20e-6,0,0]]).transpose(),
                          'corner': np.array((2550e-6,3825e-6,0.3))}
     mask = np.ones((256,256))
+    qe_mask = 1.2*np.ones((256,256), dtype=np.float32)
     patterns = np.random.rand(20,256,256)
     translations = np.random.rand(20,3)
     
@@ -153,6 +161,24 @@ def test_Ptycho2DDataset_init():
     assert t.allclose(dataset.patterns, t.as_tensor(patterns))
     assert t.allclose(dataset.translations, t.as_tensor(translations))
 
+    # Also test one with a qe_mask
+    dataset = Ptycho2DDataset(translations, patterns,
+                                entry_info=entry_info,
+                                sample_info=sample_info,
+                                wavelength=wavelength,
+                                detector_geometry=detector_geometry,
+                                mask=mask,
+                                qe_mask=qe_mask)
+
+    assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
+    assert t.all(t.eq(dataset.qe_mask,t.as_tensor(qe_mask)))
+    assert dataset.entry_info == entry_info
+    assert dataset.sample_info == sample_info
+    assert dataset.wavelength == wavelength
+    assert dataset.detector_geometry == detector_geometry
+    assert t.allclose(dataset.patterns, t.as_tensor(patterns))
+    assert t.allclose(dataset.translations, t.as_tensor(translations))
+
 
 def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
@@ -182,6 +208,9 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
         if expected['mask'] is not None:
             assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
 
+        if expected['qe_mask'] is not None:
+            assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
+
         if expected['dark'] is not None:
             assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
                               dataset.background))
@@ -221,11 +250,12 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
         if dataset.detector_geometry['corner'] is not None:
             assert 'corner' in read_dataset.detector_geometry
             
-
-        
         if dataset.mask is not None:
             assert t.all(t.eq(dataset.mask,read_dataset.mask))
 
+        if dataset.qe_mask is not None:
+            assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
+
         if dataset.background is not None:
             assert t.all(t.eq(dataset.background, read_dataset.background))
 
@@ -238,12 +268,14 @@ def test_Ptycho2DDataset_to(ptycho_cxi_1):
     
     dataset.to(dtype=t.float64)
     assert dataset.mask.dtype == t.bool
+    assert dataset.qe_mask.dtype == t.float64
     assert dataset.patterns.dtype == t.float64
     assert dataset.translations.dtype == t.float64
     # If cuda is available, check that moving the mask to CUDA works.
     if t.cuda.is_available():
         dataset.to(device='cuda:0')
         assert dataset.mask.device == t.device('cuda:0')
+        assert dataset.qe_mask.device == t.device('cuda:0')
         assert dataset.background.device == t.device('cuda:0')
         assert dataset.patterns.device == t.device('cuda:0')
         assert dataset.translations.device == t.device('cuda:0')
@@ -291,24 +323,49 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
         # May start failing if the test datasets are changed to include
         # a dataset with any dimension not even. That's a problem with the
         # test, not the code. Sorry! -Abe
+
+        masked_patterns = dataset.mask * dataset.patterns
         assert t.allclose(
             copied_dataset.patterns,
-            dataset.patterns[:,::2,::2] +
-            dataset.patterns[:,1::2,::2] +
-            dataset.patterns[:,::2,1::2] +
-            dataset.patterns[:,1::2,1::2]
+            masked_patterns[:,::2,::2] +
+            masked_patterns[:,1::2,::2] +
+            masked_patterns[:,::2,1::2] +
+            masked_patterns[:,1::2,1::2]
         )
-        
-        assert t.allclose(
-            copied_dataset.mask,
-            t.logical_and(
+
+        if dataset.qe_mask is None:
+            manually_downsampled_mask = t.logical_and(
                 t.logical_and(dataset.mask[::2,::2],
                               dataset.mask[1::2,::2]),
                 t.logical_and(dataset.mask[::2,1::2],
-                              dataset.mask[1::2,1::2]),
+                              dataset.mask[1::2,1::2])
+            )
+            assert t.allclose(
+                copied_dataset.mask,
+                manually_downsampled_mask,
+            )
+        else:
+            manually_downsampled_mask = t.logical_or(
+                t.logical_or(dataset.mask[::2,::2],
+                             dataset.mask[1::2,::2]),
+                t.logical_or(dataset.mask[::2,1::2],
+                             dataset.mask[1::2,1::2])
+            )
+            assert t.allclose(
+                copied_dataset.mask,
+                manually_downsampled_mask
             )
-        )
 
+            masked_qe_mask = dataset.mask * dataset.qe_mask
+            manually_downsampled_qe_mask = (
+                masked_qe_mask[::2,::2] + masked_qe_mask[1::2,::2]
+                + masked_qe_mask[::2,1::2] + masked_qe_mask[1::2,1::2]
+            ) / 4
+
+            assert t.allclose(
+                copied_dataset.qe_mask,
+                manually_downsampled_qe_mask
+            )
 
         
         if dataset.background is not None:
diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py
index e637509..39b64dd 100644
--- a/tests/tools/test_data.py
+++ b/tests/tools/test_data.py
@@ -60,7 +60,15 @@ def test_get_mask(test_ptycho_cxis):
         mask = data.get_mask(cxi)
         if expected['mask'] is None and mask is None:
             continue
-        assert np.all(data.get_mask(cxi) == expected['mask'])
+        assert np.all(mask == expected['mask'])
+
+        
+def test_get_qe_mask(test_ptycho_cxis):
+    for cxi, expected in test_ptycho_cxis:
+        qe_mask = data.get_qe_mask(cxi)
+        if expected['qe_mask'] is None and qe_mask is None:
+            continue
+        assert np.allclose(qe_mask, expected['qe_mask'])
 
 
 def test_get_dark(test_ptycho_cxis):
@@ -207,6 +215,18 @@ def test_add_mask(tmp_path):
     assert np.all(mask == read_mask)
 
     
+def test_add_qe_mask(tmp_path):
+    qe_mask = np.random.rand(350,199).astype(np.float32)
+
+    with data.create_cxi(tmp_path / 'test_add_qe_mask.cxi') as f:
+        data.add_qe_mask(f, qe_mask)
+
+    with h5py.File(tmp_path / 'test_add_qe_mask.cxi','r') as f:
+        read_qe_mask = data.get_qe_mask(f)
+
+    assert np.allclose(qe_mask, read_qe_mask)
+
+    
 def test_add_dark(tmp_path):
     dark = np.random.rand(350,620)
 

From 2c5587cacfae8968459980cfd1604b2f5b7bdfcd Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Wed, 22 Jan 2025 16:41:18 +0100
Subject: [PATCH 06/55] Fix an issue for datasets with colinear translations

---
 src/cdtools/datasets/ptycho_2d_dataset.py |  5 ++++-
 src/cdtools/tools/analysis/analysis.py    | 24 ++++++++++++++++++-----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index f5b92f9..46d2ba7 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -270,9 +270,12 @@ class Ptycho2DDataset(CDataset):
         """Plots the mean diffraction pattern across the dataset
 
         The output is normalized so that the summed intensity on the
-        detector is equal to the total intensity of light that passed
+        detector is roughly equal to the total intensity of light that passed
         through the sample within each detector conjugate field of view.
 
+        If the scan points are colinear (which causes issues for this
+        estimation), the mean pattern is displayed unscaled.
+
         The plot is plotted as log base 10 of the output plus log_offset.
         By default, log_offset is set equal to 1, which is a good level for
         shot-noise limited data captured in units of photons. More
diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py
index 0fd913e..9844edb 100644
--- a/src/cdtools/tools/analysis/analysis.py
+++ b/src/cdtools/tools/analysis/analysis.py
@@ -14,6 +14,7 @@ from scipy import linalg as sla
 from scipy import special
 from scipy import optimize as opt
 from scipy import spatial
+import warnings
 
 __all__ = [
     'product_svd',
@@ -1443,6 +1444,13 @@ def calc_spectral_info(dataset, nbins=50):
     the scan pattern whose area matches one detector conjugate field of
     view.
 
+    This estimation will start to deviate from the truth if the scan area
+    is not significantly larger than the illumination function, because
+    the nonzero size of the illumination function is not taken into account.
+    Furthermore, in the edge case where all the scan points are colinear,
+    the estimate will fail, and the mean diffraction pattern will be returned
+    instead
+
     Parameters
     ----------
     dataset : Ptycho2DDataset
@@ -1461,10 +1469,12 @@ def calc_spectral_info(dataset, nbins=50):
     
     """
 
-    scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
-    
-    scan_area = scan_hull.volume
-
+    try:
+        scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
+        scan_area = scan_hull.volume
+    except spatial._qhull.QhullError as e:
+        scan_area = None
+        
     ewg = cdtools.tools.initializers.exit_wave_geometry
     obj_basis = ewg(
         dataset.detector_geometry['basis'],
@@ -1477,8 +1487,12 @@ def calc_spectral_info(dataset, nbins=50):
         np.cross(obj_basis[:,0]*dataset.patterns.shape[-2],
                  obj_basis[:,1]*dataset.patterns.shape[-1])
     )
-    scale_factor = det_conj_fov_area / scan_area
 
+    if scan_area is not None:
+        scale_factor = det_conj_fov_area / scan_area
+    else:
+        warnings.warn("The scan points in this dataset are all colinear. The mean pattern will be calculated rather than a scaled mean based on the scanned area.")
+        scale_factor = 1/len(dataset)
 
     mask = dataset.mask.cpu().numpy().astype(int)
     sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor

From b6bc7d95ccdca1779f873c5e8ca03d9e3aa7c86d Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 7 Apr 2025 09:35:59 +0200
Subject: [PATCH 07/55] Change the behavior of FancyPtycho.center_probes to
 center the probe in real space even when fourier_probe is set to True

---
 src/cdtools/models/fancy_ptycho.py | 35 +++++++++++++++++++++++----
 tests/models/test_fancy_ptycho.py  | 39 ++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 0f0323e..6220299 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -629,15 +629,40 @@ class FancyPtycho(CDIModel):
 
         
     def center_probes(self, iterations=4):
-        """Centers the probes
+        """Centers the probes in real space
+
+        Takes the current guess of the illumination function and centers it
+        using a shift with periodic boundary conditions. It uses
+        cdtools.tools.image_processing.center internally to do the centering.
+        Multiple iterations of an algorithm are run, which is helpful if the
+        illumination is reconstructed near the corners and "wraps around" the
+        probe field of view.
+
+        Note that the centering is always performed in real space, even if
+        the probe array is defined in Fourier space.
         
-        Note that this does not compensate for the centering by adjusting
+        Note also that this does not compensate for the centering by adjusting
         the object, so it's a good idea to reset the object after centering
         the probes
+
+        Parameters
+        ----------
+        iterations : int
+            Default 4, how many iterations of the centering algorithm to run
         """
-        centered_probe = tools.image_processing.center(
-            self.probe.data.cpu(), iterations=iterations)
-        self.probe.data = centered_probe.to(device=self.probe.data.device)
+        if self.fourier_probe:
+            prs = tools.propagators.inverse_far_field(self.probe.detach()).cpu()
+        else:
+            prs = self.probe.detach().cpu()
+        
+        centered_prs = tools.image_processing.center(prs, iterations=iterations)
+
+        if self.fourier_probe:
+            self.probe.data = tools.propagators.far_field(
+                centered_prs.to(device=self.probe.data.device))
+        else:
+            self.probe.data = centered_prs.to(device=self.probe.data.device)
+
 
 
     def tidy_probes(self):
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 1591bf9..7ef429f 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -5,6 +5,45 @@ import torch as t
 import cdtools
 from matplotlib import pyplot as plt
 
+
+def test_center_probe(lab_ptycho_cxi):
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
+    model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=False
+    )
+    base_probe = model.probe.detach().clone()
+    model.center_probes()
+    centered_probe = model.probe.detach().clone()
+
+    fourier_model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=True,
+    )
+
+    fourier_model.probe.data = cdtools.tools.propagators.far_field(
+        base_probe
+    )
+    
+    fourier_base_probe = fourier_model.probe.detach().clone()
+    fourier_model.center_probes()
+    fourier_centered_probe = fourier_model.probe.detach().clone()
+    ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
+        fourier_centered_probe)
+
+    # So we know the code had to do something
+    assert not t.allclose(base_probe, centered_probe)
+    # And checking that they both do the same thing, whether or not
+    # fourier_probe was set to True
+    assert t.allclose(
+        centered_probe,
+        ifft_fourier_centered_probe,
+        atol=1e-4,
+        rtol=1e-3
+    )
+    
 @pytest.mark.slow
 def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 

From 4ecc782c2852d4b7d551d2c479d2f5ae7b406e4e Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Tue, 10 Jun 2025 14:11:01 +0200
Subject: [PATCH 08/55] Reply to Daynes comment's and slightly loosen a
 tolerance in one test that has been failing

---
 src/cdtools/datasets/ptycho_2d_dataset.py     |  2 +-
 .../tools/measurements/measurements.py        | 25 +++++++------------
 tests/tools/test_analysis.py                  |  4 +--
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 46d2ba7..6631cd6 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -380,7 +380,7 @@ class Ptycho2DDataset(CDataset):
         the dataset is downsampled with the data. The background is downsampled
         using the same method as the data.
 
-        If there is no quantum efficiency mask, then the mask is downsapled so
+        If there is no quantum efficiency mask, then the mask is downsampled so
         that any output pixel containing a masked pixel will be masked. If there
         is a quantum efficiency mask, then the quantum efficiency mask is
         downsampled using the same method as the data, and the mask is
diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py
index 8ce41ae..c926568 100644
--- a/src/cdtools/tools/measurements/measurements.py
+++ b/src/cdtools/tools/measurements/measurements.py
@@ -196,22 +196,15 @@ def quadratic_background(
     sim_patterns : torch.Tensor
         A real MxN array storing the wavefield's intensities
     """
-
-    if detector_slice is None:
-        raw_intensity = measurement(
-            wavefield,
-            *args,
-            epsilon=epsilon,
-            oversampling=oversampling,
-            simulate_finite_pixels=simulate_finite_pixels)
-    else:
-        raw_intensity = measurement(
-            wavefield,
-            *args,
-            detector_slice=detector_slice,
-            epsilon=epsilon,
-            oversampling=oversampling,
-            simulate_finite_pixels=simulate_finite_pixels)
+    
+    raw_intensity = measurement(
+        wavefield,
+        *args,
+        detector_slice=detector_slice,
+        epsilon=epsilon,
+        oversampling=oversampling,
+        simulate_finite_pixels=simulate_finite_pixels
+    )
 
     if qe_mask is None:
         output = raw_intensity + background**2
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index 27c6a65..f0a1cf1 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -160,7 +160,7 @@ def test_standardize():
     probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
 
     assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
-    assert np.angle(np.sum(probe)) < 1e-7
+    assert np.angle(np.sum(probe)) < 2e-7
 
     obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
     obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -232,7 +232,7 @@ def test_synthesize_reconstructions():
     probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
 
     assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
-    assert np.abs(np.angle(np.sum(probe))) < 1e-7
+    assert np.abs(np.angle(np.sum(probe))) < 2e-7
 
     obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
     obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,

From 4c784de9224431cafe92659bfd4ac7ef4fd83e3e Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Thu, 12 Jun 2025 00:02:32 +0000
Subject: [PATCH 09/55] Fixed sim_data shape change in CDIModel.compare for RPI
 measurements

---
 src/cdtools/models/base.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index b31061d..b455f45 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -872,7 +872,13 @@ class CDIModel(t.nn.Module):
             updating = True if len(axes[0].images) >= 1 else False
 
             inputs, output = dataset[idx:idx+1]
-            sim_data = self.forward(*inputs).detach().cpu().numpy()[0]
+            sim_data = self.forward(*inputs).detach().cpu().numpy()
+            # The length of sim_data.shape changes when you're doing 
+            # either a ptycho (3) or an RPI (2) reconstruction.
+            # We need to make sure that sim_data is 2D.
+            if len(sim_data.shape) > 2:
+                sim_data = sim_data[0]
+
             meas_data = output.detach().cpu().numpy()[0]
             if hasattr(self, 'mask') and self.mask is not None:
                 mask = self.mask.detach().cpu().numpy()

From b189c3c2af24ba706dc3159a9819b57115374391 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Tue, 17 Jun 2025 11:01:30 -0700
Subject: [PATCH 10/55] Add warning for 64-bit float conversion in
 Ptycho2DDataset and added to tests

---
 src/cdtools/datasets/ptycho_2d_dataset.py | 19 +++++----
 tests/test_datasets.py                    | 52 ++++++++++++++++++++---
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 6631cd6..c653a25 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -1,15 +1,16 @@
+import warnings
+from copy import copy, deepcopy
+import pathlib
+
+import h5py
 import numpy as np
 import torch as t
-from copy import copy
-import h5py
-import pathlib
+
 from cdtools.datasets import CDataset
 from cdtools.datasets.random_selection import random_selection
 from cdtools.tools import data as cdtdata
 from cdtools.tools import plotting
-from matplotlib import pyplot as plt
 from cdtools.tools import analysis
-from copy import deepcopy
 
 __all__ = ['Ptycho2DDataset']
 
@@ -164,10 +165,12 @@ class Ptycho2DDataset(CDataset):
             patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
             dataset.patterns = t.as_tensor(patterns)
             if dataset.patterns.dtype == t.float64:
-                raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request')
-            
+                # If the data is 64-bit, we need to convert it to 32-bit
+                # because 64-bit floats are not supported in reconstructions
+                dataset.patterns = dataset.patterns.to(dtype=t.float32)
+                warnings.warn('64-bit floats are not supported and precision will not be retained in reconstructions and were converted to t.float32! Please explicitly convert your data to 32-bit or submit a pull request')
             dataset.axes = axes
-            
+
             if dataset.mask is None:
                 dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool)
 
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 3bb91be..8999117 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -1,12 +1,16 @@
+import datetime
+import itertools
+import os
+from copy import deepcopy
+
+import h5py
+import numpy as np
+import pytest
+import torch as t
+
 from cdtools.datasets import CDataset, Ptycho2DDataset
 from cdtools.tools import data as cdtdata
-import numpy as np
-import torch as t
-import h5py
-import datetime
-from copy import deepcopy
-import pytest
-import itertools
+
 
 #
 # We start by testing the CDataset base class
@@ -220,6 +224,40 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
         assert t.allclose(t.tensor(expected['translations']),dataset.translations)
 
 
+def test_Ptycho2DDataset_from_cxi_64bit(test_ptycho_cxis):
+    """Test that we can load a 64-bit cxi file. Should issue
+    a warning, but still load the data."""
+
+    # create test patterns and translations
+    np.random.seed(42)
+    patterns = np.random.rand(20, 256, 256).astype(np.float64)
+    translations = np.random.rand(20, 3).astype(np.float64)
+
+    dataset = Ptycho2DDataset(translations, patterns)
+    dataset.detector_geometry = {
+        'distance': 0.1,  # in meters
+        'basis': t.tensor([
+            [-0e-06, -13.5e-06 * 4],
+            [-13.5e-06 * 4, 0e-06],
+            [0e-06, 0e-06]
+        ]),
+        'corner': None
+    }
+    dataset.wavelength = 1.6891579427792915e-09  # in meters
+    # and save to a temp file
+    dataset.to_cxi('test_Ptycho2DDataset_from_cxi_64bit.cxi')
+
+    with pytest.warns(UserWarning, match='64-bit floats'):
+        dataset_64bit = Ptycho2DDataset.from_cxi('test_Ptycho2DDataset_from_cxi_64bit.cxi')
+
+    # Check that the data is loaded correctly
+    assert dataset_64bit.patterns.dtype == t.float32
+    assert dataset_64bit.translations.dtype == t.float32
+
+    # delete the created test file
+    os.remove('test_Ptycho2DDataset_from_cxi_64bit.cxi')
+
+
 def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
     for cxi, expected in test_ptycho_cxis:
         print('loading dataset')

From d7a1ea82c98b166d07f117a6bb4a9f03fff300d5 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Tue, 17 Jun 2025 15:50:51 -0700
Subject: [PATCH 11/55] adjusted warn msg based on Daynes suggestion

---
 src/cdtools/datasets/ptycho_2d_dataset.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index c653a25..3825d6d 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -125,9 +125,9 @@ class Ptycho2DDataset(CDataset):
         self.translations = self.translations.to(*args, **kwargs)
         self.patterns = self.patterns.to(*args, **kwargs)
 
-
     # It sucks that I can't reuse the base factory method here,
     # perhaps there is a way but I couldn't figure it out.
+
     @classmethod
     def from_cxi(cls, cxi_file, cut_zeros=True, load_patterns=True):
         """Generates a new Ptycho2DDataset from a .cxi file directly
@@ -149,7 +149,7 @@ class Ptycho2DDataset(CDataset):
         """
         # If a bare string is passed
         if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
-            with h5py.File(cxi_file,'r') as f:
+            with h5py.File(cxi_file, 'r') as f:
                 return cls.from_cxi(f, cut_zeros=cut_zeros, load_patterns=load_patterns)
 
         # Generate a base dataset
@@ -168,7 +168,10 @@ class Ptycho2DDataset(CDataset):
                 # If the data is 64-bit, we need to convert it to 32-bit
                 # because 64-bit floats are not supported in reconstructions
                 dataset.patterns = dataset.patterns.to(dtype=t.float32)
-                warnings.warn('64-bit floats are not supported and precision will not be retained in reconstructions and were converted to t.float32! Please explicitly convert your data to 32-bit or submit a pull request')
+                warnings.warn(
+                    "64-bit floats are not supported and precision will not be retained in reconstructions and were converted to t.float32! "
+                    "If you would like to have 64-bit support, please open an issue or submit a pull request."
+                )
             dataset.axes = axes
 
             if dataset.mask is None:
@@ -179,9 +182,8 @@ class Ptycho2DDataset(CDataset):
             dataset.intensities = t.as_tensor(intensities, dtype=t.float32)
         except KeyError:
             dataset.intensities = None
-            
-        return dataset
 
+        return dataset
 
     def to_cxi(self, cxi_file):
         """Saves out a Ptycho2DDataset as a .cxi file

From 951636e39f5c1b05ff31aa0b0ffdd9778834612d Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Wed, 18 Jun 2025 14:09:43 -0700
Subject: [PATCH 12/55] pin current version to 0.3.0

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 59d4600..f0f09aa 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="cdtools",
-    version="0.2.0",
+    version="0.3.0",
     python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
     author="Abe Levitan",
     author_email="abraham.levitan@psi.ch",

From 2c764d4190f868b6c8bba21ab751cd63afbac4f7 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 23 Jun 2025 15:38:44 -0700
Subject: [PATCH 13/55] add pypi workflow

---
 .github/workflows/publish.yml | 30 ++++++++++++++++++++++++++++++
 setup.py                      |  2 +-
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/publish.yml

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
new file mode 100644
index 0000000..8ff3518
--- /dev/null
+++ b/.github/workflows/publish.yml
@@ -0,0 +1,30 @@
+name: Upload Python Package to PyPI when a Release is Created
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  pypi-publish:
+    name: Publish release to PyPI
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/cdtools-py
+    permissions:
+      id-token: write
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.x"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel
+    - name: Build package (setup.py)
+      run: |
+        python setup.py sdist bdist_wheel
+    - name: Publish package distributions to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
\ No newline at end of file
diff --git a/setup.py b/setup.py
index f0f09aa..74ca920 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="cdtools",
-    version="0.3.0",
+    version="0.3.1",
     python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
     author="Abe Levitan",
     author_email="abraham.levitan@psi.ch",

From 302decd922f15e8b962f16e69cf771c73cbc75b5 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 23 Jun 2025 15:46:32 -0700
Subject: [PATCH 14/55] remove version increase

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 74ca920..f0f09aa 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="cdtools",
-    version="0.3.1",
+    version="0.3.0",
     python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
     author="Abe Levitan",
     author_email="abraham.levitan@psi.ch",

From 1ecf001ad4c46f4b12836f411492e60532ff46c5 Mon Sep 17 00:00:00 2001
From: Damian Guenzing <65827185+gnzng@users.noreply.github.com>
Date: Sun, 29 Jun 2025 21:14:49 +0000
Subject: [PATCH 15/55] change name from cdtools to cdtools-py to match pypi
 naming

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index f0f09aa..8ac2e50 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@ with open("README.md", "r") as fh:
     long_description = fh.read()
 
 setuptools.setup(
-    name="cdtools",
+    name="cdtools-py",
     version="0.3.0",
     python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
     author="Abe Levitan",

From ab30afb10ef2d94b61d3c2a35148430db65c139e Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Wed, 2 Jul 2025 14:21:43 +0200
Subject: [PATCH 16/55] Update the documentation to reflect the existence of a
 PyPI package

---
 README.md                    |   8 +-
 docs/source/installation.rst |  75 +++++++++------
 docs/source/intro.rst        |   1 -
 example_environment.yml      | 181 -----------------------------------
 setup.py                     |   2 +-
 5 files changed, 52 insertions(+), 215 deletions(-)
 delete mode 100644 example_environment.yml

diff --git a/README.md b/README.md
index 696ab4f..76371bf 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,13 @@ model.compare(dataset) # See how the simulated and measured patterns compare
 plt.show()
 ```
 
-Full installation instructions and documentation can be found [here](https://cdtools-developers.github.io/cdtools/).
+CDTools can be installed via pip as the [cdtools-py](https://pypi.org/project/cdtools-py/) package on [PyPI](https://pypi.org/):
+
+```bash
+$ pip install cdtools-py
+```
+
+Further documentation is found [here](https://cdtools-developers.github.io/cdtools/).
 
 
 CDTools was developed in the [photon scattering lab](https://scattering.mit.edu/) at MIT, and further development took place within the [computational x-ray imaging group](https://www.psi.ch/en/cxi) at PSI. The code is distributed under an MIT (a.k.a. Expat) license. If you would like to publish any work that uses CDTools, please contact [Abe Levitan](mailto:abraham.levitan@psi.ch).
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index d998f93..cf86477 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -1,79 +1,92 @@
 Installation
 ============
 
-Step 1: Download
-----------------
+CDTools supports python >=3.8 and can be installed via pip as the the `cdtools-py`_ package on `PyPI`_. If you plan to contribute to the code or need a custom environment, installation from source is also possible.
 
-The source code for CDTools is hosted on `Github`_. At the moment, the repository remains private while we decide on licensing. Access can be granted upon request by contacting `Abe Levitan <alevitan@mit.edu>`_.
+.. _`cdtools-py`: https://pypi.org/project/cdtools-py/
+.. _`PyPI`: https://pypi.org/
 
-.. _`Github`: https://github.com/cdtools-developers/cdtools
+Option 1: Installation from PyPI
+--------------------------------
 
-The repository remains under active development as of early 2025.
+To install from `PyPI`_, run:
 
-Step 2: Install Dependencies
-----------------------------
+.. code:: bash
+	  
+   $ pip install cdtools-py
 
-CDTools is regularly tested with Python versions 3.8 to 3.12, so it is recommended to use one of these versions. In general, CDTools requires Python 3.7 or higher.
-
-The major dependency for CDTools is pytorch (version 1.9.0 or greater). Because the details of the installation can vary depending on platform, GPU availability, etc, it is recommended that you follow the install instructions on `the pytorch site`_ to install pytorch before installing the remaining dependencies.
+Pytorch, a major dependence of CDTools, often needs to be installed with a specific CUDA version for machine compatability. If you run into issues with pytorch, consider first installing pytorch into your environment using the instructions on `the pytorch site`_.
 
 .. _`the pytorch site`: https://pytorch.org/get-started/locally/
 
-pytorch stopped supporting installation using conda for installation, so it is recommended continue the installation using pip.
+Option 2: Installation from source
+----------------------------------
+
+
+Step 1: Download
+^^^^^^^^^^^^^^^^
+
+The source code for CDTools is hosted on `Github`_.
+
+.. _`Github`: https://github.com/cdtools-developers/cdtools
+
+
+Step 2: Install Dependencies
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The major dependency for CDTools is pytorch version 2.3.0 or greater. Because the details of the pytorch installation can vary depending on platform and GPU availability, it is recommended that you first install pytorch using the instructions on `the pytorch site`_. The remaining dependencies can be installed by running the following command from the top level directory of the git repository:
 
 .. code:: bash
 	  
    $ pip install -r requirements.txt
 
-This will install all required dependencies and verify that they meet the pytorch version requirements. Additionally, several optional dependencies used for testing and documentation will also be installed. The full set of dependencies and minimum requirements are listed below is listed below.
-
-CDTools is reguarly tested with the latest versions of the packages shown below.   
+Note that several optional dependencies used for testing and documentation will also be installed. The full set of dependencies and minimum requirements are listed below. CDTools is reguarly tested with the latest versions of these packages and with python 3.8 through 3.12.
 
 Required dependencies:
 
    * `numpy <http://www.numpy.org>`_ >= 1.0
    * `scipy <http://www.scipy.org>`_ >= 1.0
    * `matplotlib <https://matplotlib.org>`_ >= 2.0
-   * `pytorch <https://pytorch.org>`_ >= 1.9.0
+   * `pytorch <https://pytorch.org>`_ >= 2.3.0
    * `python-dateutil <https://github.com/dateutil/dateutil/>`_
    * `h5py <https://www.h5py.org/>`_ >= 2.1
 
-Optional dependencies:
+Optional dependencies for running tests:
 
    * `pytest <https://docs.pytest.org/>`_
    * `pooch <https://www.fatiando.org/pooch/latest/>`_
+
+Optional dependencies for building docs:
+
    * `sphinx <https://www.sphinx-doc.org/>`_ >= 4.3.0
    * `sphinx-argparse <https://sphinx-argparse.readthedocs.io>`_
    * `sphinx_rtd_theme <https://sphinx-rtd-theme.readthedocs.io/en/stable/>`_ >= 0.5.1
 
-The file "example_environment.yml", included in the repository's top level directory, contains an example of an environment with all dependencies properly installed on a linux machine with a GPU, circa early 2024.
-
 
 Step 3: Install
----------------
+^^^^^^^^^^^^^^^
 
-To install CDTools, run the following command from the top level directory (the directory including the setup.py file).
-
-.. code:: bash
-	  
-   $ pip install . --no-deps
-
-   
-This will install a copy of the code, as it exists at the moment of installation. If you would prefer for changes to the code to propagate to the installed version without reinstalling, install the package in developer mode:
+To install CDTools, run the following command from the top level directory of the git repository:
 
 .. code:: bash
 	  
    $ pip install -e . --no-deps
+
+   
+This will install CDTools in developer mode, so that changes to the code will propagate to the installed version immediately. This is best if you plan to actively develop CDTools. If you simply need a custom environment, you can also install CDTools in standard mode using:
+
+.. code:: bash
+	  
+   $ pip install . --no-deps
    
   
 Step 4: Run The Tests
----------------------
+^^^^^^^^^^^^^^^^^^^^^
 
-To ensure that the installation has worked correctly, it is recommended that you run the unit tests. After ensuring that pytest is installed, run the following command from the top level directory:
+To ensure that the installation has worked correctly, it is recommended that you run the unit tests. Execute the following command from the top level directory of the git repository:
 
 .. code:: bash
 
    $ pytest
 
-
-If any tests fail, make sure that you have all the noted dependencies properly installed. If you do, and things still aren't working, `send me (Abe Levitan) an email <alevitan@mit.edu>`_  and we'll get to the bottom of it. CDTools has been tested on linux and mac, on CPU, CUDA, and MPS.
+If any tests fail, make sure that you have all the noted dependencies properly installed. If you do, and things still aren't working, `open an issue on the github page <https://github.com/cdtools-developers/cdtools/issues>`_ and we'll get to the bottom of it.
diff --git a/docs/source/intro.rst b/docs/source/intro.rst
index 3a2e8d4..22bc2c1 100644
--- a/docs/source/intro.rst
+++ b/docs/source/intro.rst
@@ -44,7 +44,6 @@ The high-level interface to CDTools - datasets and models - is built on a set of
 - functions for accessing stored data in .cxi files
 - plotting tools to visualize data and reconstructions
 - basic operations, like light propagators, needed for coherent diffraction
-- tools that implement basic operations - such as light propagation - relevant to coherent diffraction.
 - analysis functions for assessing the quality of reconstructions
 
 
diff --git a/example_environment.yml b/example_environment.yml
deleted file mode 100644
index 46767a2..0000000
--- a/example_environment.yml
+++ /dev/null
@@ -1,181 +0,0 @@
-name: democdtoolsenv
-channels:
-  - pytorch
-  - nvidia
-  - conda-forge
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1=main
-  - _openmp_mutex=5.1=1_gnu
-  - alabaster=0.7.16=pyhd8ed1ab_0
-  - babel=2.14.0=pyhd8ed1ab_0
-  - blas=1.0=mkl
-  - brotli=1.0.9=h9c3ff4c_4
-  - bzip2=1.0.8=h7b6447c_0
-  - c-ares=1.19.1=h5eee18b_0
-  - ca-certificates=2024.2.2=hbcca054_0
-  - certifi=2024.2.2=pyhd8ed1ab_0
-  - charset-normalizer=2.0.4=pyhd3eb1b0_0
-  - colorama=0.4.6=pyhd8ed1ab_0
-  - commonmark=0.9.1=py_0
-  - contourpy=1.2.0=py311hdb19cb5_0
-  - cuda-cudart=12.1.105=0
-  - cuda-cupti=12.1.105=0
-  - cuda-libraries=12.1.0=0
-  - cuda-nvrtc=12.1.105=0
-  - cuda-nvtx=12.1.105=0
-  - cuda-opencl=12.3.101=0
-  - cuda-runtime=12.1.0=0
-  - cycler=0.12.1=pyhd8ed1ab_0
-  - dbus=1.13.18=hb2f20db_0
-  - docutils=0.20.1=py311h38be061_3
-  - exceptiongroup=1.2.0=pyhd8ed1ab_2
-  - expat=2.2.10=h9c3ff4c_0
-  - ffmpeg=4.3=hf484d3e_0
-  - filelock=3.13.1=py311h06a4308_0
-  - fontconfig=2.14.1=hef1e5e3_0
-  - fonttools=4.25.0=pyhd3eb1b0_0
-  - freetype=2.12.1=h4a9f257_0
-  - future=1.0.0=pyhd8ed1ab_0
-  - glib=2.78.4=h6a678d5_0
-  - glib-tools=2.78.4=h6a678d5_0
-  - gmp=6.2.1=h295c915_3
-  - gmpy2=2.1.2=py311hc9b5ff0_0
-  - gnutls=3.6.15=he1e5248_0
-  - gst-plugins-base=1.14.1=h6a678d5_1
-  - gstreamer=1.14.1=h5eee18b_1
-  - h5py=3.9.0=py311hdd6beaf_0
-  - hdf5=1.12.1=h2b7332f_3
-  - icu=58.2=hf484d3e_1000
-  - idna=3.4=py311h06a4308_0
-  - imagesize=1.4.1=pyhd8ed1ab_0
-  - importlib-metadata=7.0.1=pyha770c72_0
-  - iniconfig=2.0.0=pyhd8ed1ab_0
-  - intel-openmp=2023.1.0=hdb19cb5_46306
-  - jinja2=3.1.3=py311h06a4308_0
-  - jpeg=9e=h5eee18b_1
-  - kiwisolver=1.4.4=py311h6a678d5_0
-  - krb5=1.20.1=h143b758_1
-  - lame=3.100=h7b6447c_0
-  - lcms2=2.12=h3be6417_0
-  - ld_impl_linux-64=2.38=h1181459_1
-  - lerc=3.0=h295c915_0
-  - libclang=10.0.1=default_hb85057a_2
-  - libcublas=12.1.0.26=0
-  - libcufft=11.0.2.4=0
-  - libcufile=1.8.1.2=0
-  - libcurand=10.3.4.107=0
-  - libcurl=8.5.0=h251f7ec_0
-  - libcusolver=11.4.4.55=0
-  - libcusparse=12.0.2.55=0
-  - libdeflate=1.17=h5eee18b_1
-  - libedit=3.1.20230828=h5eee18b_0
-  - libev=4.33=h516909a_1
-  - libevent=2.1.12=hdbd6064_1
-  - libffi=3.4.4=h6a678d5_0
-  - libgcc-ng=11.2.0=h1234567_1
-  - libgfortran-ng=13.2.0=h69a702a_0
-  - libgfortran5=13.2.0=ha4646dd_0
-  - libglib=2.78.4=hdc74915_0
-  - libgomp=11.2.0=h1234567_1
-  - libiconv=1.16=h7f8727e_2
-  - libidn2=2.3.4=h5eee18b_0
-  - libjpeg-turbo=2.0.0=h9bf148f_0
-  - libllvm10=10.0.1=he513fc3_3
-  - libnghttp2=1.57.0=h2d74bed_0
-  - libnpp=12.0.2.50=0
-  - libnvjitlink=12.1.105=0
-  - libnvjpeg=12.1.1.14=0
-  - libpng=1.6.39=h5eee18b_0
-  - libpq=12.17=hdbd6064_0
-  - libssh2=1.10.0=hdbd6064_2
-  - libstdcxx-ng=11.2.0=h1234567_1
-  - libtasn1=4.19.0=h5eee18b_0
-  - libtiff=4.5.1=h6a678d5_0
-  - libunistring=0.9.10=h27cfd23_0
-  - libuuid=1.41.5=h5eee18b_0
-  - libwebp-base=1.3.2=h5eee18b_0
-  - libxcb=1.15=h7f8727e_0
-  - libxkbcommon=1.0.1=hfa300c1_0
-  - libxml2=2.9.14=h74e7548_0
-  - llvm-openmp=14.0.6=h9e868ea_0
-  - lz4-c=1.9.4=h6a678d5_0
-  - markupsafe=2.1.3=py311h5eee18b_0
-  - matplotlib=3.8.0=py311h06a4308_0
-  - matplotlib-base=3.8.0=py311ha02d727_0
-  - mkl=2023.1.0=h213fc3f_46344
-  - mkl-service=2.4.0=py311h5eee18b_1
-  - mkl_fft=1.3.8=py311h5eee18b_0
-  - mkl_random=1.2.4=py311hdb19cb5_0
-  - mpc=1.1.0=h10f8cd9_1
-  - mpfr=4.0.2=hb69a4c5_1
-  - mpmath=1.3.0=py311h06a4308_0
-  - munkres=1.1.4=pyh9f0ad1d_0
-  - ncurses=6.4=h6a678d5_0
-  - nettle=3.7.3=hbbd107a_1
-  - networkx=3.1=py311h06a4308_0
-  - nspr=4.35=h6a678d5_0
-  - nss=3.89.1=h6a678d5_0
-  - numpy=1.26.3=py311h08b1b3b_0
-  - numpy-base=1.26.3=py311hf175353_0
-  - openh264=2.1.1=h4ff587b_0
-  - openjpeg=2.4.0=h3ad879b_0
-  - openssl=3.0.13=h7f8727e_0
-  - packaging=23.2=pyhd8ed1ab_0
-  - pcre2=10.42=hebb0a14_0
-  - pillow=10.2.0=py311h5eee18b_0
-  - pip=23.3.1=py311h06a4308_0
-  - platformdirs=4.2.0=pyhd8ed1ab_0
-  - pluggy=1.4.0=pyhd8ed1ab_0
-  - ply=3.11=py_1
-  - pooch=1.8.1=pyhd8ed1ab_0
-  - pygments=2.17.2=pyhd8ed1ab_0
-  - pyparsing=2.4.7=pyhd8ed1ab_1
-  - pyqt=5.15.10=py311h6a678d5_0
-  - pyqt5-sip=12.13.0=py311h5eee18b_0
-  - pytest=8.0.1=pyhd8ed1ab_1
-  - python=3.11.7=h955ad1f_0
-  - python-dateutil=2.8.2=pyhd8ed1ab_0
-  - python_abi=3.11=2_cp311
-  - pytorch=2.2.1=py3.11_cuda12.1_cudnn8.9.2_0
-  - pytorch-cuda=12.1=ha16c6d3_5
-  - pytorch-mutex=1.0=cuda
-  - pytz=2024.1=pyhd8ed1ab_0
-  - pyyaml=6.0.1=py311h5eee18b_0
-  - qt-main=5.15.2=h327a75a_7
-  - readline=8.2=h5eee18b_0
-  - requests=2.31.0=py311h06a4308_1
-  - scipy=1.11.4=py311h08b1b3b_0
-  - setuptools=68.2.2=py311h06a4308_0
-  - sip=6.7.12=py311h6a678d5_0
-  - six=1.16.0=pyh6c4a22f_0
-  - snowballstemmer=2.2.0=pyhd8ed1ab_0
-  - sphinx=7.2.6=pyhd8ed1ab_0
-  - sphinx-argparse=0.4.0=pyhd8ed1ab_0
-  - sphinx_rtd_theme=2.0.0=pyha770c72_0
-  - sphinxcontrib-applehelp=1.0.8=pyhd8ed1ab_0
-  - sphinxcontrib-devhelp=1.0.6=pyhd8ed1ab_0
-  - sphinxcontrib-htmlhelp=2.0.5=pyhd8ed1ab_0
-  - sphinxcontrib-jquery=4.1=pyhd8ed1ab_0
-  - sphinxcontrib-jsmath=1.0.1=pyhd8ed1ab_0
-  - sphinxcontrib-qthelp=1.0.7=pyhd8ed1ab_0
-  - sphinxcontrib-serializinghtml=1.1.10=pyhd8ed1ab_0
-  - sqlite=3.41.2=h5eee18b_0
-  - sympy=1.12=py311h06a4308_0
-  - tbb=2021.8.0=hdb19cb5_0
-  - tk=8.6.12=h1ccaba5_0
-  - tomli=2.0.1=pyhd8ed1ab_0
-  - torchaudio=2.2.1=py311_cu121
-  - torchtriton=2.2.0=py311
-  - torchvision=0.17.1=py311_cu121
-  - tornado=6.3.3=py311h5eee18b_0
-  - typing_extensions=4.9.0=py311h06a4308_1
-  - tzdata=2023d=h04d1e81_0
-  - urllib3=2.1.0=py311h06a4308_0
-  - wheel=0.41.2=py311h06a4308_0
-  - xz=5.4.5=h5eee18b_0
-  - yaml=0.2.5=h7b6447c_0
-  - zipp=3.17.0=pyhd8ed1ab_0
-  - zlib=1.2.13=h5eee18b_0
-  - zstd=1.5.5=hc292b87_0
-
diff --git a/setup.py b/setup.py
index 8ac2e50..4e62c0f 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@ setuptools.setup(
     description="Tools for coherent diffractive imaging and ptychography",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    url="https://github.mit.edu/scattering/CDTools.git",
+    url="https://github.com/cdtools-developers/cdtools",
     install_requires=[
         "numpy>=1.0",
         "scipy>=1.0",

From 36de44565baca95f6555d9b1ad00fdb63a059303 Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Mon, 7 Jul 2025 15:16:12 +0000
Subject: [PATCH 17/55] Set fixed RNG seed for reconstruction pytests

---
 tests/models/test_fancy_ptycho.py  | 2 ++
 tests/models/test_simple_ptycho.py | 4 ++++
 2 files changed, 6 insertions(+)

diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 7ef429f..3f78a73 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -5,6 +5,8 @@ import torch as t
 import cdtools
 from matplotlib import pyplot as plt
 
+# Force all reconstructions to use the same RNG seed
+t.manual_seed(0)
 
 def test_center_probe(lab_ptycho_cxi):
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py
index f74dd0d..f77c589 100644
--- a/tests/models/test_simple_ptycho.py
+++ b/tests/models/test_simple_ptycho.py
@@ -1,6 +1,10 @@
 import pytest
 import cdtools
 from matplotlib import pyplot as plt
+import torch as t
+
+# Force all reconstructions to use the same RNG seed
+t.manual_seed(0)
 
 @pytest.mark.slow
 def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):

From 93f5d1969dba85ef42ac00fc8864455f6d1a925f Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 12:53:06 -0700
Subject: [PATCH 18/55] Add .flake8 configuration to ignore line length errors
 (E501)

---
 .flake8 | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 .flake8

diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..16520fc
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,2 @@
+[flake8]
+ignore = E501
\ No newline at end of file

From 38be347a1895fe527ba8d443200286d664d173e0 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 13:38:52 -0700
Subject: [PATCH 19/55] lint conftest.py

---
 tests/conftest.py | 132 +++++++++++++++++++++++-----------------------
 1 file changed, 66 insertions(+), 66 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 781dd9e..042d991 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -47,6 +47,7 @@ def pytest_collection_modifyitems(config, items):
         if "slow" in item.keywords:
             item.add_marker(skip_slow)
 
+
 @pytest.fixture
 def reconstruction_device(request):
     return request.config.getoption("--reconstruction_device")
@@ -57,7 +58,6 @@ def show_plot(request):
     return request.config.getoption("--plot")
 
 
-
 @pytest.fixture(scope='module')
 def ptycho_cxi_1():
     """Creates an example file for CXI ptychography. This file is defined
@@ -67,11 +67,11 @@ def ptycho_cxi_1():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_1','w',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_1', 'w', driver='core', backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)
-    f.create_dataset('number_of_entries',data=1)
+    f.create_dataset('number_of_entries', data=1)
 
     # Then define a bunch of metadata for entry_1
     e1f = f.create_group('entry_1')
@@ -104,19 +104,19 @@ def ptycho_cxi_1():
     s1f['concentration'] = s1e['concentration']
     s1e['mass'] = np.float32(np.random.rand())
     s1f['mass'] = s1e['mass']
-    s1e['temperature'] = np.float32(np.random.rand()*100)
+    s1e['temperature'] = np.float32(np.random.rand() * 100)
     s1f['temperature'] = s1e['temperature']
-    s1e['thickness'] = np.float32(np.random.rand()*1e-7)
+    s1e['thickness'] = np.float32(np.random.rand() * 1e-7)
     s1f['thickness'] = s1e['thickness']
     s1e['unit_cell_volume'] = np.float32(np.random.rand() * 1e-27)
     s1f['unit_cell_volume'] = s1e['unit_cell_volume']
-    s1e['unit_cell'] = np.array([1,1,1,90,90,90]).astype(np.float32)
-    s1f.create_dataset('unit_cell',data = s1e['unit_cell'])
+    s1e['unit_cell'] = np.array([1, 1, 1, 90, 90, 90]).astype(np.float32)
+    s1f.create_dataset('unit_cell', data=s1e['unit_cell'])
 
     i1f = e1f.create_group('instrument_1')
     source1f = i1f.create_group('source_1')
 
-    energy = np.float32(1.3618e-16) #Joules, = 850 eV
+    energy = np.float32(1.3618e-16)  # Joules, = 850 eV
     source1f['energy'] = energy
     expected['wavelength'] = np.float32(1.9864459e-25) / energy
     source1f['wavelength'] = expected['wavelength']
@@ -126,48 +126,48 @@ def ptycho_cxi_1():
     d1e = expected['detector']
     d1e['distance'] = np.float32(0.3)
     d1f['distance'] = d1e['distance']
-    d1e['basis'] = np.array([[0,-30e-6,0],
-                             [-20e-6,0,0]]).astype(np.float32).transpose()
-    d1f.create_dataset('basis_vectors',data=d1e['basis'])
+    d1e['basis'] = np.array([[0, -30e-6, 0],
+                             [-20e-6, 0, 0]]).astype(np.float32).transpose()
+    d1f.create_dataset('basis_vectors', data=d1e['basis'])
     d1f['x_pixel_size'] = np.float32(20e-6)
     d1f['y_pixel_size'] = np.float32(30e-6)
-    d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
+    d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
     d1f.create_dataset('corner_position', data=d1e['corner'])
 
     # Remember the format for the CXI file differs from the format used
     # internally
-    mask = np.zeros((256,256)).astype(np.int32)
-    mask[5,8] = 1
-    expected['mask'] = np.ones((256,256)).astype(bool)
-    expected['mask'][5,8] = 0
-    d1f.create_dataset('mask',data=mask)
+    mask = np.zeros((256, 256)).astype(np.int32)
+    mask[5, 8] = 1
+    expected['mask'] = np.ones((256, 256)).astype(bool)
+    expected['mask'][5, 8] = 0
+    d1f.create_dataset('mask', data=mask)
 
     # There is no specification for this in the CXI file format :(
-    qe_mask = np.ones((256,256)).astype(np.float32)
+    qe_mask = np.ones((256, 256)).astype(np.float32)
     expected['qe_mask'] = qe_mask
-    d1f.create_dataset('qe_mask',data=qe_mask)
-    
+    d1f.create_dataset('qe_mask', data=qe_mask)
+
     # Create an initial background
-    dark = np.ones((256,256)) * 0.01
+    dark = np.ones((256, 256)) * 0.01
     expected['dark'] = dark
     d1f.create_dataset('data_dark', data=dark)
 
     data1f = e1f.create_group('data_1')
 
-    data = np.random.rand(100,256,256).astype(np.float32)
+    data = np.random.rand(100, 256, 256).astype(np.float32)
     expected['data'] = data
-    d1f.create_dataset('data',data=data)
+    d1f.create_dataset('data', data=data)
     data1f['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
 
     d1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
-    expected['axes'] = ['translation','y','x']
+    expected['axes'] = ['translation', 'y', 'x']
 
     g1f = s1f.create_group('geometry_1')
-    orientation = np.array([1.,0,0,0,1,0])
+    orientation = np.array([1., 0, 0, 0, 1, 0])
     g1f.create_dataset('orientation', data=orientation)
-    s1e['orientation'] = np.array([[1.,0,0],[0,1,0],[0,0,1]])
-    translations = np.arange(300).reshape((100,3)).astype(np.float32)
-    g1f.create_dataset('translation',data=translations)    
+    s1e['orientation'] = np.array([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
+    translations = np.arange(300).reshape((100, 3)).astype(np.float32)
+    g1f.create_dataset('translation', data=translations)
     data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
     d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
     expected['translations'] = -translations
@@ -193,11 +193,11 @@ def ptycho_cxi_2():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_2','w',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_2', 'w', driver='core', backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)
-    f.create_dataset('number_of_entries',data=1)
+    f.create_dataset('number_of_entries', data=1)
 
     # Then define a bunch of metadata for entry_1
     e1f = f.create_group('entry_1')
@@ -210,13 +210,13 @@ def ptycho_cxi_2():
     s1f = e1f.create_group('sample_1')
     expected['sample info'] = {}
     s1e = expected['sample info']
-    s1e['temperature'] = np.float32(np.random.rand()*100)
+    s1e['temperature'] = np.float32(np.random.rand() * 100)
     s1f['temperature'] = s1e['temperature']
 
     i1f = e1f.create_group('instrument_1')
     source1f = i1f.create_group('source_1')
 
-    energy = np.float32(1.3618e-16) #Joules, = 850 eV
+    energy = np.float32(1.3618e-16)  # Joules, = 850 eV
     expected['wavelength'] = np.float32(1.9864459e-25) / energy
     source1f['wavelength'] = expected['wavelength']
 
@@ -224,11 +224,11 @@ def ptycho_cxi_2():
     expected['detector'] = {}
     d1e = expected['detector']
     d1e['distance'] = np.float32(0.3)
-    d1e['basis'] = np.array([[0,-30e-6,0],
-                             [-20e-6,0,0]]).astype(np.float32).transpose()
+    d1e['basis'] = np.array([[0, -30e-6, 0],
+                             [-20e-6, 0, 0]]).astype(np.float32).transpose()
     d1f['x_pixel_size'] = np.float32(20e-6)
     d1f['y_pixel_size'] = np.float32(30e-6)
-    d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
+    d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
     d1f.create_dataset('corner_position', data=d1e['corner'])
 
     # Remember the format for the CXI file differs from the format used
@@ -236,24 +236,23 @@ def ptycho_cxi_2():
     expected['mask'] = None
 
     expected['qe_mask'] = None
-    
+
     # Test with a set of dark images
-    dark = np.ones((10,256,256)) * 0.01
-    expected['dark'] = np.nanmean(dark,axis=0)
+    dark = np.ones((10, 256, 256)) * 0.01
+    expected['dark'] = np.nanmean(dark, axis=0)
     d1f.create_dataset('data_dark', data=dark)
 
+    e1f.create_group('data_1')
 
-    data1f = e1f.create_group('data_1')
-
-    data = np.random.rand(100,256,256).astype(np.float32)
+    data = np.random.rand(100, 256, 256).astype(np.float32)
     expected['data'] = data
-    d1f.create_dataset('data',data=data)
+    d1f.create_dataset('data', data=data)
 
     expected['axes'] = None
 
     g1f = s1f.create_group('geometry_1')
-    translations = np.arange(300).reshape((100,3)).astype(np.float32)
-    g1f.create_dataset('translation',data=translations)
+    translations = np.arange(300).reshape((100, 3)).astype(np.float32)
+    g1f.create_dataset('translation', data=translations)
     expected['translations'] = -translations
 
     yield f, expected
@@ -276,11 +275,11 @@ def ptycho_cxi_3():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_3','w',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_3', 'w', driver='core', backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)
-    f.create_dataset('number_of_entries',data=1)
+    f.create_dataset('number_of_entries', data=1)
 
     # Then define a bunch of metadata for entry_1
     e1f = f.create_group('entry_1')
@@ -297,7 +296,7 @@ def ptycho_cxi_3():
     i1f = e1f.create_group('instrument_1')
     source1f = i1f.create_group('source_1')
 
-    energy = np.float32(1.3618e-16) #Joules, = 850 eV
+    energy = np.float32(1.3618e-16)  # Joules, = 850 eV
     source1f['energy'] = energy
     expected['wavelength'] = np.float32(1.9864459e-25) / energy
 
@@ -306,41 +305,41 @@ def ptycho_cxi_3():
     d1e = expected['detector']
     d1e['distance'] = np.float32(0.3)
     d1f['distance'] = d1e['distance']
-    d1e['basis'] = np.array([[0,-30e-6,0],
-                             [-20e-6,0,0]]).astype(np.float32).transpose()
-    d1f.create_dataset('basis_vectors',data=d1e['basis'])
+    d1e['basis'] = np.array([[0, -30e-6, 0],
+                             [-20e-6, 0, 0]]).astype(np.float32).transpose()
+    d1f.create_dataset('basis_vectors', data=d1e['basis'])
     d1e['corner'] = None
 
     # Remember the format for the CXI file differs from the format used
     # internally
-    mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
-    mask[15,47] = 38
-    expected['mask'] = np.ones((256,256)).astype(bool)
-    expected['mask'][15,47] = 0
-    d1f.create_dataset('mask',data=mask)
+    mask = np.ones((256, 256)).astype(np.uint32) * 0x00001000
+    mask[15, 47] = 38
+    expected['mask'] = np.ones((256, 256)).astype(bool)
+    expected['mask'][15, 47] = 0
+    d1f.create_dataset('mask', data=mask)
 
     expected['qe_mask'] = None
-    
+
     expected['dark'] = None
-    
+
     data1f = e1f.create_group('data_1')
 
-    data = np.random.rand(100,256,256).astype(np.float32)
+    data = np.random.rand(100, 256, 256).astype(np.float32)
     expected['data'] = data
-    data1f.create_dataset('data',data=data)
+    data1f.create_dataset('data', data=data)
 
     data1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
-    expected['axes'] = ['translation','y','x']
+    expected['axes'] = ['translation', 'y', 'x']
 
-    translations = np.arange(300).reshape((100,3)).astype(np.float32)
-    data1f.create_dataset('translation',data=translations)
+    translations = np.arange(300).reshape((100, 3)).astype(np.float32)
+    data1f.create_dataset('translation', data=translations)
     expected['translations'] = -translations
 
     yield f, expected
 
     f.close()
 
-    
+
 @pytest.fixture(scope='module')
 def polarized_ptycho_cxi(ptycho_cxi_1):
     f, expected = ptycho_cxi_1
@@ -352,7 +351,7 @@ def polarized_ptycho_cxi(ptycho_cxi_1):
     data1f.create_dataset('polarizer_angle', data=expected['polarizer_angle'])
 
     yield f, expected
-    
+
 
 # As specific issues start to crop up with loading CXI files from different
 # beamlines, put a fixture here that replicates the issue so that we can
@@ -374,6 +373,7 @@ def gold_ball_cxi(pytestconfig):
     return str(pytestconfig.rootpath) + \
         '/examples/example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
 
+
 @pytest.fixture(scope='module')
 def lab_ptycho_cxi(pytestconfig):
     return str(pytestconfig.rootpath) + \
@@ -382,8 +382,8 @@ def lab_ptycho_cxi(pytestconfig):
 
 @pytest.fixture(scope='module')
 def example_nested_dicts(pytestconfig):
-    example_tensor = t.as_tensor(np.array([1,4.5,7]))
-    example_array = np.ones([10,20,30])
+    example_tensor = t.as_tensor(np.array([1, 4.5, 7]))
+    example_array = np.ones([10, 20, 30])
     example_scalar = 4.5
     example_single_element_array = np.array([0.3])
     example_string = 'testing'

From 4d812db68a07a4e9caa453cb62359efcf2e9a773 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 13:39:44 -0700
Subject: [PATCH 20/55] restructured conftest.py imports

---
 tests/conftest.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 042d991..850935d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,9 +1,10 @@
-import numpy as np
-import torch as t
-import h5py
-import pytest
 import datetime
 
+import h5py
+import numpy as np
+import pytest
+import torch as t
+
 
 #
 #

From f463144347148997365852b9cf50274098c552fa Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 13:49:46 -0700
Subject: [PATCH 21/55] linting test_datasets.py

---
 tests/test_datasets.py | 197 +++++++++++++++++++----------------------
 1 file changed, 89 insertions(+), 108 deletions(-)

diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 8999117..8b1d827 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -18,20 +18,20 @@ from cdtools.tools import data as cdtdata
 
 def test_CDataset_init():
     entry_info = {'start_time': datetime.datetime.now(),
-                  'title' : 'A simple test'}
+                  'title': 'A simple test'}
     sample_info = {'name': 'A test sample',
-                   'mass' : 3.4,
-                   'unit_cell' : np.array([1,1,1,87,84.5,90])}
+                   'mass': 3.4,
+                   'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
     wavelength = 1e-9
     detector_geometry = {'distance': 0.7,
-                         'basis': np.array([[0,-30e-6,0],
-                                            [-20e-6,0,0]]).transpose(),
-                         'corner': np.array((2550e-6,3825e-6,0.3))}
-    mask = np.ones((256,256))
+                         'basis': np.array([[0, -30e-6, 0],
+                                            [-20e-6, 0, 0]]).transpose(),
+                         'corner': np.array((2550e-6, 3825e-6, 0.3))}
+    mask = np.ones((256, 256))
     dataset = CDataset(entry_info, sample_info,
                        wavelength, detector_geometry, mask)
 
-    assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
+    assert t.all(t.eq(dataset.mask, t.tensor(mask.astype(bool))))
     assert dataset.entry_info == entry_info
     assert dataset.sample_info == sample_info
     assert dataset.wavelength == wavelength
@@ -53,7 +53,7 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
         else:
             assert dataset.sample_info is not None
 
-        assert np.isclose(dataset.wavelength,expected['wavelength'])
+        assert np.isclose(dataset.wavelength, expected['wavelength'])
 
         # Just check one of the loaded attributes
         assert np.isclose(dataset.detector_geometry['distance'],
@@ -64,18 +64,16 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
             assert 'corner' in dataset.detector_geometry
 
         if expected['mask'] is not None:
-            assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
+            assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
 
         if expected['qe_mask'] is not None:
-            assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
+            assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
 
         if expected['dark'] is not None:
             assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
                               dataset.background))
 
-            
 
-            
 def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
     for cxi, expected in test_ptycho_cxis:
         dataset = CDataset.from_cxi(cxi)
@@ -95,27 +93,24 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
 
         assert np.isclose(dataset.wavelength, read_dataset.wavelength)
 
-        
-         # Just check one of the loaded attributes
+        # Just check one of the loaded attributes
         assert np.isclose(dataset.detector_geometry['distance'],
                           read_dataset.detector_geometry['distance'])
         # Check that the other ones are loaded but not for fidelity
         assert 'basis' in read_dataset.detector_geometry
         if dataset.detector_geometry['corner'] is not None:
             assert 'corner' in read_dataset.detector_geometry
-            
-        
+
         if dataset.mask is not None:
-            assert t.all(t.eq(dataset.mask,read_dataset.mask))
+            assert t.all(t.eq(dataset.mask, read_dataset.mask))
 
         if dataset.qe_mask is not None:
-            assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
+            assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
 
         if dataset.background is not None:
             assert t.all(t.eq(dataset.background, read_dataset.background))
 
 
-
 def test_CDataset_to(ptycho_cxi_1):
     dataset = CDataset.from_cxi(ptycho_cxi_1[0])
 
@@ -136,28 +131,28 @@ def test_CDataset_to(ptycho_cxi_1):
 
 def test_Ptycho2DDataset_init():
     entry_info = {'start_time': datetime.datetime.now(),
-                  'title' : 'A simple test'}
+                  'title': 'A simple test'}
     sample_info = {'name': 'A test sample',
-                   'mass' : 3.4,
-                   'unit_cell' : np.array([1,1,1,87,84.5,90])}
+                   'mass': 3.4,
+                   'unit_cell': np.array([1, 1, 1, 87, 84.5, 90])}
     wavelength = 1e-9
     detector_geometry = {'distance': 0.7,
-                         'basis': np.array([[0,-30e-6,0],
-                                            [-20e-6,0,0]]).transpose(),
-                         'corner': np.array((2550e-6,3825e-6,0.3))}
-    mask = np.ones((256,256))
-    qe_mask = 1.2*np.ones((256,256), dtype=np.float32)
-    patterns = np.random.rand(20,256,256)
-    translations = np.random.rand(20,3)
-    
-    dataset = Ptycho2DDataset(translations, patterns,
-                                entry_info=entry_info,
-                                sample_info=sample_info,
-                                wavelength=wavelength,
-                                detector_geometry=detector_geometry,
-                                mask=mask)
+                         'basis': np.array([[0, -30e-6, 0],
+                                            [-20e-6, 0, 0]]).transpose(),
+                         'corner': np.array((2550e-6, 3825e-6, 0.3))}
+    mask = np.ones((256, 256))
+    qe_mask = 1.2 * np.ones((256, 256), dtype=np.float32)
+    patterns = np.random.rand(20, 256, 256)
+    translations = np.random.rand(20, 3)
 
-    assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
+    dataset = Ptycho2DDataset(translations, patterns,
+                              entry_info=entry_info,
+                              sample_info=sample_info,
+                              wavelength=wavelength,
+                              detector_geometry=detector_geometry,
+                              mask=mask)
+
+    assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
     assert dataset.entry_info == entry_info
     assert dataset.sample_info == sample_info
     assert dataset.wavelength == wavelength
@@ -167,15 +162,15 @@ def test_Ptycho2DDataset_init():
 
     # Also test one with a qe_mask
     dataset = Ptycho2DDataset(translations, patterns,
-                                entry_info=entry_info,
-                                sample_info=sample_info,
-                                wavelength=wavelength,
-                                detector_geometry=detector_geometry,
-                                mask=mask,
-                                qe_mask=qe_mask)
+                              entry_info=entry_info,
+                              sample_info=sample_info,
+                              wavelength=wavelength,
+                              detector_geometry=detector_geometry,
+                              mask=mask,
+                              qe_mask=qe_mask)
 
-    assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
-    assert t.all(t.eq(dataset.qe_mask,t.as_tensor(qe_mask)))
+    assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
+    assert t.all(t.eq(dataset.qe_mask, t.as_tensor(qe_mask)))
     assert dataset.entry_info == entry_info
     assert dataset.sample_info == sample_info
     assert dataset.wavelength == wavelength
@@ -199,7 +194,7 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
         else:
             assert dataset.sample_info is not None
 
-        assert np.isclose(dataset.wavelength,expected['wavelength'])
+        assert np.isclose(dataset.wavelength, expected['wavelength'])
 
         # Just check one of the loaded attributes
         assert np.isclose(dataset.detector_geometry['distance'],
@@ -210,18 +205,17 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
             assert 'corner' in dataset.detector_geometry
 
         if expected['mask'] is not None:
-            assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
+            assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
 
         if expected['qe_mask'] is not None:
-            assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
+            assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
 
         if expected['dark'] is not None:
             assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
                               dataset.background))
 
-            
-        assert t.allclose(t.tensor(expected['data']),dataset.patterns)
-        assert t.allclose(t.tensor(expected['translations']),dataset.translations)
+        assert t.allclose(t.tensor(expected['data']), dataset.patterns)
+        assert t.allclose(t.tensor(expected['translations']), dataset.translations)
 
 
 def test_Ptycho2DDataset_from_cxi_64bit(test_ptycho_cxis):
@@ -279,20 +273,19 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
 
         assert np.isclose(dataset.wavelength, read_dataset.wavelength)
 
-        
-         # Just check one of the loaded attributes
+        # Just check one of the loaded attributes
         assert np.isclose(dataset.detector_geometry['distance'],
                           read_dataset.detector_geometry['distance'])
         # Check that the other ones are loaded but not for fidelity
         assert 'basis' in read_dataset.detector_geometry
         if dataset.detector_geometry['corner'] is not None:
             assert 'corner' in read_dataset.detector_geometry
-            
+
         if dataset.mask is not None:
-            assert t.all(t.eq(dataset.mask,read_dataset.mask))
+            assert t.all(t.eq(dataset.mask, read_dataset.mask))
 
         if dataset.qe_mask is not None:
-            assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
+            assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
 
         if dataset.background is not None:
             assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -303,7 +296,6 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
 
 def test_Ptycho2DDataset_to(ptycho_cxi_1):
     dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
-    
     dataset.to(dtype=t.float64)
     assert dataset.mask.dtype == t.bool
     assert dataset.qe_mask.dtype == t.float64
@@ -327,8 +319,8 @@ def test_Ptycho2DDataset_ops(ptycho_cxi_1):
     assert len(dataset) == expected['data'].shape[0]
     (idx, translation), pattern = dataset[3]
     assert idx == 3
-    assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
-    assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
+    assert t.allclose(translation, t.tensor(expected['translations'][3, :]))
+    assert t.allclose(pattern, t.tensor(expected['data'][3, :, :]))
 
 
 def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
@@ -341,12 +333,12 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
         (idx, translation), pattern = dataset[3]
         assert str(translation.device) == 'cuda:0'
         assert str(pattern.device) == 'cuda:0'
-        
+
         assert idx == 3
         assert t.allclose(translation.to(device='cpu'),
-                          t.tensor(expected['translations'][3,:]))
+                          t.tensor(expected['translations'][3, :]))
         assert t.allclose(pattern.to(device='cpu'),
-                          t.tensor(expected['data'][3,:,:]))
+                          t.tensor(expected['data'][3, :, :]))
 
 
 def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
@@ -365,18 +357,15 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
         masked_patterns = dataset.mask * dataset.patterns
         assert t.allclose(
             copied_dataset.patterns,
-            masked_patterns[:,::2,::2] +
-            masked_patterns[:,1::2,::2] +
-            masked_patterns[:,::2,1::2] +
-            masked_patterns[:,1::2,1::2]
+            masked_patterns[:, ::2, ::2] + masked_patterns[:, 1::2, ::2] + masked_patterns[:, ::2, 1::2] + masked_patterns[:, 1::2, 1::2]
         )
 
         if dataset.qe_mask is None:
             manually_downsampled_mask = t.logical_and(
-                t.logical_and(dataset.mask[::2,::2],
-                              dataset.mask[1::2,::2]),
-                t.logical_and(dataset.mask[::2,1::2],
-                              dataset.mask[1::2,1::2])
+                t.logical_and(dataset.mask[::2, ::2],
+                              dataset.mask[1::2, ::2]),
+                t.logical_and(dataset.mask[::2, 1::2],
+                              dataset.mask[1::2, 1::2])
             )
             assert t.allclose(
                 copied_dataset.mask,
@@ -384,10 +373,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
             )
         else:
             manually_downsampled_mask = t.logical_or(
-                t.logical_or(dataset.mask[::2,::2],
-                             dataset.mask[1::2,::2]),
-                t.logical_or(dataset.mask[::2,1::2],
-                             dataset.mask[1::2,1::2])
+                t.logical_or(dataset.mask[::2, ::2],
+                             dataset.mask[1::2, ::2]),
+                t.logical_or(dataset.mask[::2, 1::2],
+                             dataset.mask[1::2, 1::2])
             )
             assert t.allclose(
                 copied_dataset.mask,
@@ -396,8 +385,7 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
 
             masked_qe_mask = dataset.mask * dataset.qe_mask
             manually_downsampled_qe_mask = (
-                masked_qe_mask[::2,::2] + masked_qe_mask[1::2,::2]
-                + masked_qe_mask[::2,1::2] + masked_qe_mask[1::2,1::2]
+                masked_qe_mask[::2, ::2] + masked_qe_mask[1::2, ::2] + masked_qe_mask[::2, 1::2] + masked_qe_mask[1::2, 1::2]
             ) / 4
 
             assert t.allclose(
@@ -405,15 +393,11 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
                 manually_downsampled_qe_mask
             )
 
-        
         if dataset.background is not None:
             assert t.allclose(
                 copied_dataset.background,
-                dataset.background[::2,::2] +
-                dataset.background[1::2,::2] +
-                dataset.background[::2,1::2] +
-                dataset.background[1::2,1::2]
-        )
+                dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
+            )
 
         # And then we just test the shape for a few factors, and check that
         # it doesn't fail on edge cases (e.g. factor=1)
@@ -428,10 +412,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
 
             assert np.allclose(expected_pattern_shape,
                                np.array(copied_dataset.patterns.shape))
-            
+
             assert np.allclose(np.array(dataset.mask.shape) // factor,
                                np.array(copied_dataset.mask.shape))
-            
+
             if dataset.background is not None:
                 assert np.allclose(np.array(dataset.background.shape) // factor,
                                    np.array(copied_dataset.background.shape))
@@ -468,13 +452,13 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
     copied_dataset = deepcopy(dataset)
 
     # Test 1: Complain when the the bounds of an ROI are correctly defined,
-    #   but it does not contain any sample positions inside of it. The 
-    #   translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y 
-    #   (it looks like a line scan). We select an ROI at (0, -300) which should 
+    #   but it does not contain any sample positions inside of it. The
+    #   translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y
+    #   (it looks like a line scan). We select an ROI at (0, -300) which should
     #   not contain any translation positions.
     with pytest.raises(ValueError) as excinfo:
-        copied_dataset.crop_translations(roi=(0,-5,-295,-300))
-    assert('(i.e., patterns and translations will be empty)') in str(excinfo.value)
+        copied_dataset.crop_translations(roi=(0, -5, -295, -300))
+    assert '(i.e., patterns and translations will be empty)' in str(excinfo.value)
 
     # Test 2: Draw an ROI that's centered in the middle of the x/y translation range
     #   and make sure that the first and last x/y elements in dataset.translate
@@ -482,40 +466,37 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
 
     #   Make tuples that will store the x and y positions. This will be used for
     #   making permutations of the x/y positions in the ROI later.
-    x_left, y_top = dataset.translations[10,:2]
-    x_right, y_bottom = dataset.translations[-11,:2]
+    x_left, y_top = dataset.translations[10, :2]
+    x_right, y_bottom = dataset.translations[-11, :2]
 
     x_permutations = ((x_left, x_right), (x_right, x_left))
     y_permutations = ((y_top, y_bottom), (y_bottom, y_top))
-    roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in
-                              itertools.product(x_permutations, y_permutations))
+    roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in itertools.product(x_permutations, y_permutations))
 
     #   Get the dataset
     copied_dataset.crop_translations(roi=roi_permutations[0])
 
     #   Execute the actual test
-    assert (copied_dataset.translations[0,0] in x_permutations[0]) and \
-        (copied_dataset.translations[-1,0] in x_permutations[0])
-    
-    assert (copied_dataset.translations[0,1] in y_permutations[0]) and \
-        (copied_dataset.translations[-1,1] in y_permutations[0])
+    assert (copied_dataset.translations[0, 0] in x_permutations[0]) and \
+        (copied_dataset.translations[-1, 0] in x_permutations[0])
+
+    assert (copied_dataset.translations[0, 1] in y_permutations[0]) and \
+        (copied_dataset.translations[-1, 1] in y_permutations[0])
 
     # Test 3: Check if the shape of dataset.patterns and dataset.translate is correct
     #   (designed to be 20 fewer rows here)
     #   In the future, this should include a check for dataset.intensities once
     #   an appropriate cxi file is set up for conftest.
-    expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], 
-                                            dataset.patterns.shape[-2:]])
-    
-    expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], 
-                                                dataset.translations.shape[-1:]])
+    expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], dataset.patterns.shape[-2:]])
+
+    expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], dataset.translations.shape[-1:]])
 
     assert np.allclose(np.array(copied_dataset.patterns.shape), expected_patterns_shape)
-        
+
     assert np.allclose(np.array(copied_dataset.translations.shape), expected_translations_shape)
 
     # Test 4: Make sure that we always get the same result no matter what order we
-    #   define the left/right and bottom/top values in roi, provided that roi[:2] 
+    #   define the left/right and bottom/top values in roi, provided that roi[:2]
     #   and roi[2:] correspond with the x and y coordinates, respectively.
 
     #   Check each permutation
@@ -526,6 +507,6 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
         copied_dataset.crop_translations(roi=roi)
 
         # Check if the contents of dataset.patterns and dataset.translate is correct
-        assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10,:])
+        assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10, :])
 
-        assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])
+        assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])

From 745f7b73e6bafeae7303675436680e5e7f211fef Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 13:55:54 -0700
Subject: [PATCH 22/55] Fix background calculation in
 test_Ptycho2DDataset_downsample

---
 tests/test_datasets.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 8b1d827..4d17f0b 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -396,7 +396,7 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
         if dataset.background is not None:
             assert t.allclose(
                 copied_dataset.background,
-                dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
+                dataset.background[::2, ::2] + dataset.background[1::2, ::2] + dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
             )
 
         # And then we just test the shape for a few factors, and check that

From ecb8ee4864378875bffed3c7fd4b85ae78d29b6b Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 13:59:35 -0700
Subject: [PATCH 23/55] linting test_simple_ptycho.py

---
 tests/models/test_simple_ptycho.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py
index f77c589..b6b1868 100644
--- a/tests/models/test_simple_ptycho.py
+++ b/tests/models/test_simple_ptycho.py
@@ -1,17 +1,18 @@
 import pytest
-import cdtools
-from matplotlib import pyplot as plt
 import torch as t
 
+import cdtools
+
 # Force all reconstructions to use the same RNG seed
 t.manual_seed(0)
 
+
 @pytest.mark.slow
 def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
 
     model = cdtools.models.SimplePtycho.from_dataset(dataset)
-    
+
     model.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
@@ -19,10 +20,10 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
-            
+
     if show_plot:
         model.inspect(dataset)
         model.compare(dataset)
-        
+
     # If this fails, the reconstruction got worse
     assert model.loss_history[-1] < 0.013

From 96d8b10792da152b3aaa6babbba6e18808c59730 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 14:11:50 -0700
Subject: [PATCH 24/55] linting test_fancy_ptycho.py

---
 tests/models/test_fancy_ptycho.py | 34 +++++++++++++++----------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 3f78a73..e5422d7 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -1,13 +1,12 @@
 import pytest
-import cdtools
 import torch as t
 
 import cdtools
-from matplotlib import pyplot as plt
 
 # Force all reconstructions to use the same RNG seed
 t.manual_seed(0)
 
+
 def test_center_probe(lab_ptycho_cxi):
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
     model = cdtools.models.FancyPtycho.from_dataset(
@@ -28,8 +27,8 @@ def test_center_probe(lab_ptycho_cxi):
     fourier_model.probe.data = cdtools.tools.propagators.far_field(
         base_probe
     )
-    
-    fourier_base_probe = fourier_model.probe.detach().clone()
+
+    fourier_model.probe.detach().clone()
     fourier_model.center_probes()
     fourier_centered_probe = fourier_model.probe.detach().clone()
     ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
@@ -45,26 +44,27 @@ def test_center_probe(lab_ptycho_cxi):
         atol=1e-4,
         rtol=1e-3
     )
-    
+
+
 @pytest.mark.slow
 def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 
     print('\nTesting performance on the standard transmission ptycho dataset')
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
-    
+
     model = cdtools.models.FancyPtycho.from_dataset(
         dataset,
-        n_modes=3, 
+        n_modes=3,
         oversampling=2,
         exponentiate_obj=True,
         dm_rank=2,
         probe_support_radius=120,
-        propagation_distance=5e-3, 
-        units='mm', 
+        propagation_distance=5e-3,
+        units='mm',
         obj_view_crop=-50,
-        use_qe_mask=True, # test this in the case where no qe mask is defined
+        use_qe_mask=True,  # test this in the case where no qe mask is defined
     )
-    
+
     print('Running reconstruction on provided reconstruction_device,',
           reconstruction_device)
     model.to(device=reconstruction_device)
@@ -75,11 +75,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
 
-    for loss in model.Adam_optimize(50, dataset,  lr=0.005, batch_size=50):
+    for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
-            
+
     model.tidy_probes()
 
     if show_plot:
@@ -94,7 +94,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
 
     print('\nTesting performance on the standard gold balls dataset')
-    
+
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
 
     pad = 10
@@ -119,7 +119,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
           reconstruction_device)
     model.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
-    
+
     for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
@@ -135,7 +135,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
-            
+
     model.tidy_probes()
 
     if show_plot:
@@ -146,5 +146,3 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     # and choosing a rough value. If it triggers this assertion error,
     # something changed to make the final quality worse!
     assert model.loss_history[-1] < 0.0001
-
-

From 58246966f4c2aa2c47dafa0aadfba54d238c3893 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 14:21:17 -0700
Subject: [PATCH 25/55] Add W503 to flake8 ignore list, line break before
 binary operator

---
 .flake8 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.flake8 b/.flake8
index 16520fc..9214cb9 100644
--- a/.flake8
+++ b/.flake8
@@ -1,2 +1,2 @@
 [flake8]
-ignore = E501
\ No newline at end of file
+ignore = E501, W503
\ No newline at end of file

From b00f39e1fe5b028d0f0942228c74f9d2289b719d Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 14:28:23 -0700
Subject: [PATCH 26/55] linting test_analysis.py and test_data.py

---
 tests/tools/test_analysis.py | 477 +++++++++++++++++++----------------
 tests/tools/test_data.py     | 163 ++++++------
 2 files changed, 336 insertions(+), 304 deletions(-)

diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index f0a1cf1..ad06757 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -1,8 +1,7 @@
 import numpy as np
+import torch as t
 from scipy import linalg as la
 from scipy.sparse import linalg as spla
-import torch as t
-from itertools import combinations
 
 from cdtools.tools import analysis, initializers
 
@@ -15,10 +14,10 @@ def test_product_svd():
     A = np.random.rand(*shape_A) + 1j * np.random.rand(*shape_A)
     B = np.random.rand(*shape_B) + 1j * np.random.rand(*shape_B)
 
-    AB = np.matmul(A,B)
+    AB = np.matmul(A, B)
     U_1, S_1, Vh_1 = t.linalg.svd(t.as_tensor(AB), full_matrices=False)
 
-    U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A),t.as_tensor(B))
+    U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A), t.as_tensor(B))
     check_AB = U_2 @ t.diag_embed(S_2).to(dtype=Vh_2.dtype) @ Vh_2
 
     # So, at a minimum, U S Vh = AB
@@ -28,23 +27,23 @@ def test_product_svd():
     # singular vector, so all we can ask for in the comparison is that the
     # magnitudes here are
     assert np.allclose(S_1[:rank].numpy(), S_2.numpy())
-    prod_U = U_1[:,:rank].transpose(0,1).conj() @ U_2
-    prod_Vh = Vh_1[:rank,:] @ Vh_2.transpose(0,1).conj()
+    prod_U = U_1[:, :rank].transpose(0, 1).conj() @ U_2
+    prod_Vh = Vh_1[:rank, :] @ Vh_2.transpose(0, 1).conj()
     assert np.allclose(t.abs(prod_U).numpy(), np.eye(rank))
     assert np.allclose(t.abs(prod_Vh).numpy(), np.eye(rank))
     # Confirms that the phases are consistent between the two, I think
     # it's redundant with the first check but I'm not sure
     assert np.allclose(prod_Vh.numpy(), prod_U.numpy())
-    
+
     # test with numpy
-    U_3, S_3, Vh_3 = analysis.product_svd(A,B)
+    U_3, S_3, Vh_3 = analysis.product_svd(A, B)
     assert isinstance(U_3, np.ndarray)
     assert isinstance(S_3, np.ndarray)
     assert isinstance(Vh_3, np.ndarray)
-    
+
     assert np.allclose(S_1[:rank].numpy(), S_3)
-    prod_U = U_1[:,:rank].transpose(0,1).numpy().conj() @ U_3
-    prod_Vh = Vh_1[:rank,:].numpy() @ Vh_3.transpose().conj()
+    prod_U = U_1[:, :rank].transpose(0, 1).numpy().conj() @ U_3
+    prod_Vh = Vh_1[:rank, :].numpy() @ Vh_3.transpose().conj()
     assert np.allclose(np.abs(prod_U), np.eye(rank))
     assert np.allclose(np.abs(prod_Vh), np.eye(rank))
     # Confirms that the phases are consistent between the two, I think
@@ -55,42 +54,49 @@ def test_product_svd():
 def test_orthogonalize_probes():
 
     op = analysis.orthogonalize_probes
-    
+
     probe_xs = np.arange(64) - 32
     probe_ys = np.arange(76) - 38
     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 + 1j)),
-                       3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
-                       1*np.exp(-probe_Rs**2 / (2 * 15**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)),
+        ]
+    )
 
     weight_matrix_none = None
-    weight_matrix_single = np.random.randn(1,3) + 1j * np.random.randn(1,3)
-    weight_matrix_small = np.random.randn(2,3) + 1j * np.random.randn(2,3)
-    weight_matrix_medium = np.random.randn(3,3) + 1j * np.random.randn(3,3)
-    weight_matrix_large = np.random.randn(7,3) + 1j * np.random.randn(7,3)
+    weight_matrix_single = np.random.randn(1, 3) + 1j * np.random.randn(1, 3)
+    weight_matrix_small = np.random.randn(2, 3) + 1j * np.random.randn(2, 3)
+    weight_matrix_medium = np.random.randn(3, 3) + 1j * np.random.randn(3, 3)
+    weight_matrix_large = np.random.randn(7, 3) + 1j * np.random.randn(7, 3)
 
     weight_matrices = [
         weight_matrix_none,
         weight_matrix_single,
         weight_matrix_small,
         weight_matrix_medium,
-        weight_matrix_large
+        weight_matrix_large,
     ]
-    
+
     for weight_matrix in weight_matrices:
-        ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix,
-                                     return_reexpressed_weights=True)
+        ortho_probes_np, rwm_np = op(
+            probes, weight_matrix=weight_matrix, return_reexpressed_weights=True
+        )
         assert isinstance(ortho_probes_np, np.ndarray)
         assert isinstance(rwm_np, np.ndarray)
-        
-        probes_t = t.as_tensor(probes)
-        wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
-                else weight_matrix)
 
-        ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t,
-                                   return_reexpressed_weights=True)
+        probes_t = t.as_tensor(probes)
+        wm_t = (
+            t.as_tensor(weight_matrix) if weight_matrix is not None else weight_matrix
+        )
+
+        ortho_probes_t, rwm_t = op(
+            probes_t, weight_matrix=wm_t, return_reexpressed_weights=True
+        )
         assert t.is_tensor(ortho_probes_t)
         assert t.is_tensor(rwm_t)
 
@@ -105,15 +111,17 @@ def test_orthogonalize_probes():
             realized_probes = probes
 
         calculated_probes = np.tensordot(rwm_np, ortho_probes_np, axes=1)
-            
+
         assert np.allclose(realized_probes, calculated_probes)
-        
+
         # Now we test if the orthogonalized probes are orthogonalized
         reshaped_probes = ortho_probes_np.reshape(
-            (ortho_probes_np.shape[0],
-             ortho_probes_np.shape[1] * ortho_probes_np.shape[2]))
-        products = np.matmul(reshaped_probes,
-                             reshaped_probes.conj().transpose())
+            (
+                ortho_probes_np.shape[0],
+                ortho_probes_np.shape[1] * ortho_probes_np.shape[2],
+            )
+        )
+        products = np.matmul(reshaped_probes, reshaped_probes.conj().transpose())
 
         if weight_matrix is not None:
             output_nmodes = min(weight_matrix.shape[0], probes.shape[0])
@@ -121,23 +129,22 @@ def test_orthogonalize_probes():
             output_nmodes = probes.shape[0]
 
         for i in range(output_nmodes):
-            for j in range(i+1, output_nmodes):
-                assert np.isclose(products[i,j], 0)
+            for j in range(i + 1, output_nmodes):
+                assert np.isclose(products[i, j], 0)
 
         # And now we test if they multiply to the same density matrix as
         # the original probes + weight matrix
         reshaped_realized_probes = realized_probes.reshape(
-            (realized_probes.shape[0],
-             realized_probes.shape[1] * realized_probes.shape[2]))
-        
+            (
+                realized_probes.shape[0],
+                realized_probes.shape[1] * realized_probes.shape[2],
+            )
+        )
+
         dm_original = np.matmul(
-            reshaped_realized_probes.conj().transpose(),
-            reshaped_realized_probes
-        )
-        dm_output = np.matmul(
-            reshaped_probes.conj().transpose(),
-            reshaped_probes
+            reshaped_realized_probes.conj().transpose(), reshaped_realized_probes
         )
+        dm_output = np.matmul(reshaped_probes.conj().transpose(), reshaped_probes)
         assert np.allclose(dm_original, dm_output)
 
         # And finally, we confirm that what we have are the eigenvectors/values
@@ -150,53 +157,58 @@ def test_orthogonalize_probes():
         # is undefined
         assert np.allclose(np.abs(cross_products), np.abs(products))
 
-        
+
 def test_standardize():
 
     # Start by making a probe and object that should meet the standardization
     # conditions
-    probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
-    probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
+    probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
+    probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
     probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
 
-    assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
+    assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
     assert np.angle(np.sum(probe)) < 2e-7
 
-    obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
-    obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
-                      (obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
+    obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
+    obj_slice = np.s_[
+        (obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
+        (obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
+    ]
 
     obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
-    assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
+    assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
 
-    
     # Then make a nonstandard version of them and standardize it
     # First, don't add a phase ramp and test
-    test_probe = probe * 37.6 * np.exp(1j*0.35)
-    test_obj = obj / 37.6 * np.exp(1j*1.43)
+    test_probe = probe * 37.6 * np.exp(1j * 0.35)
+    test_obj = obj / 37.6 * np.exp(1j * 1.43)
     s_probe, s_obj = analysis.standardize(test_probe, test_obj)
     assert np.allclose(probe, s_probe)
     assert np.allclose(obj, s_obj)
 
     # Test that it works on torch tensors
-    s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64))
+    s_probe, s_obj = analysis.standardize(
+        t.as_tensor(test_probe, dtype=t.complex64),
+        t.as_tensor(test_obj, dtype=t.complex64),
+    )
     s_probe = s_probe.numpy()
     s_obj = s_obj.numpy()
     assert np.allclose(probe, s_probe)
     assert np.allclose(obj, s_obj)
 
-    
     # Then do one with a phase ramp
-    phase_ramp_dir = (np.random.rand(2) - 0.5)
+    phase_ramp_dir = np.random.rand(2) - 0.5
 
-    probe_Xs, probe_Ys = np.mgrid[:probe.shape[0],:probe.shape[1]]
-    phase_ramp = np.exp(1j*probe_Ys * phase_ramp_dir[1]+
-                        1j*probe_Xs * phase_ramp_dir[0])
+    probe_Xs, probe_Ys = np.mgrid[: probe.shape[0], : probe.shape[1]]
+    phase_ramp = np.exp(
+        1j * probe_Ys * phase_ramp_dir[1] + 1j * probe_Xs * phase_ramp_dir[0]
+    )
     test_probe = test_probe * phase_ramp
 
-    obj_Xs, obj_Ys = np.mgrid[:obj.shape[0],:obj.shape[1]]
-    obj_phase_ramp = np.exp(-1j*obj_Ys * phase_ramp_dir[1]+
-                        -1j*obj_Xs * phase_ramp_dir[0])
+    obj_Xs, obj_Ys = np.mgrid[: obj.shape[0], : obj.shape[1]]
+    obj_phase_ramp = np.exp(
+        -1j * obj_Ys * phase_ramp_dir[1] + -1j * obj_Xs * phase_ramp_dir[0]
+    )
     test_obj = test_obj * obj_phase_ramp
 
     s_probe, s_obj = analysis.standardize(test_probe, test_obj, correct_ramp=True)
@@ -205,12 +217,12 @@ def test_standardize():
     assert np.max(s_obj - obj) / np.max(np.abs(obj)) < 1e-4
 
     # Finally a test with the phase ramp and multiple probes
-    subdominant_probe = 0.1*np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
+    subdominant_probe = (0.1 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5)))
     subdominant_probe = subdominant_probe * np.exp(-1j * np.angle(np.sum(subdominant_probe)))
     test_subdominant_probe = subdominant_probe * 37.6
     test_subdominant_probe = test_subdominant_probe * phase_ramp
 
-    incoh_probe = np.array([test_probe,test_subdominant_probe])
+    incoh_probe = np.array([test_probe, test_subdominant_probe])
 
     s_probe, s_obj = analysis.standardize(incoh_probe, test_obj, correct_ramp=True)
 
@@ -219,7 +231,6 @@ def test_standardize():
     assert np.max(s_probe[1] - subdominant_probe) / np.max(np.abs(subdominant_probe)) < 1e-4
 
 
-
 def test_synthesize_reconstructions():
     # I can only really test for a lack of failures, so I think my plan
     # will be to create a dataset that just needs to be added and see that
@@ -227,140 +238,153 @@ def test_synthesize_reconstructions():
 
     # Start by making a probe and object that should meet the standardization
     # conditions
-    probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
-    probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
+    probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
+    probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
     probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
 
-    assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
+    assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
     assert np.abs(np.angle(np.sum(probe))) < 2e-7
 
-    obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
-    obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
-                      (obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
+    obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
+    obj_slice = np.s_[
+        (obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
+        (obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
+    ]
 
     obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
-    assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
+    assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
 
     # Now I make stacks of identical probes and objects
-    probes = [probe,probe,probe,probe]
+    probes = [probe, probe, probe, probe]
     probe = np.copy(probe)
-    objects = [obj,obj,obj,obj]
+    objects = [obj, obj, obj, obj]
     obj = np.copy(obj)
 
-    s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes,objects)
+    s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes, objects)
     assert np.max(s_probe - probe) < 2e-5
     assert np.max(s_obj - obj) < 2e-5
     for t_obj in obj_stack:
         assert np.max(t_obj - obj) < 5e-5
-        
-    
+
 
 def test_calc_consistency_prtf():
 
     # Create an object with a specific structure
-    obj = 30 * np.random.rand(1030,1040) * np.exp(1j * (np.random.rand(1030,1040) - 0.5))
+    obj = (30 * np.random.rand(1030, 1040) * np.exp(1j * (np.random.rand(1030, 1040) - 0.5)))
 
     #
     synth_obj = np.sqrt(0.7) * obj
 
-    obj_stack = [obj,obj,obj,obj]
+    obj_stack = [obj, obj, obj, obj]
+
+    basis = np.array([[0, 2, 0], [3, 0, 0]])
 
-    basis = np.array([[0,2,0],
-                      [3,0,0]])
-    
     freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis)
     assert np.allclose(prtf, 0.7)
-    
+
     freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis, nbins=30)
     assert np.allclose(prtf, 0.7)
 
     # Check that it also works with torch input
     t_synth_obj = t.as_tensor(synth_obj)
     t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
-    freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
+    freqs, prtf = analysis.calc_consistency_prtf(
+        t_synth_obj, t_obj_stack, basis, nbins=30
+    )
     assert np.allclose(prtf.numpy(), 0.7)
 
     # And also when the basis is in torch
     t_synth_obj = t.as_tensor(synth_obj)
     t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
-    freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
+    freqs, prtf = analysis.calc_consistency_prtf(
+        t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30
+    )
     assert np.allclose(prtf.numpy(), 0.7)
 
-    
     # Check that is uses the right number of bins
     assert len(prtf) == 30
     assert len(freqs) == 30
-    
+
     # Check that the maximum frequency is correct for the basis
-    assert np.isclose(freqs[-1]-freqs[-2] + freqs[-1], np.sqrt(1/4**2 + 1/6**2))
-    
+    assert np.isclose(freqs[-1] - freqs[-2] + freqs[-1], np.sqrt(1 / 4**2 + 1 / 6**2))
+
 
 def test_calc_deconvolved_cross_correlation():
 
-    obj1 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
-    obj2 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
-    
+    obj1 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
+    obj2 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
+
     cor_fft = np.fft.fft2(obj1) * np.conj(np.fft.fft2(obj2))
-    
+
     # Not sure if this is more or less stable than just the correlation
     # maximum - requires some testing
     np_cor = np.fft.ifft2(cor_fft / np.abs(cor_fft))
     # test with numpy inputs
-    test_cor = analysis.calc_deconvolved_cross_correlation(obj1,obj2, im_slice=np.s_[:,:])
+    test_cor = analysis.calc_deconvolved_cross_correlation(
+        obj1, obj2, im_slice=np.s_[:, :]
+    )
     assert np.allclose(test_cor, np_cor)
-    
+
     # test with pytorch inputs
     obj1_t = t.as_tensor(obj1)
     obj2_t = t.as_tensor(obj2)
-    test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:])
-    
+    test_cor_t = analysis.calc_deconvolved_cross_correlation(
+        obj1_t, obj2_t, im_slice=np.s_[:, :]
+    )
+
     assert np.allclose(test_cor_t.numpy(), np_cor)
 
-    
+
 def test_calc_frc():
 
-    obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
-    obj2 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
+    obj1 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
+    obj2 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
 
-    basis = np.array([[0,2,0],
-                      [3,0,0]])
+    basis = np.array([[0, 2, 0], [3, 0, 0]])
 
     nbins = 100
     snr = 2
-    
-    cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])) * \
-        np.fft.fftshift(np.conj(np.fft.fft2(obj2[10:-10,20:-20])))
-    
-    F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])))**2
-    F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10,20:-20])))**2
-    
 
-    di = np.linalg.norm(basis[:,0]) 
-    dj = np.linalg.norm(basis[:,1])
-    
-    i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0],d=di))
-    j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1],d=dj))
-    
-    Js,Is = np.meshgrid(j_freqs,i_freqs)
-    Rs = np.sqrt(Is**2+Js**2)
+    cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20])) * np.fft.fftshift(
+        np.conj(np.fft.fft2(obj2[10:-10, 20:-20]))
+    )
 
-    numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft)
-    denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1)
-    denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2)
-    n_pix, bins = np.histogram(Rs,bins=nbins)
+    F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20]))) ** 2
+    F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10, 20:-20]))) ** 2
+
+    di = np.linalg.norm(basis[:, 0])
+    dj = np.linalg.norm(basis[:, 1])
+
+    i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0], d=di))
+    j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1], d=dj))
+
+    Js, Is = np.meshgrid(j_freqs, i_freqs)
+    Rs = np.sqrt(Is**2 + Js**2)
+
+    numerator, bins = np.histogram(Rs, bins=nbins, weights=cor_fft)
+    denominator_F1, bins = np.histogram(Rs, bins=nbins, weights=F1)
+    denominator_F2, bins = np.histogram(Rs, bins=nbins, weights=F2)
+    n_pix, bins = np.histogram(Rs, bins=nbins)
     bins = bins[:-1]
-    
-    frc = numerator / np.sqrt(denominator_F1*denominator_F2)
-        # This moves from combined-image SNR to single-image SNR
+
+    frc = numerator / np.sqrt(denominator_F1 * denominator_F2)
+    # This moves from combined-image SNR to single-image SNR
     snr /= 2
-    
-    threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
-        (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
-    
+
+    threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / (
+        1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)
+    )
+
     test_bins, test_frc, test_threshold = analysis.calc_frc(
-        obj1, obj2, basis, im_slice=np.s_[10:-10,20:-20],
-        nbins=100, snr=2, limit='corner')
-    
+        obj1,
+        obj2,
+        basis,
+        im_slice=np.s_[10:-10, 20:-20],
+        nbins=100,
+        snr=2,
+        limit="corner",
+    )
+
     assert np.allclose(bins, test_bins)
     assert np.allclose(frc, test_frc)
     assert np.allclose(threshold, test_threshold)
@@ -374,33 +398,40 @@ def test_calc_frc():
         obj1_torch,
         obj2_torch,
         basis_torch,
-        im_slice=np.s_[10:-10,20:-20], nbins=100, snr=2, limit='corner')
+        im_slice=np.s_[10:-10, 20:-20],
+        nbins=100,
+        snr=2,
+        limit="corner",
+    )
 
     assert np.allclose(bins, test_bins_t.numpy())
     assert np.allclose(frc, test_frc_t.numpy())
     assert np.allclose(threshold, test_threshold_t.numpy())
 
-    
+
 def test_calc_rms_error():
-    field_1 = t.rand(14,19, dtype=t.complex64)
-    field_2 = t.rand(14,19, dtype=t.complex64)
+    field_1 = t.rand(14, 19, dtype=t.complex64)
+    field_2 = t.rand(14, 19, dtype=t.complex64)
 
     # Check that the calculation is insensitive to phase
-    assert t.allclose(analysis.calc_rms_error(field_1, field_2),
-                      analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
+    assert t.allclose(
+        analysis.calc_rms_error(field_1, field_2),
+        analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2),
+    )
 
-    # And that it is sensitive to phase if we turn off the 
+    # And that it is sensitive to phase if we turn off the
     assert not t.allclose(
         analysis.calc_rms_error(field_1, field_2, align_phases=False),
-        analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
-                                align_phases=False))
+        analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2, align_phases=False),
+    )
 
     # Check that the result is positive
     assert analysis.calc_rms_error(field_1, field_2) > 0
 
     # And that it is a smaller number with align_phases on
-    assert (analysis.calc_rms_error(field_1, field_2) <=
-            analysis.calc_rms_error(field_1, field_2, align_phases=False))
+    assert analysis.calc_rms_error(field_1, field_2) <= analysis.calc_rms_error(
+        field_1, field_2, align_phases=False
+    )
 
     # Now we check against an explicit implementation:
     gamma = field_1 * t.conj(field_2)
@@ -408,114 +439,126 @@ def test_calc_rms_error():
 
     # This is an alternate way of doing the calculation. Actually, would this
     # be a better implementation anyway? Probably no difference tbh.
-    rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
-                                t.mean(t.abs(field_2)**2) -
-                                2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
-    assert t.allclose(rms_error_nophase,
-                      analysis.calc_rms_error(field_1, field_2))
+    rms_error_nophase = t.sqrt(
+        (
+            t.mean(t.abs(field_1) ** 2)
+            + t.mean(t.abs(field_2) ** 2)
+            - 2 * t.abs(t.mean(field_1 * t.conj(field_2)))
+        )
+    )
+    assert t.allclose(rms_error_nophase, analysis.calc_rms_error(field_1, field_2))
 
-    rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
-                              t.mean(t.abs(field_2)**2) -
-                              2 * t.real(t.mean(field_1 * t.conj(field_2)))))
+    rms_error_phase = t.sqrt(
+        (
+            t.mean(t.abs(field_1) ** 2)
+            + t.mean(t.abs(field_2) ** 2)
+            - 2 * t.real(t.mean(field_1 * t.conj(field_2)))
+        )
+    )
 
-    assert t.allclose(rms_error_phase,
-                      analysis.calc_rms_error(field_1, field_2,
-                                              align_phases=False))
+    assert t.allclose(
+        rms_error_phase, analysis.calc_rms_error(field_1, field_2, align_phases=False)
+    )
 
     # Now let's test that it works along a dimension:
-    
-    field_1 = t.rand(3,14,19, dtype=t.complex64)
-    field_2 = t.rand(3,14,19, dtype=t.complex64)
+
+    field_1 = t.rand(3, 14, 19, dtype=t.complex64)
+    field_2 = t.rand(3, 14, 19, dtype=t.complex64)
     result = analysis.calc_rms_error(field_1, field_2, normalize=True)
-    assert (result.shape == t.Size([3]))
-    
+    assert result.shape == t.Size([3])
+
     for i in range(3):
-        assert t.allclose(analysis.calc_rms_error(field_1[i],
-                                                  field_2[i],
-                                                  normalize=True),
-                          result[i])
-    
+        assert t.allclose(
+            analysis.calc_rms_error(field_1[i], field_2[i], normalize=True), result[i]
+        )
+
 
 def test_calc_fidelity():
 
-    fields_1 = t.rand(2,30,17, dtype=t.complex128)
-    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+    fields_1 = t.rand(2, 30, 17, dtype=t.complex128)
+    fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
 
-    dm_1 = t.reshape(fields_1, (2,-1))
-    dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
-    dm_2 = t.reshape(fields_2, (3,-1))
-    dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
+    dm_1 = t.reshape(fields_1, (2, -1))
+    dm_1 = t.tensordot(dm_1.transpose(0, 1), dm_1.conj(), dims=1).numpy()
+    dm_2 = t.reshape(fields_2, (3, -1))
+    dm_2 = t.tensordot(dm_2.transpose(0, 1), dm_2.conj(), dims=1).numpy()
 
     sqrt_dm_1 = la.sqrtm(dm_1).astype(dm_1.dtype)
-    inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1,dm_2), sqrt_dm_1))
-    inner_mat = inner_mat.astype(dm_1.dtype) #la.sqrtm doubles the precision
-    fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
+    inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1, dm_2), sqrt_dm_1))
+    inner_mat = inner_mat.astype(dm_1.dtype)  # la.sqrtm doubles the precision
+    fidelity = t.as_tensor(np.abs(np.trace(inner_mat)) ** 2)
 
     assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
 
     # Check that it reduces to the overlap for coherent fields
-    fields_1 = t.rand(1,30,17, dtype=t.complex128)
-    fields_2 = t.rand(1,30,17, dtype=t.complex128)
+    fields_1 = t.rand(1, 30, 17, dtype=t.complex128)
+    fields_2 = t.rand(1, 30, 17, dtype=t.complex128)
 
-    assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
-                     analysis.calc_fidelity(fields_1, fields_2))
+    assert t.isclose(
+        t.abs(t.sum(fields_1 * fields_2.conj())) ** 2,
+        analysis.calc_fidelity(fields_1, fields_2),
+    )
     # Checking that it works with extra dimensions
-    fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
-    fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
-    field_3 = t.rand(1,30,17, dtype=t.complex128)
+    fields_1 = t.rand(3, 3, 30, 17, dtype=t.complex128)
+    fields_2 = t.rand(3, 1, 30, 17, dtype=t.complex128)
+    field_3 = t.rand(1, 30, 17, dtype=t.complex128)
 
     fidelities = analysis.calc_fidelity(fields_1, fields_2)
     fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
     for i in range(3):
-        assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
-                         fidelities[i])
-        assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
-                         fidelities_2[i])
+        assert t.isclose(
+            analysis.calc_fidelity(fields_1[i], fields_2[i]), fidelities[i]
+        )
+        assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3), fidelities_2[i])
 
     # Check that the diensionality argument works
-    fields_1 = t.rand(3,2,12, dtype=t.complex128)
-    fields_2 = t.rand(3,2,12, dtype=t.complex128)
+    fields_1 = t.rand(3, 2, 12, dtype=t.complex128)
+    fields_2 = t.rand(3, 2, 12, dtype=t.complex128)
 
-    assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
-            == t.Size([3]))
-    
-    fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
-    fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
+    assert analysis.calc_fidelity(fields_1, fields_2, dims=1).shape == t.Size([3])
+
+    fields_1 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
+    fields_2 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
+
+    assert analysis.calc_fidelity(fields_1, fields_2, dims=3).shape == t.Size([3])
 
-    assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
-            == t.Size([3]))
 
 def test_calc_generalized_rms_error():
 
     # Test that it matches the rms error for coherent fields
-    
-    fields_1 = t.rand(1,30,17, dtype=t.complex128)
-    fields_2 = t.rand(1,30,17, dtype=t.complex128)
-    
-    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
-                     analysis.calc_rms_error(fields_1[0], fields_2[0],
-                                             align_phases=True))
+
+    fields_1 = t.rand(1, 30, 17, dtype=t.complex128)
+    fields_2 = t.rand(1, 30, 17, dtype=t.complex128)
+
+    assert t.isclose(
+        analysis.calc_generalized_rms_error(fields_1, fields_2),
+        analysis.calc_rms_error(fields_1[0], fields_2[0], align_phases=True),
+    )
 
     # Test that it is independent of field order
-    fields_1 = t.rand(5,30,17, dtype=t.complex128)
-    fields_2 = t.rand(3,30,17, dtype=t.complex128)
+    fields_1 = t.rand(5, 30, 17, dtype=t.complex128)
+    fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
     fields_3 = fields_2.flip(0)
-    
-    assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
-                     analysis.calc_generalized_rms_error(fields_1, fields_3))
+
+    assert t.isclose(
+        analysis.calc_generalized_rms_error(fields_1, fields_2),
+        analysis.calc_generalized_rms_error(fields_1, fields_3),
+    )
 
     # Test with leading dimensions
-    fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
-    fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
-    
-    assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
-            == t.Size([3,4]))
+    fields_1 = t.rand(3, 4, 2, 10, 17, dtype=t.complex128)
+    fields_2 = t.rand(3, 4, 3, 10, 17, dtype=t.complex128)
+
+    assert analysis.calc_generalized_rms_error(fields_1, fields_2).shape == t.Size(
+        [3, 4]
+    )
 
     # And test with different number of dimensions dims
     # Test that it is independent of field order
-    fields_1 = t.rand(3,6,17, dtype=t.complex128)
-    fields_2 = t.rand(3,1,17, dtype=t.complex128)
+    fields_1 = t.rand(3, 6, 17, dtype=t.complex128)
+    fields_2 = t.rand(3, 1, 17, dtype=t.complex128)
     fields_3 = fields_2.flip(0)
 
-    assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
-
+    assert analysis.calc_generalized_rms_error(
+        fields_1, fields_2, dims=1
+    ).shape == t.Size([3])
diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py
index 39b64dd..7f2571c 100644
--- a/tests/tools/test_data.py
+++ b/tests/tools/test_data.py
@@ -1,60 +1,55 @@
-from cdtools.tools import data
-import numpy as np
-import torch as t
-import h5py
-import pytest
-import os
 import datetime
 import numbers
-import pathlib
 
+import h5py
+import numpy as np
+import torch as t
 
+from cdtools.tools import data
 
 #
 # We start with a bunch of tests of the data loading capabilities
 #
 
 
-
 def test_get_entry_info(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         entry_info = data.get_entry_info(cxi)
         for key in expected['entry metadata']:
             assert entry_info[key] == expected['entry metadata'][key]
 
-        
+
 def test_get_sample_info(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         sample_info = data.get_sample_info(cxi)
         if sample_info is None and  \
-           ('sample info' not in expected or
-            expected['sample info'] is None):
+           ('sample info' not in expected or expected['sample info'] is None):
             # Valid if no sample info is defined at all
             continue
         for key in expected['sample info']:
-            if isinstance(expected['sample info'][key],np.ndarray):
+            if isinstance(expected['sample info'][key], np.ndarray):
                 assert np.allclose(sample_info[key],
                                    expected['sample info'][key])
             else:
                 assert sample_info[key] == expected['sample info'][key]
-                
-    
+
+
 def test_get_wavelength(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
-        assert np.isclose(expected['wavelength'],data.get_wavelength(cxi))
+        assert np.isclose(expected['wavelength'], data.get_wavelength(cxi))
 
 
 def test_get_detector_geometry(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         distance, basis, corner = data.get_detector_geometry(cxi)
-        assert np.isclose(distance,expected['detector']['distance'])
-        assert np.allclose(basis,expected['detector']['basis'])
+        assert np.isclose(distance, expected['detector']['distance'])
+        assert np.allclose(basis, expected['detector']['basis'])
         if isinstance(expected['detector']['corner'], np.ndarray):
             assert np.allclose(corner, expected['detector']['corner'])
         else:
             assert corner == expected['detector']['corner']
 
-            
+
 def test_get_mask(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         mask = data.get_mask(cxi)
@@ -62,7 +57,7 @@ def test_get_mask(test_ptycho_cxis):
             continue
         assert np.all(mask == expected['mask'])
 
-        
+
 def test_get_qe_mask(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         qe_mask = data.get_qe_mask(cxi)
@@ -78,8 +73,8 @@ def test_get_dark(test_ptycho_cxis):
             assert expected['dark'] is None
         else:
             assert np.allclose(dark, expected['dark'])
-            
-        
+
+
 def test_get_data(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         patterns, axes = data.get_data(cxi)
@@ -99,8 +94,6 @@ def test_get_ptycho_translations(test_ptycho_cxis):
         assert np.allclose(data.get_ptycho_translations(cxi),
                            expected['translations'])
 
-
-
 #
 # Then, write a test for the data saving. It should create a .cxi file
 # using the data seving tools, and then check that when read with the
@@ -110,13 +103,13 @@ def test_get_ptycho_translations(test_ptycho_cxis):
 
 def test_create_cxi(tmp_path):
     data.create_cxi(tmp_path / 'test_create.cxi')
-    with h5py.File(tmp_path / 'test_create.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_create.cxi', 'r') as f:
         assert f['cxi_version'][()] == 160
         assert 'entry_1' in f
 
-        
+
 def test_add_entry_info(tmp_path):
-    entry_info = {'experiment_identifier':'test of cxi file writing tools',
+    entry_info = {'experiment_identifier': 'test of cxi file writing tools',
                   'title': 'my cool experiment',
                   'start_time': datetime.datetime.now(),
                   'end_time': datetime.datetime.now()}
@@ -124,12 +117,11 @@ def test_add_entry_info(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_entry_info.cxi') as f:
         data.add_entry_info(f, entry_info)
 
-
-    with h5py.File(tmp_path / 'test_add_entry_info.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_entry_info.cxi', 'r') as f:
         read_entry_info = data.get_entry_info(f)
 
     print(read_entry_info)
-        
+
     for key in entry_info:
         if isinstance(entry_info[key], np.ndarray):
             assert np.allclose(entry_info[key], read_entry_info[key])
@@ -138,18 +130,18 @@ def test_add_entry_info(tmp_path):
 
 
 def test_add_sample_info(tmp_path):
-    sample_info = {'name':'A nice fake sample',
+    sample_info = {'name': 'A nice fake sample',
                    'concentration': 10,
                    'mass': 5.3,
                    'temperature': 76,
                    'description': 'A very nice sample',
-                   'unit_cell': np.array([1,1,1,90.,90.,90.])}
+                   'unit_cell': np.array([1, 1, 1, 90., 90., 90.])}
 
     with data.create_cxi(tmp_path / 'test_add_sample_info.cxi') as f:
         data.add_sample_info(f, sample_info)
 
-    with h5py.File(tmp_path / 'test_add_sample_info.cxi','r') as f:
-        read_sample_info = data.get_sample_info(f)    
+    with h5py.File(tmp_path / 'test_add_sample_info.cxi', 'r') as f:
+        read_sample_info = data.get_sample_info(f)
 
     for key in sample_info:
         if isinstance(sample_info[key], np.ndarray):
@@ -158,7 +150,7 @@ def test_add_sample_info(tmp_path):
             assert np.isclose(sample_info[key], read_sample_info[key])
         else:
             assert sample_info[key] == read_sample_info[key]
-                      
+
 
 def test_add_source(tmp_path):
     wavelength = 1e-9
@@ -167,26 +159,26 @@ def test_add_source(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_source.cxi') as f:
         data.add_source(f, wavelength)
 
-    with h5py.File(tmp_path / 'test_add_source.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_source.cxi', 'r') as f:
         # Check this directly since we want to make sure it saved
         # the wavelength and energy
         read_wavelength = f['entry_1/instrument_1/source_1/wavelength'][()]
         read_energy = f['entry_1/instrument_1/source_1/energy'][()]
 
-    assert np.isclose( wavelength, read_wavelength)
-    assert np.isclose( energy, read_energy)
+    assert np.isclose(wavelength, read_wavelength)
+    assert np.isclose(energy, read_energy)
+
 
-    
 def test_add_detector(tmp_path):
     distance = 0.34
-    basis = np.array([[0,-30e-6,0],
-                      [-20e-6,0,0]]).astype(np.float32).transpose()
-    corner = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
-    
+    basis = np.array([[0, -30e-6, 0],
+                      [-20e-6, 0, 0]]).astype(np.float32).transpose()
+    corner = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
+
     with data.create_cxi(tmp_path / 'test_add_detector.cxi') as f:
         data.add_detector(f, distance, basis, corner=corner)
 
-    with h5py.File(tmp_path / 'test_add_detector.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_detector.cxi', 'r') as f:
         # Check this directly since we want to make sure it saved
         # the pixel sizes
         d1 = f['entry_1/instrument_1/detector_1']
@@ -198,112 +190,110 @@ def test_add_detector(tmp_path):
 
     assert np.isclose(distance, read_distance)
     assert np.allclose(basis, read_basis)
-    assert np.isclose(np.linalg.norm(basis[:,1]), read_x_pix)
-    assert np.isclose(np.linalg.norm(basis[:,0]), read_y_pix)
-    assert np.allclose(corner,read_corner)
-   
+    assert np.isclose(np.linalg.norm(basis[:, 1]), read_x_pix)
+    assert np.isclose(np.linalg.norm(basis[:, 0]), read_y_pix)
+    assert np.allclose(corner, read_corner)
+
 
 def test_add_mask(tmp_path):
-    mask = (np.random.rand(350,600) > 0.1).astype(np.uint8)
+    mask = (np.random.rand(350, 600) > 0.1).astype(np.uint8)
 
     with data.create_cxi(tmp_path / 'test_add_mask.cxi') as f:
         data.add_mask(f, mask)
 
-    with h5py.File(tmp_path / 'test_add_mask.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_mask.cxi', 'r') as f:
         read_mask = data.get_mask(f)
 
     assert np.all(mask == read_mask)
 
-    
+
 def test_add_qe_mask(tmp_path):
-    qe_mask = np.random.rand(350,199).astype(np.float32)
+    qe_mask = np.random.rand(350, 199).astype(np.float32)
 
     with data.create_cxi(tmp_path / 'test_add_qe_mask.cxi') as f:
         data.add_qe_mask(f, qe_mask)
 
-    with h5py.File(tmp_path / 'test_add_qe_mask.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_qe_mask.cxi', 'r') as f:
         read_qe_mask = data.get_qe_mask(f)
 
     assert np.allclose(qe_mask, read_qe_mask)
 
-    
+
 def test_add_dark(tmp_path):
-    dark = np.random.rand(350,620)
+    dark = np.random.rand(350, 620)
 
     with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f:
         data.add_dark(f, dark)
 
-    with h5py.File(tmp_path / 'test_add_dark.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_dark.cxi', 'r') as f:
         read_dark = data.get_dark(f)
 
     print(dark.shape)
     assert np.allclose(dark, read_dark)
 
 
-
 def test_add_data(tmp_path):
     # First test from numpy, with axes
-    fake_data = np.random.rand(100,256,256)
-    axes = ['translation','y','x']
+    fake_data = np.random.rand(100, 256, 256)
+    axes = ['translation', 'y', 'x']
 
     with data.create_cxi(tmp_path / 'test_add_data.cxi') as f:
         data.add_data(f, fake_data, axes)
 
-    with h5py.File(tmp_path / 'test_add_data.cxi','r') as f:
+    with h5py.File(tmp_path / 'test_add_data.cxi', 'r') as f:
         # Check this directly since we want to make sure it saved
         # it in all the places it should have
         read_data_1 = f['entry_1/data_1/data'][()]
         read_data_2 = f['entry_1/instrument_1/detector_1/data'][()]
         read_axes = str(f['entry_1/instrument_1/detector_1/data'].attrs['axes'].decode())
-    
+
     assert np.allclose(fake_data, read_data_1)
     assert np.allclose(fake_data, read_data_2)
     assert 'translation:y:x' == read_axes
 
     # Then test from torch, without axes
     fake_data = t.from_numpy(fake_data)
-    
+
     with data.create_cxi(tmp_path / 'test_add_data_torch.cxi') as f:
         data.add_data(f, fake_data)
 
-    with h5py.File(tmp_path / 'test_add_data_torch.cxi','r') as f:    
+    with h5py.File(tmp_path / 'test_add_data_torch.cxi', 'r') as f:
         read_data, axes = data.get_data(f)
 
-    assert np.allclose(fake_data.numpy(),read_data)
+    assert np.allclose(fake_data.numpy(), read_data)
+
 
 def test_add_shot_to_shot_info(tmp_path):
-    
+
     analyzer = np.random.rand(100)
 
     with data.create_cxi(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
         data.add_shot_to_shot_info(f, analyzer, 'analyzer_angle')
-    
+
     with h5py.File(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
         # Check this directly since we want to make sure it saved
         # it in all the places it should have
         read_analyzer_1 = f['entry_1/data_1/analyzer_angle'][()]
-        read_analyzer_2 = \
-                       f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
+        read_analyzer_2 = f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
         read_analyzer_3 = f['entry_1/sample_1/geometry_1/analyzer_angle'][()]
 
     assert np.allclose(analyzer, read_analyzer_1)
     assert np.allclose(analyzer, read_analyzer_2)
     assert np.allclose(analyzer, read_analyzer_3)
-    
-        
+
+
 def test_add_ptycho_translations(tmp_path):
-    
-    translations = np.random.rand(3,100)
+
+    translations = np.random.rand(3, 100)
 
     with data.create_cxi(tmp_path / 'test_add_ptycho_translations.cxi') as f:
         data.add_ptycho_translations(f, translations)
-    
-    with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi','r') as f:
+
+    with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi', 'r') as f:
         # Check this directly since we want to make sure it saved
         # it in all the places it should have
         read_translations_1 = f['entry_1/data_1/translation'][()]
-        read_translations_2 = \
-                       f['entry_1/instrument_1/detector_1/translation'][()]
+        read_translations_2 = f['entry_1/instrument_1/detector_1/translation'][()]
         read_translations_3 = f['entry_1/sample_1/geometry_1/translation'][()]
 
     assert np.allclose(-translations, read_translations_1)
@@ -312,8 +302,7 @@ def test_add_ptycho_translations(tmp_path):
 
 
 def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
-    ### Tests both nested_dict_to_h5 and h5_to_nested_dict
-
+    # Tests both nested_dict_to_h5 and h5_to_nested_dict
     def check_dict_equality(truth, to_test):
         for key in truth.keys():
             if isinstance(truth[key], dict):
@@ -328,23 +317,24 @@ def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
                 assert truth[key] == to_test[key]
             else:
                 assert 0
-            
+
     for test_dict in example_nested_dicts:
         filename = tmp_path / 'example_dataset.h5'
         data.nested_dict_to_h5(filename, test_dict)
         roundtrip = data.h5_to_nested_dict(filename)
         check_dict_equality(test_dict, roundtrip)
-                
-    
+
+
 def test_h5_to_nested_dict(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         # Just test that it runs without errors for these ones.
         # A round-trip test is in test_nested_dict_to_h5
-        d = data.h5_to_nested_dict(cxi)
+        data.h5_to_nested_dict(cxi)
+
 
 def test_nested_dict_to_numpy(example_nested_dicts):
 
-    def check_dict_numpyness(truth, to_test): 
+    def check_dict_numpyness(truth, to_test):
         for key in truth.keys():
             if isinstance(truth[key], dict):
                 check_dict_numpyness(truth[key], to_test[key])
@@ -361,12 +351,11 @@ def test_nested_dict_to_numpy(example_nested_dicts):
 
     for test_dict in example_nested_dicts:
         numpy_dict = data.nested_dict_to_numpy(test_dict)
-        check_dict_numpyness(test_dict, numpy_dict)            
-        
-    
-def test_nested_dict_to_torch(example_nested_dicts):
+        check_dict_numpyness(test_dict, numpy_dict)
 
-    def check_dict_torchiness(truth, to_test): 
+
+def test_nested_dict_to_torch(example_nested_dicts):
+    def check_dict_torchiness(truth, to_test):
         for key in truth.keys():
             if isinstance(truth[key], dict):
                 check_dict_torchiness(truth[key], to_test[key])

From 88df4c079484f401bcb5812d0eabef2b4d952f8a Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 14:40:17 -0700
Subject: [PATCH 27/55] linting test_image_processing.py and
 test_initializers.py

---
 tests/tools/test_image_processing.py | 115 ++++++++++++-------------
 tests/tools/test_initializers.py     | 124 +++++++++++++--------------
 2 files changed, 115 insertions(+), 124 deletions(-)

diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py
index 83c34d8..a6385cd 100644
--- a/tests/tools/test_image_processing.py
+++ b/tests/tools/test_image_processing.py
@@ -1,18 +1,19 @@
 import numpy as np
 import torch as t
+from scipy import ndimage
 
 from cdtools.tools import image_processing, interactions
-from scipy import ndimage
+
 
 def test_centroid():
     # Test single im
-    im = t.rand((30,40))
+    im = t.rand((30, 40))
     sp_centroid = ndimage.center_of_mass(im.numpy())
     centroid = image_processing.centroid(im)
     assert t.allclose(centroid, t.Tensor(sp_centroid))
-    
+
     # Test stack o' ims
-    ims = t.rand((5,30,40))
+    ims = t.rand((5, 30, 40))
     sp_centroids = [ndimage.center_of_mass(im.numpy())
                     for im in ims]
     centroids = image_processing.centroid(ims)
@@ -21,33 +22,33 @@ def test_centroid():
 
 def test_centroid_sq():
     # Test single im
-    im = t.rand((30,40))
+    im = t.rand((30, 40))
     sp_centroid = ndimage.center_of_mass(im.numpy()**2)
     centroid = image_processing.centroid_sq(im)
     assert t.allclose(centroid, t.Tensor(sp_centroid))
 
     # Test complex with multiple ims
-    ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
+    ims = t.rand((5, 30, 40)) + 1j * t.rand((5, 30, 40))
     np_ims = ims.numpy()
     sp_centroids = [ndimage.center_of_mass(np.abs(im)**2)
                     for im in np_ims]
     centroids = image_processing.centroid_sq(ims, comp=True)
     assert t.allclose(centroids, t.Tensor(np.array(sp_centroids)))
 
-    
+
 def test_sinc_subpixel_shift():
 
-    im = np.zeros((512,512), dtype=np.complex128)
-    im[256,256] = 1
+    im = np.zeros((512, 512), dtype=np.complex128)
+    im[256, 256] = 1
 
     # test it by creating a single pixel object and seeing that it is
     # shifted correctly
     xs = np.arange(512) - 256
-    Ys,Xs = np.meshgrid(xs,xs)
-    sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
+    Ys, Xs = np.meshgrid(xs, xs)
+    sinc_im = np.sinc(Xs - 0.3) * np.sinc(Ys - 0.6)
 
     torch_im = t.as_tensor(im)
-    test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
+    test_im = image_processing.sinc_subpixel_shift(torch_im, (0.3, 0.6))
 
     # The fidelity isn't great due to the FFT-based approach, so we need
     # a pretty relaxed condition
@@ -57,84 +58,82 @@ def test_sinc_subpixel_shift():
 def test_find_pixel_shift():
 
     # Test two real ims
-    big_im = t.rand((30,70))
-    im1 = big_im[3:,:-20]
-    im2 = big_im[:-3,20:]
-    assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([-3,20]))
+    big_im = t.rand((30, 70))
+    im1 = big_im[3:, :-20]
+    im2 = big_im[:-3, 20:]
+    assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([-3, 20]))
 
     # Test a real and complex im
-    big_im = t.rand((30,70))
-    im1 = big_im[:-5,10:].to(dtype=t.complex64)
-    im2 = big_im[5:,:-10]
-    assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
-    assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
-    
+    big_im = t.rand((30, 70))
+    im1 = big_im[:-5, 10:].to(dtype=t.complex64)
+    im2 = big_im[5:, :-10]
+    assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, -10]))
+    assert t.all(image_processing.find_pixel_shift(im2, im1) == t.LongTensor([-5, 10]))
+
     # Test two complex ims
-    big_im = t.rand((45,45)) + 1j * t.rand((45,45))
-    im1 = big_im[:-5,:-4]
-    im2 = big_im[5:,4:]
-    assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
+    big_im = t.rand((45, 45)) + 1j * t.rand((45, 45))
+    im1 = big_im[:-5, :-4]
+    im2 = big_im[5:, 4:]
+    assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, 4]))
 
 
 def test_find_subpixel_shift():
     # We can do this by creating a test probe and a test object
-    test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
-    test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
+    test_probe = t.rand((70, 70)) + 1j * t.rand((70, 70))
+    test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
+
+    shift = t.tensor((0.8, 0.75))
 
-    shift = t.tensor((0.8,0.75))
-    
     im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
-    
-    retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
+
+    retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0, 0), resolution=50)
     # tolerance of 0.03 on this measurement
     assert t.all(t.abs(shift - retrieved_shift) < 0.03)
 
-    
+
 def test_find_shift():
 
     # We can do this by creating a test probe and a test object
-    test_probe = t.rand((200,200)) + 1j * t.rand((200,200)) 
-    test_obj = t.ones((300,300)) + 1j *  t.rand((300,300)) 
+    test_probe = t.rand((200, 200)) + 1j * t.rand((200, 200))
+    test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
+
+    shift = t.tensor((0.8, 0.75))
 
-    shift = t.tensor((0.8,0.75))
-    
     im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
-                                     multiple_modes=False)[:-40,:-6]
+                                     multiple_modes=False)[:-40, :-6]
 
-    retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
+    retrieved_shift = image_processing.find_shift(im, test_probe[40:, 6:], resolution=50)
     # tolerance of 0.03 on this measurement
-    assert t.all(t.abs(shift + t.Tensor((40,6)) - retrieved_shift) < 0.03)
+    assert t.all(t.abs(shift + t.Tensor((40, 6)) - retrieved_shift) < 0.03)
+
 
-    
 def test_convolve_1d():
-    test_image = np.random.rand(400,300)
-    #test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
-    xs = np.linspace(-100,100,300)
-    kernel = 1/(1+xs**2)
+    test_image = np.random.rand(400, 300)
+    # test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
+    xs = np.linspace(-100, 100, 300)
+    kernel = 1 / (1 + xs**2)
 
     # First, we test with everything real, dim=1
     convolved = image_processing.convolve_1d(t.as_tensor(test_image),
-                                             t.as_tensor(kernel),dim=1)
+                                             t.as_tensor(kernel), dim=1)
 
-    np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
-    assert np.allclose(convolved.numpy(),np_result)
+    np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
+    assert np.allclose(convolved.numpy(), np_result)
 
-
-    xs = np.linspace(-100,100,400)
-    kernel = 1/(1+xs**2)
+    xs = np.linspace(-100, 100, 400)
+    kernel = 1 / (1 + xs**2)
 
     # Then with dim=0, and a non-fftshifted kernel
     convolved = image_processing.convolve_1d(t.as_tensor(test_image),
                                              t.as_tensor(np.fft.ifftshift(kernel)),
                                              fftshift_kernel=False)
 
-    np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
-    assert np.allclose(convolved.numpy(),np_result)
+    np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0))
+    assert np.allclose(convolved.numpy(), np_result)
 
     # And finally with complex input
-    convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
-                                             t.as_tensor(kernel,dtype=t.complex64)).numpy()
-
-    np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
-    assert np.allclose(convolved,np_result)
+    convolved = image_processing.convolve_1d(t.as_tensor(test_image, dtype=t.complex64),
+                                             t.as_tensor(kernel, dtype=t.complex64)).numpy()
 
+    np_result = np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0)
+    assert np.allclose(convolved, np_result)
diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py
index 53c284a..5430271 100644
--- a/tests/tools/test_initializers.py
+++ b/tests/tools/test_initializers.py
@@ -1,26 +1,27 @@
-from cdtools.tools import initializers
-from cdtools.datasets import Ptycho2DDataset
 import numpy as np
 import torch as t
 
+from cdtools.tools import initializers
+from cdtools.datasets import Ptycho2DDataset
+
+
 def test_exit_wave_geometry():
     # First test a simple case where nothing need change
-    basis = t.Tensor([[0,-30e-6,0],
-                      [-20e-6,0,0]]).transpose(0,1)
-    shape = t.Size([73,56])
+    basis = t.Tensor([[0, -30e-6, 0],
+                      [-20e-6, 0, 0]]).transpose(0, 1)
+    shape = t.Size([73, 56])
     wavelength = 1e-9
     distance = 1.
     rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
-    
-    assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
-    assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
-    
+
+    assert t.allclose(rs_basis[0, 1], t.Tensor([-8.928571428571428e-07]))
+    assert t.allclose(rs_basis[1, 0], t.Tensor([-4.5662100456621004e-07]))
 
 
 def test_calc_object_setup():
     # First just try a simple case
-    probe_shape = t.Size([120,57])
-    translations = t.rand((30,2)) * 300
+    probe_shape = t.Size([120, 57])
+    translations = t.rand((30, 2)) * 300
     t_max = t.max(translations, dim=0)[0]
     t_min = t.min(translations, dim=0)[0]
     obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations)
@@ -28,65 +29,58 @@ def test_calc_object_setup():
     assert t.allclose(min_translation, t_min)
 
     assert obj_shape == t.Size(exp_shape)
-    
+
     # Then add some padding
     padding = 5
     obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations, padding=padding)
     assert t.allclose(min_translation, t_min - padding)
     assert obj_shape == t.Size(exp_shape + 2 * padding)
-        
 
-    
+
 def test_gaussian():
     # Generate gaussian as a numpy array (square array)
     shape = [10, 10]
     sigma = [2.5, 2.5]
-    
-    center = ((shape[0]-1)/2, (shape[1]-1)/2)
+
+    center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
-                          -0.5*((y-center[0])/sigma[0])**2)
+    np_result = 10 * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
     init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
     assert np.allclose(init_result, np_result)
 
     # Generate gaussian as a numpy array (rectangular array)
     shape = [10, 5]
     sigma = [2.5, 3]
-    center = ((shape[0]-1)/2, (shape[1]-1)/2)
+    center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
-                          -0.5*((y-center[0])/sigma[0])**2)
+    np_result = np.exp(-0.5 * ((x - center[1]) / sigma[1])**2
+                       - 0.5 * ((y - center[0]) / sigma[0])**2)
     init_result = initializers.gaussian(shape, sigma).numpy()
     assert np.allclose(init_result, np_result)
-                                    
+
     # Generate gaussian with curvature
     shape = [20, 30]
     sigma = [2.5, 5]
-    curvature = [1,0.6]
-    center = ((shape[0]-1)/2 + 3, (shape[1]-1)/2 - 1.4)
+    curvature = [1, 0.6]
+    center = ((shape[0] - 1) / 2 + 3, (shape[1] - 1) / 2 - 1.4)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = (10+0j)*np.exp(-0.5*((x-center[1])/sigma[1])**2
-                               -0.5*((y-center[0])/sigma[0])**2)
-    np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
-                        +0.5j*curvature[0]*(y-center[0])**2)
-    init_result = initializers.gaussian(shape, sigma, center=center,
-                            curvature=curvature, amplitude=10).numpy()
+    np_result = (10 + 0j) * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
+    np_result *= np.exp(0.5j * curvature[1] * (x - center[1])**2 + 0.5j * curvature[0] * (y - center[0])**2)
+    init_result = initializers.gaussian(shape, sigma, center=center, curvature=curvature, amplitude=10).numpy()
     assert np.allclose(init_result, np_result)
-    
+
 
 def test_gaussian_probe(ptycho_cxi_1):
-    
     dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
 
     det_basis = t.Tensor(dataset.detector_geometry['basis'])
     det_shape = t.Size(dataset.patterns.shape[-2:])
     wavelength = dataset.wavelength
     distance = dataset.detector_geometry['distance']
-    
     basis = initializers.exit_wave_geometry(det_basis,
-                                             det_shape,
-                                             wavelength,
-                                             distance)
+                                            det_shape,
+                                            wavelength,
+                                            distance)
     # Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction
     # Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir
 
@@ -94,15 +88,13 @@ def test_gaussian_probe(ptycho_cxi_1):
     sigma = 5e-7
 
     # Build a stage explicitly with numpy to compare against
-    x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy()
-    y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy()
-    Xs,Ys = np.meshgrid(x,y)
-    Rs = np.sqrt(Xs**2+Ys**2)
-
-    
+    x = (np.arange(256) - 127.5) * (-basis[0, 1]).numpy()
+    y = (np.arange(256) - 127.5) * (-basis[1, 0]).numpy()
+    Xs, Ys = np.meshgrid(x, y)
+    Rs = np.sqrt(Xs**2 + Ys**2)
 
     # Now we first test the non-propagated probe
-    np_probe = np.exp(-1/(2*sigma**2) * Rs**2)
+    np_probe = np.exp(- 1 / (2 * sigma**2) * Rs**2)
 
     normalization = 0
     for params, im in dataset:
@@ -110,27 +102,26 @@ def test_gaussian_probe(ptycho_cxi_1):
     normalization /= len(dataset)
 
     normalization_1 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
-    
+
     probe = initializers.gaussian_probe(
         dataset, basis, det_shape, sigma).numpy()
- 
-    assert np.allclose(probe, normalization_1*np_probe)
+
+    assert np.allclose(probe, normalization_1 * np_probe)
 
     # And then a propagated probe
-    z = 1e-4 #nm
+    z = 1e-4  # nm
     k = 2 * np.pi / wavelength
-    w0 = np.sqrt(2)*sigma
+    w0 = np.sqrt(2) * sigma
     zr = np.pi * w0**2 / wavelength
     wz = w0 * np.sqrt(1 + (z / zr)**2)
-    Rz =  z * (1 + (zr / z)**2)
-    np_probe  =  np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
+    Rz = z * (1 + (zr / z)**2)
+    np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
+
+    normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
 
-    normalization_2 = np.sqrt(normalization /  np.sum(np.abs(np_probe)**2))
-    
     probe = initializers.gaussian_probe(dataset, basis, det_shape, sigma,
                                         propagation_distance=z).numpy()
-    
-    assert np.allclose(probe, normalization_2*np_probe)
+    assert np.allclose(probe, normalization_2 * np_probe)
 
 
 def test_SHARP_style_probe(ptycho_cxi_1):
@@ -148,11 +139,13 @@ def test_SHARP_style_probe(ptycho_cxi_1):
                                             wavelength,
                                             distance)
 
+    assert basis.shape == t.Size([3, 2])
+
     probe = initializers.SHARP_style_probe(dataset)
-    assert probe.shape == t.Size([256,256])
+    assert probe.shape == t.Size([256, 256])
 
     probe = initializers.SHARP_style_probe(dataset, propagation_distance=20e-6)
-    assert probe.shape == t.Size([256,256])
+    assert probe.shape == t.Size([256, 256])
 
 
 def test_RPI_spectral_init():
@@ -160,28 +153,27 @@ def test_RPI_spectral_init():
     # 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]
+    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.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
+    background = t.as_tensor(np.random.rand(*pattern.shape), dtype=t.float32) * 0.05
 
     probe = t.as_tensor(probe)
     pattern = t.as_tensor(pattern)
-    
+
     obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
-    assert list(obj.shape) == [1]+obj_shape
+    assert list(obj.shape) == [1] + obj_shape
 
     obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
                                          n_modes=2, mask=mask)
-    assert list(obj.shape) == [2]+obj_shape
+    assert list(obj.shape) == [2] + obj_shape
 
     obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
                                          n_modes=2, background=background)
-    assert list(obj.shape) == [2]+obj_shape
+    assert list(obj.shape) == [2] + obj_shape
 
     obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
                                          n_modes=2, mask=mask,
                                          background=background)
-    assert list(obj.shape) == [2]+obj_shape
-
+    assert list(obj.shape) == [2] + obj_shape

From 65e83753a8aef67a66fc5fa97ed48f4d311f5cf3 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 14:58:11 -0700
Subject: [PATCH 28/55] linting test_interactions.py, test_losses.py, and
 test_measurements.py

---
 tests/tools/test_interactions.py | 212 ++++++++++++++-----------------
 tests/tools/test_losses.py       |  70 +++++-----
 tests/tools/test_measurements.py |  83 ++++++------
 3 files changed, 167 insertions(+), 198 deletions(-)

diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py
index ae3d6bd..563d8f6 100644
--- a/tests/tools/test_interactions.py
+++ b/tests/tools/test_interactions.py
@@ -1,9 +1,10 @@
-from cdtools.tools import interactions
 import numpy as np
-import torch as t
-from numpy import fft
 from numpy.fft import fftshift, ifftshift
+from numpy import fft
 import pytest
+import torch as t
+
+from cdtools.tools import interactions
 
 
 # Have a random probe and a random object and test the two
@@ -14,59 +15,60 @@ import pytest
 
 @pytest.fixture(scope='module')
 def random_probe():
-    return np.random.rand(256,256) * np.exp(2j * np.pi * np.random.rand(256,256)) 
+    return np.random.rand(256, 256) * np.exp(2j * np.pi * np.random.rand(256, 256))
+
 
 @pytest.fixture(scope='module')
 def random_obj():
-    return np.random.rand(900,900) * np.exp(2j * np.pi * np.random.rand(900,900)) 
+    return np.random.rand(900, 900) * np.exp(2j * np.pi * np.random.rand(900, 900))
+
 
 @pytest.fixture(scope='module')
 def single_pixel_probe(scope='module'):
-    probe = np.zeros((256,256), dtype=np.complex128)
-    probe[128,128] = 1
+    probe = np.zeros((256, 256), dtype=np.complex128)
+    probe[128, 128] = 1
     return probe
 
 
 def test_translations_to_pixel():
     # 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))
+    basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
+    translations = t.rand((10, 3))
     output = interactions.translations_to_pixel(basis, translations)
-    assert t.allclose(output, -translations[:,:2].flip(1))
-    
+    assert t.allclose(output, -translations[:, :2].flip(1))
+
     # Next, try a case with a single translation
     translation = t.rand((3))
     output = interactions.translations_to_pixel(basis, translation)
     assert t.allclose(output, -translation[:2].flip(0))
-    
+
     # Then, try a case with no surface normal but with a real conversion
-    basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
-    translations = t.rand((10,3))
+    basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
+    translations = t.rand((10, 3))
     output = interactions.translations_to_pixel(basis, translations)
     basis_vectors_inv = t.pinverse(basis)
-    translations[:,2] = 0 # manually project off z component
-    assert t.allclose(output, t.mm(translations,basis_vectors_inv.t()))
-    
+    translations[:, 2] = 0  # manually project off z component
+    assert t.allclose(output, t.mm(translations, basis_vectors_inv.t()))
+
     # Finally, try a case with a known surface normal (reflection)
-    basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
-    surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
-    translations = t.rand((10,3))
+    basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
+    surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
+    translations = t.rand((10, 3))
     output = interactions.translations_to_pixel(basis, translations,
                                                 surface_normal=surface_normal)
-    exp_translations = t.stack((-translations[:,1],translations[:,0]),dim=1)
+    exp_translations = t.stack((-translations[:, 1], translations[:, 0]), dim=1)
     assert t.allclose(output, exp_translations)
 
 
-
 def test_pixel_to_translations():
     # 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))
-    translations[:,2] = 0
+    basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
+    translations = t.rand((10, 3))
+    translations[:, 2] = 0
     output = interactions.translations_to_pixel(basis, translations)
     roundtrip = interactions.pixel_to_translations(basis, output)
     assert t.allclose(translations, roundtrip)
-    
+
     # Next, try a case with a single translation
     translation = t.rand((3))
     translation[2] = 0
@@ -74,65 +76,59 @@ def test_pixel_to_translations():
     roundtrip = interactions.pixel_to_translations(basis, output)
     assert t.allclose(translation, roundtrip)
 
-    
     # Then, try a case with no surface normal but with a real conversion
-    basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
-    translations = t.rand((10,3))
-    translations[:,2] = 0 # manually project off z component
+    basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
+    translations = t.rand((10, 3))
+    translations[:, 2] = 0  # manually project off z component
     output = interactions.translations_to_pixel(basis, translations)
     roundtrip = interactions.pixel_to_translations(basis, output)
     assert t.allclose(translations, roundtrip)
-    
+
     # Finally, try a case with a known surface normal (reflection)
-    basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
-    surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
-    translations = t.rand((10,3))
-    translations[:,2] = 0 # manually project off z component
+    basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
+    surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
+    translations = t.rand((10, 3))
+    translations[:, 2] = 0  # manually project off z component
     output = interactions.translations_to_pixel(basis, translations,
                                                 surface_normal=surface_normal)
     roundtrip = interactions.pixel_to_translations(basis, output,
-                                                surface_normal=surface_normal)
+                                                   surface_normal=surface_normal)
     assert t.allclose(translations, roundtrip)
 
 
-
 def test_project_translations_to_sample():
     # 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))
+    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())
+    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))
+    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())
+    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
-    translations = np.random.rand(10,2) * 500
-    exit_waves_np = [random_probe * \
-                     random_obj[tr[0]:tr[0]+random_probe.shape[0],
-                                tr[1]:tr[1]+random_probe.shape[1]] for
+    translations = np.random.rand(10, 2) * 500
+    exit_waves_np = [random_probe * random_obj[tr[0]:tr[0] + random_probe.shape[0],
+                     tr[1]:tr[1] + random_probe.shape[1]] for
                      tr in np.round(translations).astype(int)]
     exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
                                                 t.as_tensor(random_obj),
@@ -146,14 +142,13 @@ def test_ptycho_2D_round(random_probe, random_obj):
     assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0])
 
 
-
 def test_ptycho_2D_linear(single_pixel_probe, random_obj):
 
     # For this one, I just want to check one translation, but
     # I need to check both formats
-    translations = np.array([[46.7,53.2]])
-    translation = np.array([46.7,53.2])
-     
+    translations = np.array([[46.7, 53.2]])
+    translation = np.array([46.7, 53.2])
+
     exit_waves_probe = interactions.ptycho_2D_linear(
         t.as_tensor(single_pixel_probe),
         t.as_tensor(random_obj),
@@ -167,14 +162,9 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
         shift_probe=True)
 
     # Check that the outputs match
-    assert t.allclose(exit_waves_probe[0],exit_wave_probe)
+    assert t.allclose(exit_waves_probe[0], exit_wave_probe)
 
-    
-    exit_waves_obj = interactions.ptycho_2D_linear(
-        t.as_tensor(single_pixel_probe),
-        t.as_tensor(random_obj),
-        t.tensor(translations),
-        shift_probe=False)
+    exit_waves_obj = interactions.ptycho_2D_linear(t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), t.tensor(translations), shift_probe=False)
 
     exit_wave_obj = interactions.ptycho_2D_linear(
         t.as_tensor(single_pixel_probe),
@@ -183,38 +173,37 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
         shift_probe=False)
 
     # Check that the outputs match
-    assert t.allclose(exit_waves_obj[0],exit_wave_obj)
+    assert t.allclose(exit_waves_obj[0], exit_wave_obj)
 
     # For the shifted probe, we should find 4 pixels with intensity
     exit_waves_probe = t.as_tensor(exit_waves_probe)[0]
 
-    probe_shift = np.array([[0.3*0.8,0.3*0.2],
-                            [0.7*0.8,0.7*0.2]])
-    obj_section = random_obj[128+46:128+48,
-                             128+53:128+55]
-    exit_section = exit_waves_probe[128:130,128:130]
+    probe_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
+                            [0.7 * 0.8, 0.7 * 0.2]])
+    obj_section = random_obj[128 + 46:128 + 48,
+                             128 + 53:128 + 55]
+    exit_section = exit_waves_probe[128:130, 128:130]
     assert np.allclose(probe_shift * obj_section, exit_section)
-    
+
     # For the shifted obj, we should find one pixel with intensity
     exit_waves_obj = t.as_tensor(exit_waves_obj)[0]
-    obj_shift = np.array([[0.3*0.8,0.3*0.2],
-                            [0.7*0.8,0.7*0.2]])
-    obj_section = random_obj[128+46:128+48,
-                             128+53:128+55]
-    exit_pixel = exit_waves_obj[128,128]
-    assert np.isclose(np.sum(obj_shift * obj_section),exit_pixel)
-    
+    obj_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
+                          [0.7 * 0.8, 0.7 * 0.2]])
+    obj_section = random_obj[128 + 46:128 + 48,
+                             128 + 53:128 + 55]
+    exit_pixel = exit_waves_obj[128, 128]
+    assert np.isclose(np.sum(obj_shift * obj_section), exit_pixel)
+
     # Test for a single translation
 
 
 def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
-    
-    
+
     # For this one, I just want to check one translation, but
     # I need to check both formats
-    translations = np.array([[46.7,53.2]])
-    translation = np.array([46.7,53.2])
-     
+    translations = np.array([[46.7, 53.2]])
+    translation = np.array([46.7, 53.2])
+
     exit_waves_probe = interactions.ptycho_2D_sinc(
         t.as_tensor(single_pixel_probe),
         t.as_tensor(random_obj),
@@ -228,32 +217,31 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
         shift_probe=True)
 
     # Check that the outputs match
-    assert t.allclose(exit_waves_probe[0],exit_wave_probe)
-
+    assert t.allclose(exit_waves_probe[0], exit_wave_probe)
 
     # Now we explicitly define what the sinc interpolated array should
     # look like
     xs = np.arange(256) - 128
-    Ys,Xs = np.meshgrid(xs,xs)
+    Ys, Xs = np.meshgrid(xs, xs)
     sinc_probe = np.sinc(Xs) * np.sinc(Ys)
     # Just check that the unshifted probe is correct
     assert np.allclose(single_pixel_probe, sinc_probe)
 
-    sinc_shifted_probe = np.sinc(Xs-0.7) * np.sinc(Ys-0.2)
-    obj_section = random_obj[46:46+256,
-                             53:53+256]
+    sinc_shifted_probe = np.sinc(Xs - 0.7) * np.sinc(Ys - 0.2)
+    obj_section = random_obj[46:46 + 256,
+                             53:53 + 256]
     exit_wave_np = sinc_shifted_probe * obj_section
-    
+
     exit_wave_torch = exit_wave_probe.numpy()
 
     # The fidelity isn't great due to the FFT-based approach, so we need
     # a pretty relaxed condition
-    assert np.max(np.abs(exit_wave_np-exit_wave_torch)) < 0.005
+    assert np.max(np.abs(exit_wave_np - exit_wave_torch)) < 0.005
 
 
 def test_RPI_interaction(random_probe, random_obj):
 
-    random_obj1 = random_obj[:79,:68] * 0 + 1
+    random_obj1 = random_obj[:79, :68] * 0 + 1
     random_probe1 = random_probe * 0 + 1
     t_random_obj1 = t.as_tensor(random_obj1)
     t_random_probe1 = t.as_tensor(random_probe1)
@@ -261,36 +249,32 @@ def test_RPI_interaction(random_probe, random_obj):
 
     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]//2 - 79//2:
-             -(random_probe1.shape[0]-79 - (random_probe1.shape[0]//2 - 79//2)),
-             (random_probe1.shape[1]-68)//2:
-             (random_probe1.shape[1]-68)//2 + 68] = obj1_fourier
+    obj1_ups[random_probe1.shape[0] // 2 - 79 // 2:
+             -(random_probe1.shape[0] - 79 - (random_probe1.shape[0] // 2 - 79 // 2)),
+             (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'))
+                                                 norm='ortho'))
 
-    output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1]
-                       / (random_obj1.shape[-2] * random_obj1.shape[-1]))
+    output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] / (random_obj1.shape[-2] * random_obj1.shape[-1]))
 
     assert np.allclose(t_output1, output1)
-    
-    random_obj2 = np.stack([random_obj[:64,:89]]*3)
-    random_probe2 = random_probe[3:,5:]
+
+    random_obj2 = np.stack([random_obj[:64, :89]] * 3)
+    random_probe2 = random_probe[3:, 5:]
     t_random_obj2 = t.as_tensor(random_obj2)
     t_random_probe2 = t.as_tensor(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
+    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'))
+                                                 norm='ortho'))
 
-    output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1]
-                       / (random_obj2.shape[-2] * random_obj2.shape[-1]))
+    output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] / (random_obj2.shape[-2] * random_obj2.shape[-1]))
 
-    
     assert np.allclose(t_output2, output2)
-    
diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py
index 633027d..b36b85d 100644
--- a/tests/tools/test_losses.py
+++ b/tests/tools/test_losses.py
@@ -1,79 +1,73 @@
-from cdtools.tools import losses
 import numpy as np
 import torch as t
 
+from cdtools.tools import losses
+
 
 # The idea here is to use a simple numpy calculation of the various
 # objective functions to check the torch implementations and make sure
 # that any optimizations in the future don't change the results
 
+
 def test_amplitude_mse():
-    
     # Make some fake data
-    data = np.random.rand(10,100,100)
+    data = np.random.rand(10, 100, 100)
     # And add some noise to it
-    sim = data + 0.1 * np.random.rand(10,100,100)
+    sim = data + 0.1 * np.random.rand(10, 100, 100)
     # and define a simple mask that needs to be broadcast
-    mask = (np.random.rand(100,100) > 0.1).astype(bool)
+    mask = (np.random.rand(100, 100) > 0.1).astype(bool)
 
     # First, test without a mask
     np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
-    #np_result /=  data.size
-    torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim))
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
+    # np_result /=  data.size
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
     # Then, test with a mask
     np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
-    #np_result /=  np.count_nonzero(mask * np.ones_like(data))
-    torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim),
-                         mask = t.from_numpy(mask))
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
+    # np_result /=  np.count_nonzero(mask * np.ones_like(data))
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
 
 def test_intensity_mse():
     # Make some fake data
-    data = np.random.rand(10,100,100)
+    data = np.random.rand(10, 100, 100)
     # And add some noise to it
-    sim = data + 0.1 * np.random.rand(10,100,100)
+    sim = data + 0.1 * np.random.rand(10, 100, 100)
     # and define a simple mask that needs to be broadcast
-    mask = (np.random.rand(100,100) > 0.1).astype(bool)
-
+    mask = (np.random.rand(100, 100) > 0.1).astype(bool)
 
     # First, test without a mask
     np_result = np.sum((data - sim)**2)
-    np_result /=  data.size
-    torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim))
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
+    np_result /= data.size
+    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim))
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
     # Then, test with a mask
     np_result = np.sum(mask * (data - sim)**2)
-    np_result /=  np.count_nonzero(mask * np.ones_like(data))
-    torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim),
-                         mask = t.from_numpy(mask))
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
-    
+    np_result /= np.count_nonzero(mask * np.ones_like(data))
+    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
+
 
 def test_poisson_nll():
     # Make some fake data
-    data = np.random.rand(10,100,100)
+    data = np.random.rand(10, 100, 100)
     # And add some noise to it
-    sim = data + 0.1 * np.random.rand(10,100,100)
+    sim = data + 0.1 * np.random.rand(10, 100, 100)
     # and define a simple mask that needs to be broadcast
-    mask = (np.random.rand(100,100) > 0.1).astype(bool)
-
+    mask = (np.random.rand(100, 100) > 0.1).astype(bool)
 
     # First, test without a mask
     np_result = np.sum(sim - data * np.log(sim))
-    np_result /=  data.size
-    torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), eps=0)
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
+    np_result /= data.size
+    torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), eps=0)
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
     # Then, test with a mask
     np_result = np.sum(mask * (sim - data * np.log(sim)))
-    np_result /=  np.count_nonzero(mask * np.ones_like(data))
-    torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim),
-                                      mask = t.from_numpy(mask), eps=0)
-    assert np.isclose(np_result, np.take(torch_result.numpy(),0))
-
-
-
+    np_result /= np.count_nonzero(mask * np.ones_like(data))
+    torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim),
+                                      mask=t.from_numpy(mask), eps=0)
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py
index bfa06de..93ac72c 100644
--- a/tests/tools/test_measurements.py
+++ b/tests/tools/test_measurements.py
@@ -1,99 +1,90 @@
-from cdtools.tools import measurements
 import torch as t
 import numpy as np
 
+from cdtools.tools import measurements
+
 
 def test_intensity():
-    wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
-    epsilon=1e-6
+    wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
+    epsilon = 1e-6
     np_result = np.abs(wavefields.numpy())**2 + epsilon
-    assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon),
+    assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon),
                       t.as_tensor(np_result))
 
     # Test single field case
-    assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon),
+    assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon),
                       t.as_tensor(np_result[0]))
 
-    
-    det_slice = np.s_[3:,5:8]
-    assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon),
-                      t.as_tensor(np_result[(np.s_[:],)+det_slice]))
-    
+    det_slice = np.s_[3:, 5:8]
+    assert t.allclose(measurements.intensity(wavefields, det_slice, epsilon=epsilon),
+                      t.as_tensor(np_result[(np.s_[:],) + det_slice]))
+
     # Test single field case
-    assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon),
+    assert t.allclose(measurements.intensity(wavefields[0], det_slice, epsilon=epsilon),
                       t.as_tensor(np_result[0][det_slice]))
 
-
     # With oversampling on
-    np_oversampling_result = (np_result[:,::2,::2] + \
-        np_result[:,1::2,::2] + \
-        np_result[:,::2,1::2] + \
-        np_result[:,1::2,1::2]) / 4
+    np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
 
     # With multiple fields
-    assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2),
+    assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon, oversampling=2),
                       t.as_tensor(np_oversampling_result,))
 
     # With a single field
-    assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2),
+    assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon, oversampling=2),
                       t.as_tensor(np_oversampling_result[0],))
-    
+
 
 def test_incoherent_sum():
 
     # With no explicit slice given
-    
-    wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10))
-    epsilon=1e-6
-    np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon
-    assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
+
+    wavefields = t.rand((5, 4, 10, 10)) + 1j * t.rand((5, 4, 10, 10))
+    epsilon = 1e-6
+    np_result = np.sum(np.abs(wavefields.numpy())**2, axis=-3) + epsilon
+    assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon),
                       t.as_tensor(np_result))
     # Test single field case
-    assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon),
+    assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon),
                       t.as_tensor(np_result[0]))
 
-
     # With a slice given
-    det_slice = np.s_[3:,5:8]
-    assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
-                      t.as_tensor(np_result[(np.s_[:],)+det_slice]))
+    det_slice = np.s_[3:, 5:8]
+    assert t.allclose(measurements.incoherent_sum(wavefields, det_slice, epsilon=epsilon),
+                      t.as_tensor(np_result[(np.s_[:],) + det_slice]))
     # Test single field case
-    assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon),
+    assert t.allclose(measurements.incoherent_sum(wavefields[0, :], det_slice, epsilon=epsilon),
                       t.as_tensor(np_result[0][det_slice]))
 
     # With oversampling on
-    np_oversampling_result = (np_result[:,::2,::2] + \
-        np_result[:,1::2,::2] + \
-        np_result[:,::2,1::2] + \
-        np_result[:,1::2,1::2]) / 4
+    np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
 
     # With multiple fields
-    assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2),
+    assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon, oversampling=2),
                       t.as_tensor(np_oversampling_result,))
 
     # With a single field
-    assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2),
+    assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon, oversampling=2),
                       t.as_tensor(np_oversampling_result[0],))
 
 
 def test_quadratic_background():
     # test with intensity
-    wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
-    epsilon=1e-6
-    background = t.rand((10,10))
+    wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
+    epsilon = 1e-6
+    background = t.rand((10, 10))
     np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
-    det_slice = np.s_[3:,5:8]
+    det_slice = np.s_[3:, 5:8]
 
-    result = measurements.quadratic_background(wavefields,background[det_slice],
+    result = measurements.quadratic_background(wavefields, background[det_slice],
                                                detector_slice=det_slice,
                                                epsilon=epsilon,
                                                measurement=measurements.intensity)
-    assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice]))
-    
-    
+    assert t.allclose(result, t.tensor(np_result[(np.s_[:],) + det_slice]))
+
     # test with incoherent sum but no slice and no stack
-    wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10))
-    np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0)
+    wavefields = t.rand((4, 10, 10)) + 1j * t.rand((4, 10, 10))
+    np_result = np.sum(np.abs(wavefields.numpy())**2, axis=0)
     np_result += background.numpy()**2
     result = measurements.quadratic_background(wavefields, background,
                                                epsilon=epsilon,

From fc5a76ad6118c77269329219cd2cf7d04f1f7455 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 7 Jul 2025 15:13:11 -0700
Subject: [PATCH 29/55] linting test_plotting.py and test_propagators.py

---
 tests/tools/test_plotting.py    |  27 ++-
 tests/tools/test_propagators.py | 405 +++++++++++++++-----------------
 2 files changed, 209 insertions(+), 223 deletions(-)

diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py
index b377ef6..baca33d 100644
--- a/tests/tools/test_plotting.py
+++ b/tests/tools/test_plotting.py
@@ -1,47 +1,50 @@
-from cdtools.tools import plotting
-from cdtools.tools import initializers
 import numpy as np
 import torch as t
 import scipy.datasets
 import matplotlib.pyplot as plt
 
+from cdtools.tools import plotting
+from cdtools.tools import initializers
+
+
 def test_plot_amplitude(show_plot):
     # Test with tensor
-    im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
-    plotting.plot_amplitude(im, basis = np.array([[0,-1], [-1,0], [0,0]]), title = 'Test Amplitude')
+    im = t.as_tensor(scipy.datasets.ascent(), dtype=t.complex128)
+    plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude')
     if show_plot:
         plt.show()
 
     # Test with numpy array
     im = scipy.datasets.ascent().astype(np.complex128)
-    plotting.plot_amplitude(im, title = 'Test Amplitude')
+    plotting.plot_amplitude(im, title='Test Amplitude')
     if show_plot:
         plt.show()
 
 
 def test_plot_phase(show_plot):
     # Test with tensor
-    im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
-    plotting.plot_phase(im, title = 'Test Phase')
+    im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
+    plotting.plot_phase(im, title='Test Phase')
     if show_plot:
         plt.show()
 
     # Test with numpy array
-    im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
-    plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[0,-1], [-1,0], [0,0]]))
+    im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy()
+    plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
     if show_plot:
         plt.show()
 
+
 def test_plot_colorized(show_plot):
     # Test with tensor
-    gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
+    gaussian = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
     im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
-    plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[0,-1], [-1,0], [0,0]]))
+    plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
     if show_plot:
         plt.show()
 
     # Test with numpy array
     im = im.numpy()
-    plotting.plot_colorized(im, title = 'Test Colorize')
+    plotting.plot_colorized(im, title='Test Colorize')
     if show_plot:
         plt.show()
diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py
index ec55140..d571b41 100644
--- a/tests/tools/test_propagators.py
+++ b/tests/tools/test_propagators.py
@@ -1,7 +1,3 @@
-from cdtools.tools import initializers
-from cdtools.tools import propagators
-from cdtools.tools import image_processing
-
 import numpy as np
 import torch as t
 import pytest
@@ -9,35 +5,37 @@ import scipy.datasets
 from scipy import stats
 from matplotlib import pyplot as plt
 
+from cdtools.tools import initializers
+from cdtools.tools import propagators
+from cdtools.tools import image_processing
+
 
 @pytest.fixture(scope='module')
 def exit_waves_1():
     # Import scipy test image and add a random phase
-    obj = scipy.datasets.ascent()[0:64,0:64].astype(np.complex128)
-    arr = np.random.random_sample((64,64))
-    obj *= (arr+(1-arr**2)**.5*1j)
+    obj = scipy.datasets.ascent()[0:64, 0:64].astype(np.complex128)
+    arr = np.random.random_sample((64, 64))
+    obj *= (arr + (1 - arr**2)**.5 * 1j)
     obj = t.as_tensor(obj)
 
     # Construct wavefront from image
     probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
     return probe * obj
 
-    
 
 def test_far_field(exit_waves_1):
     # Far field diffraction patterns calculated by numpy with zero frequency in center
-    np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))
-
-    assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
+    np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho'))
 
+    assert (np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
 
 
 def test_inverse_far_field(exit_waves_1):
     # We want the inverse far field to map back to the exit waves with no intensity corrections
     # Far field result for exit waves calculated with numpy
-    far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')))
+    far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho')))
 
-    assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
+    assert (np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
 
 
 def test_generate_high_NA_k_intensity_map():
@@ -45,9 +43,9 @@ def test_generate_high_NA_k_intensity_map():
     # We need to generate a plausible scenario. I will start
     # by using the initializer to generate a reasonable exit wave geometry
     # and detector pair
-    basis = t.Tensor([[0,-30e-6,0],
-                      [-20e-6,0,0]]).transpose(0,1)
-    shape = t.Size([478,573])
+    basis = t.Tensor([[0, -30e-6, 0],
+                      [-20e-6, 0, 0]]).transpose(0, 1)
+    shape = t.Size([478, 573])
     wavelength = 1e-9
     distance = 1
     rs_basis = \
@@ -60,7 +58,7 @@ def test_generate_high_NA_k_intensity_map():
     # generate a good test exit wave
     i = (np.arange(478) - 240)
     j = (np.arange(573) - 270)
-    Is,Js = np.meshgrid(i,j,indexing='ij')
+    Is, Js = np.meshgrid(i, j, indexing='ij')
     wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
     t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
 
@@ -74,37 +72,36 @@ def test_generate_high_NA_k_intensity_map():
     # Checking first that for a low-NA propagation they give the same result
     # 1e-4 tolerance seems to be reasonable in this comparison given my
     # exploration with the code
-    #assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
+    # assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
 
     # Now I will explore some results with a tilted sample
-    #print(rs_basis)
-    #print(rs_basis_tilted)
-    distance = 0.01#6e-3
-    rs_basis = \
-        initializers.exit_wave_geometry(basis, shape, wavelength, distance)
+    # print(rs_basis)
+    # print(rs_basis_tilted)
+    distance = 0.01  # 6e-3
+    rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
     rs_basis_tilted = rs_basis.clone()
-    rs_basis_tilted[2,1] = rs_basis_tilted[0,1]
+    rs_basis_tilted[2, 1] = rs_basis_tilted[0, 1]
 
-    
     k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
         rs_basis_tilted, basis, shape, distance, wavelength,
         dtype=t.float32)
-    
+
     high_NA_propagated = propagators.high_NA_far_field(
         t_wavefield, k_map, intensity_map=intensity_map)
     low_NA_propagated = propagators.far_field(t_wavefield)
 
     low_NA = low_NA_propagated.numpy()
     high_NA = high_NA_propagated.numpy()
-    
-    #plt.close('all')
-    #plt.imshow(np.abs(low_NA)) 
-    #plt.figure()
-    #plt.imshow(np.abs(high_NA))
-    #plt.colorbar()
-    #plt.imshow(np.abs(wavefield))
-    #plt.show()
 
+    print(low_NA.shape, high_NA.shape)
+
+    # plt.close('all')
+    # plt.imshow(np.abs(low_NA))
+    # plt.figure()
+    # plt.imshow(np.abs(high_NA))
+    # plt.colorbar()
+    # plt.imshow(np.abs(wavefield))
+    # plt.show()
 
     # Now I want to test that it doesn't crash for wavefields of various shapes
     propagators.high_NA_far_field(t_wavefield.unsqueeze(0),
@@ -114,8 +111,9 @@ def test_generate_high_NA_k_intensity_map():
 
     # I believe this works, but I still would like to get a second method for
     # simulating at least one diffraction pattern as an independent check
-    
-    #assert 0
+
+    # assert 0
+
 
 def test_near_field_direction():
     #
@@ -135,19 +133,19 @@ def test_near_field_direction():
 
     x = (np.arange(901) - 400)
     y = (np.arange(1200) - 500)
-    Ys,Xs = np.meshgrid(y,x)
-    Rs = np.sqrt(Xs**2+Ys**2)
-    
-    E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)),dtype=t.complex64)
+    Ys, Xs = np.meshgrid(y, x)
+    Rs = np.sqrt(Xs**2 + Ys**2)
+
+    E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)), dtype=t.complex64)
     E0_real = propagators.inverse_far_field(E0_fourier)
     # This is in the top-left corner in Fourier space
 
-    wavelength = 3e-9 #nm
-    z = 1000e-9
+    wavelength = 3e-9  # nm
+    z = 1000e-9  # nm
     asp = propagators.generate_angular_spectrum_propagator(
-        E0_real.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex64)
+        E0_real.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex64)
 
-    Ez_real = propagators.near_field(E0_real,asp)
+    Ez_real = propagators.near_field(E0_real, asp)
 
     centroid = image_processing.centroid(t.abs(Ez_real))
 
@@ -155,35 +153,35 @@ def test_near_field_direction():
     assert centroid[0] < Ez_real.shape[0] // 2
     # Assert it's in left half
     assert centroid[1] < Ez_real.shape[1] // 2
-    
-    #plt.imshow(t.abs(E0_fourier))
-    #plt.figure()
-    #plt.imshow(t.abs(E0_real))    
-    #plt.figure()
-    #plt.imshow(t.abs(Ez_real))
-    #plt.show()
 
-    
+    # plt.imshow(t.abs(E0_fourier))
+    # plt.figure()
+    # plt.imshow(t.abs(E0_real))
+    # plt.figure()
+    # plt.imshow(t.abs(Ez_real))
+    # plt.show()
+
+
 def test_near_field():
 
     # The strategy is to compare the propagation of a gaussian beam to
     # the propagation in the paraxial approximation.
-    
+
     x = (np.arange(901) - 450) * 1.5e-9
     y = (np.arange(1200) - 600) * 1e-9
-    Ys,Xs = np.meshgrid(y,x)
-    Rs = np.sqrt(Xs**2+Ys**2)
+    Ys, Xs = np.meshgrid(y, x)
+    Rs = np.sqrt(Xs**2 + Ys**2)
 
-    wavelength = 3e-9 #nm
-    sigma = 20e-9 #nm
-    z = 1000e-9 #nm
+    wavelength = 3e-9  # nm
+    sigma = 20e-9  # nm
+    z = 1000e-9  # nm
 
     k = 2 * np.pi / wavelength
-    w0 = np.sqrt(2)*sigma
+    w0 = np.sqrt(2) * sigma
     zr = np.pi * w0**2 / wavelength
     wz = w0 * np.sqrt(1 + (z / zr)**2)
-    Rz =  z * (1 + (zr / z)**2)
-    
+    Rz = z * (1 + (zr / z)**2)
+
     E0 = np.exp(-Rs**2 / w0**2)
 
     # The analytical expression for propagation of a gaussian beam in the
@@ -208,52 +206,49 @@ def test_near_field():
     #    If we choose e^(-ikx) to represent light propagating along K, the
     #    answer is no, and we find we have to use the inverse FT instead.
     #    Thus, e^(ikx) is the right choice here.
-    
-    Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * ( z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
+
+    Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * (z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
     Ez_nozphase = Ez * np.exp(-1j * k * z)
-    
 
     # First we check it normally
     asp = propagators.generate_angular_spectrum_propagator(
-        E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
+        E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
+
+    Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
 
-    
-    Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
-    
     # Check for at least 10^-3 relative accuracy in this scenario
-    assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
-
+    assert np.max(np.abs(Ez_nozphase - Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
 
     Emz = np.conj(Ez_nozphase)
 
-    Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
+    Emz_t = propagators.inverse_near_field(t.as_tensor(E0), asp).numpy()
 
     # Again, 10^-3 is about all the accuracy we can expect
-    assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
+    assert np.max(np.abs(Emz - Emz_t)) < 1e-3 * np.max(np.abs(Emz))
 
     # Then, we check that the bandlimiting at least does something
     asp = propagators.generate_angular_spectrum_propagator(
-        E0.shape,(1.5e-9,1e-9),wavelength,z,
+        E0.shape, (1.5e-9, 1e-9), wavelength, z,
         dtype=t.complex128, bandlimit=0.3)
 
-    assert asp[140,0] == 0
-    assert asp[0,180] == 0
-    assert asp[130,0] != 0
-    assert asp[0,175] != 0
-    
+    assert asp[140, 0] == 0
+    assert asp[0, 180] == 0
+    assert asp[130, 0] != 0
+    assert asp[0, 175] != 0
+
     # Then, we check that automatic differentiation works
-    z = t.tensor([z],requires_grad=True)
-    spacing = t.tensor((1.5e-9,1e-9), requires_grad=True)
+    z = t.tensor([z], requires_grad=True)
+    spacing = t.tensor((1.5e-9, 1e-9), requires_grad=True)
     wavelength = t.tensor([wavelength], requires_grad=True)
 
     asp = propagators.generate_angular_spectrum_propagator(
         E0.shape, spacing, wavelength, z)
 
-    t.real(asp[10,10]).backward()
+    t.real(asp[10, 10]).backward()
     assert z.grad != 0
     assert spacing.grad[0] != 0
-    assert wavelength.grad !=0
-    
+    assert wavelength.grad != 0
+
 
 def test_generalized_near_field():
 
@@ -264,26 +259,25 @@ def test_generalized_near_field():
     # First, we should do a test with the phase ramp along the z direction
     # explicitly included
 
-    basis= np.array([[0,-1.5e-9],[-1e-9,0],[0,0]])
-    i_vec,j_vec = np.arange(901) - 450 ,np.arange(1200)-600
-    Is, Js = np.meshgrid(i_vec,j_vec,indexing='ij')
-    Xs_0,Ys_0,Zs_0 = np.tensordot(basis,np.stack([Is,Js]),axes=1)
-    #x = (np.arange(901) - 450) * 1.5e-9
-    #y = (np.arange(1200) - 600) * 1e-9
-    #Xs_0,Ys_0 = np.meshgrid(x,y)
-    #Zs_0 = np.zeros(Xs_0.shape)
+    basis = np.array([[0, -1.5e-9], [-1e-9, 0], [0, 0]])
+    i_vec, j_vec = np.arange(901) - 450, np.arange(1200) - 600
+    Is, Js = np.meshgrid(i_vec, j_vec, indexing='ij')
+    Xs_0, Ys_0, Zs_0 = np.tensordot(basis, np.stack([Is, Js]), axes=1)
+    # x = (np.arange(901) - 450) * 1.5e-9
+    # y = (np.arange(1200) - 600) * 1e-9
+    # Xs_0,Ys_0 = np.meshgrid(x,y)
+    # Zs_0 = np.zeros(Xs_0.shape)
+
+    Positions = np.stack([Xs_0, Ys_0, Zs_0])
 
-    Positions = np.stack([Xs_0,Ys_0,Zs_0])
-    
     # assert 0
-    wavelength = 3e-9 #nm
-    sigma = 20e-9 #nm
-    z = 1000e-9 #nm
+    wavelength = 3e-9  # nm
+    sigma = 20e-9  # nm
+    z = 1000e-9  # nm
 
     k = 2 * np.pi / wavelength
-    w0 = np.sqrt(2)*sigma
+    w0 = np.sqrt(2) * sigma
     zr = np.pi * w0**2 / wavelength
-    
 
     # The analytical expression for propagation of a gaussian beam in the
     # paraxial approx
@@ -292,140 +286,132 @@ def test_generalized_near_field():
 
     def get_inv_R(Zs):
         return Zs / (Zs**2 + zr**2)
-    
+
     def get_E(Xs, Ys, Zs, correct=True):
         # Again, this follows the convention opposite from Wikipedia. See
         # the note in test_angular_spectrum_propagator.
-        
+
         # if correct is True, remove the e^(ikz) dependence
-        
+
         Rs_sq = Xs**2 + Ys**2
         Wzs = get_w(Zs)
         E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\
-            np.exp(1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \
-                   - 1j * np.arctan(Zs / zr))
-        
+            np.exp(1j * k * (Zs + Rs_sq * get_inv_R(Zs) / 2) - 1j * np.arctan(Zs / zr))
+
         # This removes the z-dependence of the phase
-        if correct: 
+        if correct:
             E = E * np.exp(-1j * k * Zs)
         return E
-    
+
     def check_equiv(analytical, numerical):
-        phase = np.angle(np.mean(numerical.conj()*analytical))
-        comp = np.exp(1j*phase) * numerical
-        return (np.max(np.abs(analytical-comp))
-                < 1e-3 * np.max(np.abs(analytical)))
-    
+        phase = np.angle(np.mean(numerical.conj() * analytical))
+        comp = np.exp(1j * phase) * numerical
+        return (np.max(np.abs(analytical - comp)) < 1e-3 * np.max(np.abs(analytical)))
+
     # We make some rotation matrices to test
-    
+
     # This tests the straight ahead case
-    I = np.eye(3)
+    IdentityMatrix = np.eye(3)
 
     # This tests a rotation about the y axis
     th = np.deg2rad(5)
-    Ry = np.array([[np.cos(th),0,np.sin(th)],
-                   [0,1,0],
-                   [-np.sin(th),0,np.cos(th)]])
-    
+    Ry = np.array([[np.cos(th), 0, np.sin(th)],
+                   [0, 1, 0],
+                   [-np.sin(th), 0, np.cos(th)]])
+
     # This tests a rotation about two axes
     phi = np.deg2rad(2)
-    Rx = np.array([[1,0,0],
-                   [0,np.cos(phi),-np.sin(phi)],
-                   [0,np.sin(phi),np.cos(phi)]])
-    Rboth = np.matmul(Rx,Ry)
+    Rx = np.array([[1, 0, 0],
+                   [0, np.cos(phi), -np.sin(phi)],
+                   [0, np.sin(phi), np.cos(phi)]])
+    Rboth = np.matmul(Rx, Ry)
 
     # This tests a shearing
     shear = 0.23
-    Rshear = np.array([[1,shear,0],
-                       [0,1,0],
-                       [0,0,1]])
+    Rshear = np.array([[1, shear, 0],
+                       [0, 1, 0],
+                       [0, 0, 1]])
 
-    
     # This tests an inversion of the axes
-    Rinv = np.array([[-1,0,0],
-                     [0,-1,0],
-                     [0,0,-1]])
+    Rinv = np.array([[-1, 0, 0],
+                     [0, -1, 0],
+                     [0, 0, -1]])
 
     # This tests a reflection about the y-z plane
-    Rrefl = np.array([[-1,0,0],
-                      [0,1,0],
-                      [0,0,-1]])
-    
-    # This tests a shearing and a rotation together
-    Rall = np.matmul(Rrefl,np.matmul(Rboth,Rshear))
+    Rrefl = np.array([[-1, 0, 0],
+                      [0, 1, 0],
+                      [0, 0, -1]])
 
+    # This tests a shearing and a rotation together
+    Rall = np.matmul(Rrefl, np.matmul(Rboth, Rshear))
 
     # And we make some propagation vectors to test:
 
     # This is along the z direction
-    z_dir = np.array([0,0,1])
+    z_dir = np.array([0, 0, 1])
 
     # This checks that it's not sensitive to the magnitude
-    z_dir_large = np.array([0,0,10])
-    
+    z_dir_large = np.array([0, 0, 10])
+
     # And finally some offset vectors
 
     # This checks straight ahead
-    z_offset = np.array([0,0,z])
+    z_offset = np.array([0, 0, z])
 
-    
     # This checks with an offset in x and y
-    shear_offset = np.array([0.1*z,-0.03*z,z])
+    shear_offset = np.array([0.1 * z, -0.03 * z, z])
 
     # This checks with an offset in x and y, with negative z
-    shear_back_offset = np.array([0.1*z,-0.03*z,-z])
+    shear_back_offset = np.array([0.1 * z, -0.03 * z, -z])
 
-    
-    rot_mats = [Rrefl,I,Rinv, Rboth, Rboth,Rboth, Rall, Rall, Rall, I, Rall]
-    offset_vecs = [z_offset]*8 + [shear_offset] + [shear_back_offset]*2
-    propagation_vecs = ['perp','offset',z_dir,
-                        'perp','offset',z_dir_large,
-                        'perp','offset',z_dir_large,
+    rot_mats = [Rrefl, IdentityMatrix, Rinv, Rboth, Rboth, Rboth, Rall, Rall, Rall, IdentityMatrix, Rall]
+    offset_vecs = [z_offset] * 8 + [shear_offset] + [shear_back_offset] * 2
+    propagation_vecs = ['perp', 'offset', z_dir,
+                        'perp', 'offset', z_dir_large,
+                        'perp', 'offset', z_dir_large,
                         z_dir, z_dir_large]
-    purposes = ['standard']*3 + ['both-rot']*3 + ['shear-rot']*3 + ['backward']*2
-    
-    #rot_mats = [Ry.transpose()]
-    #offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
+    purposes = ['standard'] * 3 + ['both-rot'] * 3 + ['shear-rot'] * 3 + ['backward'] * 2
+
+    # rot_mats = [Ry.transpose()]
+    # offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
     #                  np.dot(Ry.transpose(),basis)[:,1])
-    #offset /= np.linalg.norm(offset) / 3e-6
-    #offset_vecs = [-offset]#[shear_offset]
-    #propagation_vecs = [z_dir]
-    #purposes=['meh']
-    
-    for purpose,rot_mat,offset_vec, propagation_vec \
-        in zip(purposes,rot_mats,offset_vecs,propagation_vecs):
-        
+    # offset /= np.linalg.norm(offset) / 3e-6
+    # offset_vecs = [-offset]#[shear_offset]
+    # propagation_vecs = [z_dir]
+    # purposes=['meh']
+
+    for purpose, rot_mat, offset_vec, propagation_vec in zip(purposes, rot_mats, offset_vecs, propagation_vecs):
         print('Testing', purpose)
-        Xs,Ys,Zs_0 = np.tensordot(rot_mat,Positions,axes=1)
+        Xs, Ys, Zs_0 = np.tensordot(rot_mat, Positions, axes=1)
         new_basis = np.dot(rot_mat, basis)
-        Xs_prop, Ys_prop, Zs_prop = np.stack([Xs,Ys,Zs_0]) \
-            + offset_vec[:,None,None]
+        Xs_prop, Ys_prop, Zs_prop = np.stack([Xs, Ys, Zs_0]) \
+            + offset_vec[:, None, None]
+
+        print('Propagate Along', propagation_vec)
 
-        print('Propagate Along',propagation_vec)
-        
         if str(propagation_vec) == 'perp':
-            E0 = get_E(Xs,Ys,Zs_0, correct=False)
-            Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False)
+            E0 = get_E(Xs, Ys, Zs_0, correct=False)
+            Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=False)
             asp = propagators.generate_generalized_angular_spectrum_propagator(
-                E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128)
+                E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128)
         elif str(propagation_vec) == 'offset':
-            E0 = get_E(Xs,Ys,Zs_0, correct=True)
-            Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
+            E0 = get_E(Xs, Ys, Zs_0, correct=True)
+            Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
             asp = propagators.generate_generalized_angular_spectrum_propagator(
-                E0.shape,new_basis,wavelength,offset_vec,
+                E0.shape, new_basis, wavelength, offset_vec,
                 dtype=t.complex128, propagate_along_offset=True)
         else:
-            E0 = get_E(Xs,Ys,Zs_0, correct=True)
-            Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
+            E0 = get_E(Xs, Ys, Zs_0, correct=True)
+            Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
             asp = propagators.generate_generalized_angular_spectrum_propagator(
-                    E0.shape,new_basis,wavelength,offset_vec,
-                    dtype=t.complex128, propagation_vector=propagation_vec)
+                E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128,
+                propagation_vector=propagation_vec)
+
+        Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
 
-        Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
-        
         # Check for at least 10^-3 relative accuracy in this scenario
         if not check_equiv(Ez, Ez_t):
-        #if True:
+            # if True:
             plt.close('all')
             plt.figure()
             plt.imshow(np.angle(E0))
@@ -443,70 +429,67 @@ def test_generalized_near_field():
             plt.imshow(np.angle(Ez_t))
             plt.title('Angle of numerically calculated Ez')
             plt.figure()
-            plt.imshow(np.abs(Ez-Ez_t))#/np.max(np.abs(Ez)))
+            plt.imshow(np.abs(Ez - Ez_t))
             plt.title('Magnitude of difference')
             plt.show()
-            
+
         assert check_equiv(Ez, Ez_t)
-        
-            
-        Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
-        
-        assert check_equiv(E0,Em0_t)
-            
+
+        Em0_t = propagators.inverse_near_field(t.as_tensor(Ez), asp).numpy()
+
+        assert check_equiv(E0, Em0_t)
+
         print('Test Successful')
 
     # One final test, to see if any of a few arbitrary rotations will
     # change the predicted propagation if everything else is kept
     # constant
     Rrands = [stats.ortho_group.rvs(3) for i in range(3)]
-    Xs,Ys,Zs_0 = np.tensordot(Rboth,Positions,axes=1)
+    Xs, Ys, Zs_0 = np.tensordot(Rboth, Positions, axes=1)
     new_basis = np.dot(rot_mat, basis)
     offset_vec = shear_back_offset
     propagation_vec = z_dir
-    
-    E0 = get_E(Xs,Ys,Zs_0, correct=True)
+
+    E0 = get_E(Xs, Ys, Zs_0, correct=True)
     asp = propagators.generate_generalized_angular_spectrum_propagator(
-                    E0.shape, new_basis, wavelength,offset_vec,
-                    dtype=t.complex128, propagation_vector=propagation_vec)
-    Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
-        
+        E0.shape, new_basis, wavelength, offset_vec,
+        dtype=t.complex128, propagation_vector=propagation_vec)
+    Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
 
     for Rrand in Rrands:
-        Xs,Ys,Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth,Positions,axes=1),axes=1)
+        Xs, Ys, Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth, Positions, axes=1), axes=1)
         rot_offset = np.dot(Rrand, offset_vec)
         rot_basis = np.dot(Rrand, new_basis)
         rot_prop = np.dot(Rrand, propagation_vec)
         asp = propagators.generate_generalized_angular_spectrum_propagator(
-                    E0.shape, rot_basis, wavelength, rot_offset,
-                    dtype=t.complex128, propagation_vector=rot_prop)
-        Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
+            E0.shape, rot_basis, wavelength, rot_offset,
+            dtype=t.complex128, propagation_vector=rot_prop)
+        Ez_rot_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
+
+        assert np.max(np.abs(Ez_t - Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
 
-        assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
-        
 
 def test_inverse_near_field():
-    
+
     x = (np.arange(800) - 400) * 1.5e-9
     y = (np.arange(1200) - 600) * 1e-9
-    Ys,Xs = np.meshgrid(y,x)
-    Rs = np.sqrt(Xs**2+Ys**2)
-    
-    wavelength = 3e-9 #nm
-    sigma = 20e-9 #nm
-    z = 1000e-9 #nm
+    Ys, Xs = np.meshgrid(y, x)
+    Rs = np.sqrt(Xs**2 + Ys**2)
 
-    w0 = np.sqrt(2)*sigma
+    wavelength = 3e-9  # nm
+    sigma = 20e-9  # nm
+    z = 1000e-9  # nm
+
+    w0 = np.sqrt(2) * sigma
     E0 = np.exp(-Rs**2 / w0**2)
 
     asp = propagators.generate_angular_spectrum_propagator(
-        E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
+        E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
 
+    E0 = t.as_tensor(E0, dtype=t.complex128)
+    E_prop = propagators.near_field(E0, asp)
 
-    E0 = t.as_tensor(E0,dtype=t.complex128)
-    E_prop = propagators.near_field(E0,asp)
-    
     E_backprop = propagators.inverse_near_field(E_prop, asp)
 
     # We just want to check that it actually is the inverse
-    assert t.all(t.isclose(E0,E_backprop))
+    assert t.all(t.isclose(E0, E_backprop))

From 0785eae552de0dcfc13d54af73cac78a5a56b4b1 Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Fri, 1 Aug 2025 21:37:05 +0000
Subject: [PATCH 30/55] Created Reconstructors class to replace optimization
 methods in CDIModel

---
 src/cdtools/__init__.py                |   4 +-
 src/cdtools/reconstructors/__init__.py |  16 ++
 src/cdtools/reconstructors/adam.py     | 160 ++++++++++++
 src/cdtools/reconstructors/base.py     | 325 +++++++++++++++++++++++++
 src/cdtools/reconstructors/lbfgs.py    | 134 ++++++++++
 src/cdtools/reconstructors/sgd.py      | 157 ++++++++++++
 6 files changed, 794 insertions(+), 2 deletions(-)
 create mode 100644 src/cdtools/reconstructors/__init__.py
 create mode 100644 src/cdtools/reconstructors/adam.py
 create mode 100644 src/cdtools/reconstructors/base.py
 create mode 100644 src/cdtools/reconstructors/lbfgs.py
 create mode 100644 src/cdtools/reconstructors/sgd.py

diff --git a/src/cdtools/__init__.py b/src/cdtools/__init__.py
index 9684207..9132209 100644
--- a/src/cdtools/__init__.py
+++ b/src/cdtools/__init__.py
@@ -4,9 +4,9 @@ import warnings
 warnings.filterwarnings("ignore",
                         message='To copy construct from a tensor, ')
 
-__all__ = ['tools', 'datasets', 'models']
+__all__ = ['tools', 'datasets', 'models', 'reconstructors']
 
 from cdtools import tools
 from cdtools import datasets
 from cdtools import models
-
+from cdtools import reconstructors
diff --git a/src/cdtools/reconstructors/__init__.py b/src/cdtools/reconstructors/__init__.py
new file mode 100644
index 0000000..84b96ab
--- /dev/null
+++ b/src/cdtools/reconstructors/__init__.py
@@ -0,0 +1,16 @@
+"""This module contains optimizers for performing reconstructions
+
+"""
+
+# We define __all__ to be sure that import * only imports what we want
+__all__ = [
+    'Reconstructor',
+    'Adam',
+    'LBFGS',
+    'SGD'
+]
+
+from cdtools.reconstructors.base import Reconstructor
+from cdtools.reconstructors.adam import Adam
+from cdtools.reconstructors.lbfgs import LBFGS
+from cdtools.reconstructors.sgd import SGD
diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
new file mode 100644
index 0000000..5a489a0
--- /dev/null
+++ b/src/cdtools/reconstructors/adam.py
@@ -0,0 +1,160 @@
+"""This module contains the Adam Reconstructor subclass for performing
+optimization ('reconstructions') on ptychographic/CDI models using
+the Adam optimizer.
+
+The Reconstructor class is designed to resemble so-called
+'Trainer' classes that (in the language of the AI/ML folks) handles
+the 'training' of a model given some dataset and optimizer.
+"""
+import torch as t
+from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+from cdtools.models import CDIModel
+from typing import Tuple, List, Union
+from cdtools.reconstructors import Reconstructor
+
+__all__ = ['Adam']
+
+
+class Adam(Reconstructor):
+    """
+    The Adam Reconstructor subclass handles the optimization ('reconstruction')
+    of ptychographic models and datasets using the Adam optimizer.
+
+    Parameters
+    ----------
+    model: CDIModel
+        Model for CDI/ptychography reconstruction.
+    dataset: Ptycho2DDataset
+        The dataset to reconstruct against.
+    subset : list(int) or int
+        Optional, a pattern index or list of pattern indices to use.
+    schedule : bool
+        Optional, create a learning rate scheduler
+        (torch.optim.lr_scheduler._LRScheduler).
+
+    Important attributes:
+    - **model** -- Always points to the core model used.
+    - **optimizer** -- This class by default uses `torch.optim.Adam` to perform
+        optimizations.
+    - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
+        `optimize` method.
+    - **data_loader** -- A torch.utils.data.DataLoader that is defined by
+        calling the `setup_dataloader` method.
+    """
+    def __init__(self,
+                 model: CDIModel,
+                 dataset: Ptycho2DDataset,
+                 subset: List[int] = None):
+
+        super().__init__(model, dataset, subset)
+
+        # Define the optimizer for use in this subclass
+        self.optimizer = t.optim.Adam(self.model.parameters())
+
+    def adjust_optimizer(self,
+                         lr: int = 0.005,
+                         betas: Tuple[float] = (0.9, 0.999),
+                         amsgrad: bool = False):
+        """
+        Change hyperparameters for the utilized optimizer.
+
+        Parameters
+        ----------
+        lr : float
+            Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
+            is typically the highest possible value with any chance of being
+            stable.
+        betas : tuple
+            Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
+        amsgrad : bool
+            Optional, whether to use the AMSGrad variant of this algorithm.
+        """
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+            param_group['betas'] = betas
+            param_group['amsgrad'] = amsgrad
+
+    def optimize(self,
+                 iterations: int,
+                 batch_size: int = 15,
+                 lr: float = 0.005,
+                 betas: Tuple[float] = (0.9, 0.999),
+                 schedule: bool = False,
+                 amsgrad: bool = False,
+                 regularization_factor: Union[float, List[float]] = None,
+                 thread: bool = True,
+                 calculation_width: int = 10,
+                 shuffle: bool = True):
+        """
+        Runs a round of reconstruction using the Adam optimizer
+
+        Formerly `CDIModel.Adam_optimize`
+
+        This calls the Reconstructor.optimize superclass method
+        (formerly `CDIModel.AD_optimize`) to run a round of reconstruction
+        once the dataloader and optimizer hyperparameters have been
+        set up.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run.
+        batch_size : int
+            Optional, the size of the minibatches to use.
+        lr : float
+            Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
+            is typically the highest possible value with any chance of being
+            stable.
+        betas : tuple
+            Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
+        schedule : bool
+            Optional, create a learning rate scheduler
+            (torch.optim.lr_scheduler._LRScheduler).
+        amsgrad : bool
+            Optional, whether to use the AMSGrad variant of this algorithm.
+        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.
+        calculation_width : int
+            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.
+        shuffle : bool
+            Optional, enable/disable shuffling of the dataset. This option
+            is intended for diagnostic purposes and should be left as True.
+        """
+        # Update the training history
+        self.model.training_history += (
+            f'Planning {iterations} epochs of Adam, with a learning rate = '
+            f'{lr}, batch size = {batch_size}, regularization_factor = '
+            f'{regularization_factor}, and schedule = {schedule}.\n'
+        )
+
+        # 1) The subset statement is contained in Reconstructor.__init__
+
+        # 2) Set up / re-initialize the data laoder
+        self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
+
+        # 3) The optimizer is created in self.__init__, but the
+        #    hyperparameters need to be set up with self.adjust_optimizer
+        self.adjust_optimizer(lr=lr,
+                              betas=betas,
+                              amsgrad=amsgrad)
+
+        # 4) Set up the scheduler
+        if schedule:
+            self.scheduler = \
+                t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
+                                                       factor=0.2,
+                                                       threshold=1e-9)
+        else:
+            self.scheduler = None
+
+        # 5) This is analagous to making a call to CDIModel.AD_optimize
+        return super(Adam, self).optimize(iterations,
+                                          regularization_factor,
+                                          thread,
+                                          calculation_width)
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
new file mode 100644
index 0000000..65b2771
--- /dev/null
+++ b/src/cdtools/reconstructors/base.py
@@ -0,0 +1,325 @@
+"""This module contains the base Reconstructor class for performing
+optimization ('reconstructions') on ptychographic/CDI models.
+
+The Reconstructor class is designed to resemble so-called
+'Trainer' classes that (in the language of the AI/ML folks) handles
+the 'training' of a model given some dataset and optimizer.
+
+The subclasses of Reconstructor are required to implement
+their own data loaders and optimizer adjusters
+"""
+
+import torch as t
+from torch.utils import data as td
+import threading
+import queue
+import time
+from cdtools.datasets import CDataset
+from cdtools.models import CDIModel
+from typing import List, Union
+
+__all__ = ['Reconstructor']
+
+
+class Reconstructor:
+    """
+    Reconstructor handles the optimization ('reconstruction') of ptychographic
+    models given a CDIModel (or subclass) and corresponding CDataset.
+
+    This is a base model that defines all functions Reconstructor subclasses
+    must implement.
+
+    Parameters
+    ----------
+    model: CDIModel
+        Model for CDI/ptychography reconstruction
+    dataset: CDataset
+        The dataset to reconstruct against
+    subset : list(int) or int
+        Optional, a pattern index or list of pattern indices to use
+
+    Important attributes:
+    - **model** -- Always points to the core model used.
+    - **optimizer** -- A `torch.optim.Optimizer` that must be defined when
+        initializing the Reconstructor subclass.
+    - **scheduler** -- A `torch.optim.lr_scheduler` that may be defined during
+        the `optimize` method.
+    - **data_loader** -- A torch.utils.data.DataLoader that is defined by
+        calling the `setup_dataloader` method.
+    """
+    def __init__(self,
+                 model: CDIModel,
+                 dataset: CDataset,
+                 subset: Union[int, List[int]] = None):
+        # Store parameters as attributes of Reconstructor
+        self.subset = subset
+
+        # Initialize attributes that must be defined by the subclasses
+        self.optimizer = None
+        self.scheduler = None
+        self.data_loader = None
+
+        # Store the original model
+        self.model = model
+
+        # Store the dataset
+        if subset is not None:
+            # if subset is just one pattern, turn into a list for convenience
+            if isinstance(subset, int):
+                subset = [subset]
+            dataset = td.Subset(dataset, subset)
+        self.dataset = dataset
+
+    def setup_dataloader(self,
+                         batch_size: int = None,
+                         shuffle: bool = True):
+        """
+        Sets up / re-initializes the dataloader.
+
+        Parameters
+        ----------
+        batch_size : int
+            Optional, the size of the minibatches to use
+        shuffle : bool
+            Optional, enable/disable shuffling of the dataset. This option
+            is intended for diagnostic purposes and should be left as True.
+        """
+        if batch_size is not None:
+            self.data_loader = td.DataLoader(self.dataset,
+                                             batch_size=batch_size,
+                                             shuffle=shuffle)
+        else:
+            self.data_loader = td.Dataloader(self.dataset)
+
+    def adjust_optimizer(self, **kwargs):
+        """
+        Change hyperparameters for the utilized optimizer.
+
+        For each optimizer, the keyword arguments should be manually defined
+        as parameters.
+        """
+        raise NotImplementedError()
+
+    def _run_epoch(self,
+                   stop_event: threading.Event = None,
+                   regularization_factor: Union[float, List[float]] = None,
+                   calculation_width: int = 10):
+        """
+        Runs one full epoch of the reconstruction. Intended to be called
+        by Reconstructor.optimize.
+
+        Parameters
+        ----------
+        stop_event : threading.Event
+            Default None, causes the reconstruction to stop when an exception
+            occurs in Optimizer.optimize.
+        regularization_factor : float or list(float)
+            Optional, if the model has a regularizer defined, the set of
+            parameters to pass the regularizer method
+        calculation_width : int
+            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.
+
+        Returns
+        ------
+        loss : float
+            The summed loss over the latest epoch, divided by the total
+            diffraction pattern intensity
+        """
+
+        # 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 self.data_loader:
+            normalization += t.sum(patterns).cpu().numpy()
+            N += 1
+
+            def closure():
+                self.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]),
+                                calculation_width)]
+                pattern_chunks = [patterns[i:i + calculation_width]
+                                  for i in range(0, len(inputs[0]),
+                                  calculation_width)]
+
+                total_loss = 0
+
+                for inp, pats in zip(input_chunks, pattern_chunks):
+                    # This 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.model.forward(*inp)
+
+                    # Calculate the loss
+                    if hasattr(self, 'mask'):
+                        loss = self.model.loss(pats,
+                                               sim_patterns,
+                                               mask=self.model.mask)
+                    else:
+                        loss = self.model.loss(pats,
+                                               sim_patterns)
+
+                    # And accumulate the gradients
+                    loss.backward()
+
+                    # Normalize the accumulating total loss
+                    total_loss += loss.detach() // self.model.world_size
+
+                # 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.model, 'regularizer'):
+
+                    loss = self.model.regularizer(regularization_factor)
+                    loss.backward()
+
+                return total_loss
+
+            # This takes the step for this minibatch
+            loss += self.optimizer.step(closure).detach().cpu().numpy()
+
+        loss /= normalization
+
+        # We step the scheduler after the full epoch
+        if self.scheduler is not None:
+            self.scheduler.step(loss)
+
+        self.model.loss_history.append(loss)
+        self.model.epoch = len(self.model.loss_history)
+        self.model.latest_iteration_time = time.time() - t0
+        self.model.training_history += self.model.report() + '\n'
+        return loss
+
+    def optimize(self,
+                 iterations: int,
+                 regularization_factor: Union[float, List[float]] = None,
+                 thread: bool = True,
+                 calculation_width: int = 10):
+        """
+        Runs a round of reconstruction using the provided optimizer
+
+        Formerly CDIModel.AD_optimize
+
+        This is the basic automatic differentiation reconstruction tool
+        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.
+
+        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.
+        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.
+        calculation_width : int
+            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.
+
+        Yields
+        ------
+        loss : float
+            The summed loss over the latest epoch, divided by the total
+            diffraction pattern intensity.
+        """
+
+        # We store the current optimizer as a model parameter so that
+        # it can be saved and loaded for checkpointing
+        self.current_optimizer = self.optimizer
+
+        # If we don't want to run in a different thread, this is easy
+        if not thread:
+            for it in range(iterations):
+                if self.model.skip_computation():
+                    self.epoch = self.epoch + 1
+                    if len(self.model.loss_history) >= 1:
+                        yield self.model.loss_history[-1]
+                    else:
+                        yield float('nan')
+                    continue
+
+                yield self._run_epoch(regularization_factor=regularization_factor, # noqa
+                                      calculation_width=calculation_width)
+
+        # 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(
+                        self._run_epoch(stop_event=stop_event,
+                                        regularization_factor=regularization_factor, # noqa
+                                        calculation_width=calculation_width))
+                except Exception as e:
+                    # If something bad happens, put the exception into the
+                    # result queue
+                    result_queue.put(e)
+
+            # And this actually starts and monitors the thread
+            for it in range(iterations):
+                if self.model.skip_computation():
+                    self.model.epoch = self.model.epoch + 1
+                    if len(self.model.loss_history) >= 1:
+                        yield self.model.loss_history[-1]
+                    else:
+                        yield float('nan')
+                    continue
+
+                calc = threading.Thread(target=target,
+                                        name='calculator',
+                                        daemon=True)
+                try:
+                    calc.start()
+                    while calc.is_alive():
+                        if hasattr(self.model, 'figs'):
+                            self.model.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
+
+                res = result_queue.get()
+
+                # If something went wrong in the thead, we'll get an exception
+                if isinstance(res, Exception):
+                    raise res
+
+                yield res
+
+        # And finally, we unset the current optimizer:
+        self.current_optimizer = None
diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py
new file mode 100644
index 0000000..0b51dfd
--- /dev/null
+++ b/src/cdtools/reconstructors/lbfgs.py
@@ -0,0 +1,134 @@
+"""This module contains the LBFGS Reconstructor subclass for performing
+optimization ('reconstructions') on ptychographic/CDI models using
+the LBFGS optimizer.
+
+The Reconstructor class is designed to resemble so-called
+'Trainer' classes that (in the language of the AI/ML folks) handles
+the 'training' of a model given some dataset and optimizer.
+"""
+import torch as t
+from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+from cdtools.models import CDIModel
+from typing import List, Union
+from cdtools.reconstructors import Reconstructor
+
+__all__ = ['LBFGS']
+
+
+class LBFGS(Reconstructor):
+    """
+    The LBFGS Reconstructor subclass handles the optimization
+    ('reconstruction') of ptychographic models and datasets using the LBFGS
+    optimizer.
+
+    Parameters
+    ----------
+    model: CDIModel
+        Model for CDI/ptychography reconstruction.
+    dataset: Ptycho2DDataset
+        The dataset to reconstruct against.
+    subset : list(int) or int
+        Optional, a pattern index or list of pattern indices to use.
+    schedule : bool
+        Optional, create a learning rate scheduler
+        (torch.optim.lr_scheduler._LRScheduler).
+
+    Important attributes:
+    - **model** -- Always points to the core model used.
+    - **optimizer** -- This class by default uses `torch.optim.LBFGS` to
+        perform optimizations.
+    - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during
+        the `optimize` method.
+    - **data_loader** -- A torch.utils.data.DataLoader that is defined by
+        calling the `setup_dataloader` method.
+    """
+    def __init__(self,
+                 model: CDIModel,
+                 dataset: Ptycho2DDataset,
+                 subset: List[int] = None):
+
+        super().__init__(model, dataset, subset)
+
+        # Define the optimizer for use in this subclass
+        self.optimizer = t.optim.LBFGS(self.model.parameters())
+
+    def adjust_optimizer(self,
+                         lr: int = 0.005,
+                         history_size: int = 2,
+                         line_search_fn: str = None):
+        """
+        Change hyperparameters for the utilized optimizer.
+
+        Parameters
+        ----------
+        lr : float
+            Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
+            is typically the highest possible value with any chance of being
+            stable.
+        history_size : int
+            Optional, the length of the history to use.
+        line_search_fn : str
+            Optional, either `strong_wolfe` or None
+        """
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+            param_group['history_size'] = history_size
+            param_group['line_search_fn'] = line_search_fn
+
+    def optimize(self,
+                 iterations: int,
+                 lr: float = 0.1,
+                 history_size: int = 2,
+                 regularization_factor: Union[float, List[float]] = None,
+                 thread: bool = True,
+                 calculation_width: int = 10,
+                 line_search_fn: str = None):
+        """
+        Runs a round of reconstruction using the LBFGS optimizer
+
+        Formerly `CDIModel.LBFGS_optimize`
+
+        This algorithm is often less stable that Adam, however in certain
+        situations or geometries it can be shockingly efficient. Like all
+        the other optimization routines, it is defined as a generator
+        function which yields the average loss each epoch.
+
+        NOTE: There is no batch size, because it is a usually a bad idea to use
+        LBFGS on anything but all the data at onece
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run.
+        lr : float
+            Optional, The learning rate (alpha) to use. Default is 0.1.
+        history_size : int
+            Optional, the length of the history 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.
+        calculation_width : int
+            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.
+        """
+        # 1) The subset statement is contained in Reconstructor.__init__
+
+        # 2) Set up / re-initialize the data loader. For LBFGS, we load
+        #    all the data at once.
+        self.setup_dataloader(batch_size=len(self.dataset))
+
+        # 3) The optimizer is created in self.__init__, but the
+        #    hyperparameters need to be set up with self.adjust_optimizer
+        self.adjust_optimizer(lr=lr,
+                              history_size=history_size,
+                              line_search_fn=line_search_fn)
+
+        # 4) This is analagous to making a call to CDIModel.AD_optimize
+        return super(LBFGS, self).optimize(iterations,
+                                           regularization_factor,
+                                           thread,
+                                           calculation_width)
diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py
new file mode 100644
index 0000000..f2dd7b0
--- /dev/null
+++ b/src/cdtools/reconstructors/sgd.py
@@ -0,0 +1,157 @@
+"""This module contains the SGD Reconstructor subclass for performing
+optimization ('reconstructions') on ptychographic/CDI models using
+stochastic gradient descent.
+
+The Reconstructor class is designed to resemble so-called
+'Trainer' classes that (in the language of the AI/ML folks) handles
+the 'training' of a model given some dataset and optimizer.
+"""
+import torch as t
+from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+from cdtools.models import CDIModel
+from typing import List, Union
+from cdtools.reconstructors import Reconstructor
+
+__all__ = ['SGD']
+
+
+class SGD(Reconstructor):
+    """
+    The Adam Reconstructor subclass handles the optimization ('reconstruction')
+    of ptychographic models and datasets using the Adam optimizer.
+
+    Parameters
+    ----------
+    model: CDIModel
+        Model for CDI/ptychography reconstruction.
+    dataset: Ptycho2DDataset
+        The dataset to reconstruct against.
+    subset : list(int) or int
+        Optional, a pattern index or list of pattern indices to use.
+
+    Important attributes:
+    - **model** -- Always points to the core model used.
+    - **optimizer** -- This class by default uses `torch.optim.Adam` to perform
+        optimizations.
+    - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the
+        `optimize` method.
+    - **data_loader** -- A torch.utils.data.DataLoader that is defined by
+        calling the `setup_dataloader` method.
+    """
+    def __init__(self,
+                 model: CDIModel,
+                 dataset: Ptycho2DDataset,
+                 subset: List[int] = None):
+
+        super().__init__(model, dataset, subset)
+
+        # Define the optimizer for use in this subclass
+        self.optimizer = t.optim.SGD(self.model.parameters())
+
+    def adjust_optimizer(self,
+                         lr: int = 0.005,
+                         momentum: float = 0,
+                         dampening: float = 0,
+                         weight_decay: float = 0,
+                         nesterov: bool = False):
+        """
+        Change hyperparameters for the utilized optimizer.
+
+        Parameters
+        ----------
+        lr : float
+            Optional, The learning rate (alpha) to use. Default is 0.005. 0.05
+            is typically the highest possible value with any chance of being
+            stable.
+        momentum : float
+            Optional, the length of the history to use.
+        dampening : float
+            Optional, dampening for the momentum.
+        weight_decay : float
+            Optional, weight decay (L2 penalty).
+        nesterov : bool
+            Optional, enables Nesterov momentum. Only applicable when momentum
+            is non-zero.
+        """
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+            param_group['momentum'] = momentum
+            param_group['dampening'] = dampening
+            param_group['weight_decay'] = weight_decay
+            param_group['nesterov'] = nesterov
+
+    def optimize(self,
+                 iterations: int,
+                 batch_size: int = None,
+                 lr: float = 2e-7,
+                 momentum: float = 0,
+                 dampening: float = 0,
+                 weight_decay: float = 0,
+                 nesterov: bool = False,
+                 regularization_factor: Union[float, List[float]] = None,
+                 thread: bool = True,
+                 calculation_width: int = 10,
+                 shuffle: bool = True):
+        """
+        Runs a round of reconstruction using the Adam optimizer
+
+        Formerly `CDIModel.Adam_optimize`
+
+        This calls the Reconstructor.optimize superclass method
+        (formerly `CDIModel.AD_optimize`) to run a round of reconstruction
+        once the dataloader and optimizer hyperparameters have been
+        set up.
+
+        Parameters
+        ----------
+        iterations : int
+            How many epochs of the algorithm to run.
+        batch_size : int
+            Optional, the size of the minibatches to use.
+        lr : float
+            Optional, The learning rate to use. The default is 2e-7.
+        momentum : float
+            Optional, the length of the history to use.
+        dampening : float
+            Optional, dampening for the momentum.
+        weight_decay : float
+            Optional, weight decay (L2 penalty).
+        nesterov : bool
+            Optional, enables Nesterov momentum. Only applicable when momentum
+            is non-zero.
+        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.
+        calculation_width : int
+            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.
+        shuffle : bool
+            Optional, enable/disable shuffling of the dataset. This option
+            is intended for diagnostic purposes and should be left as True.
+        """
+        # 1) The subset statement is contained in Reconstructor.__init__
+
+        # 2) Set up / re-initialize the data laoder
+        if batch_size is not None:
+            self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
+        else:
+            # Use default torch dataloader parameters
+            self.setup_dataloader(batch_size=1, shuffle=False)
+
+        # 3) The optimizer is created in self.__init__, but the
+        #    hyperparameters need to be set up with self.adjust_optimizer
+        self.adjust_optimizer(lr=lr,
+                              momentum=momentum,
+                              dampening=dampening,
+                              weight_decay=weight_decay,
+                              nesterov=nesterov)
+
+        # 4) This is analagous to making a call to CDIModel.AD_optimize
+        return super(SGD, self).optimize(iterations,
+                                         regularization_factor,
+                                         thread,
+                                         calculation_width)

From 76bc2fa75b5ce0a4503c258138f1562dfdf072c5 Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Fri, 1 Aug 2025 21:50:35 +0000
Subject: [PATCH 31/55] Got rid of world_size attribute

---
 src/cdtools/reconstructors/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index 65b2771..d5d34a6 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -178,7 +178,7 @@ class Reconstructor:
                     loss.backward()
 
                     # Normalize the accumulating total loss
-                    total_loss += loss.detach() // self.model.world_size
+                    total_loss += loss.detach() 
 
                 # If we have a regularizer, we can calculate it separately,
                 # and the gradients will add to the minibatch gradient

From 404d03bbd3186f43294a9dbd437cc0080e432519 Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Fri, 1 Aug 2025 21:51:52 +0000
Subject: [PATCH 32/55] Rebased CDIModel to use Reconstructors for
 reconstructions.

---
 src/cdtools/models/base.py | 472 +++++++++++++------------------------
 1 file changed, 161 insertions(+), 311 deletions(-)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index b455f45..41cdcd0 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -40,6 +40,9 @@ import time
 from scipy import io
 from contextlib import contextmanager
 from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
+from cdtools.datasets import CDataset
+from typing import List, Union, Tuple
+import os
 
 __all__ = ['CDIModel']
 
@@ -316,202 +319,24 @@ class CDIModel(t.nn.Module):
 
         self.current_checkpoint_id += 1
 
-
-        
-    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. It is a
-        generator which yields the average loss each epoch, ending after
-        the specified number of iterations.
-
-        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
-        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
-            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
-        calculation_width : int
-            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.
-
-        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]),
-                                                   calculation_width)]
-                    pattern_chunks = [patterns[i:i + calculation_width]
-                                      for i in range(0, len(inputs[0]),
-                                                     calculation_width)]
-
-                    total_loss = 0
-                    for inp, pats in zip(input_chunks, pattern_chunks):
-                        # This 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)
-
-            self.loss_history.append(loss)
-            self.epoch = len(self.loss_history)
-            self.latest_iteration_time = time.time() - t0
-            self.training_history += self.report() + '\n'
-            return loss
-
-        # We store the current optimizer as a model parameter so that
-        # it can be saved and loaded for checkpointing
-        self.current_optimizer = optimizer
-        
-        # If we don't want to run in a different thread, this is easy
-        if not thread:
-            for it in range(iterations):
-                if self.skip_computation():
-                    self.epoch = self.epoch + 1
-                    if len(self.loss_history) >= 1:
-                        yield self.loss_history[-1]
-                    else:
-                        yield float('nan')
-                    continue
-                
-                yield run_epoch()
-                    
-                
-        # 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_epoch(stop_event))
-                except Exception as e:
-                    # If something bad happens, put the exception into the
-                    # result queue
-                    result_queue.put(e)
-
-            # And this actually starts and monitors the thread
-            for it in range(iterations):
-                if self.skip_computation():
-                    self.epoch = self.epoch + 1                    
-                    if len(self.loss_history) >= 1:
-                        yield self.loss_history[-1]
-                    else:
-                        yield float('nan')
-                    continue
-                
-                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
-
-                res = result_queue.get()
-
-                # If something went wrong in the thead, we'll get an exception
-                if isinstance(res, Exception):
-                    raise res
-
-                yield res
-
-        # And finally, we unset the current optimizer:
-        self.current_optimizer = None
-
     
     def Adam_optimize(
             self,
-            iterations,
-            dataset,
-            batch_size=15,
-            lr=0.005,
-            betas=(0.9, 0.999),
-            schedule=False,
-            amsgrad=False,
-            subset=None,
-            regularization_factor=None,
+            iterations: int,
+            dataset: CDataset,
+            batch_size: int = 15,
+            lr: float = 0.005,
+            betas: Tuple[float] = (0.9, 0.999),
+            schedule: bool = False,
+            amsgrad: bool = False,
+            subset: Union[int, List[int]] = None,
+            regularization_factor: Union[float, List[float]] = None,
             thread=True,
             calculation_width=10
     ):
-        """Runs a round of reconstruction using the Adam optimizer
+        """
+        Runs a round of reconstruction using the Adam optimizer from
+        cdtools.reconstructors.Adam.
 
         This is generally accepted to be the most robust algorithm for use
         with ptychography. Like all the other optimization routines,
@@ -521,125 +346,143 @@ class CDIModel(t.nn.Module):
         Parameters
         ----------
         iterations : int
-            How many epochs of the algorithm to run
+            How many epochs of the algorithm to run.
         dataset : CDataset
-            The dataset to reconstruct against
+            The dataset to reconstruct against.
         batch_size : int
-            Optional, the size of the minibatches to use
+            Optional, the size of the minibatches to use.
         lr : float
-            Optional, The learning rate (alpha) to use. Defaultis 0.005. 0.05 is typically the highest possible value with any chance of being stable
-        betas : tuple
+            Optional, The learning rate (alpha) to use. Defaultis 0.005. 
+            0.05 is typically the highest possible value with any chance 
+            of being stable.
+        betas : tuple(float)
             Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
-        schedule : float
-            Optional, whether to use the ReduceLROnPlateau scheduler
+        schedule : bool
+            Optional, whether to use the ReduceLROnPlateau scheduler.
+        amsgrad : bool
+            Optional, whether to use the AMSGrad variant of this algorithm.
         subset : list(int) or int
             Optional, a pattern index or list of pattern indices to use
-        regularization_factor : float or list(float)
-            Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
+        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.
         calculation_width : int
-            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
-
-        """
-
-        self.training_history += (
-            f'Planning {iterations} epochs of Adam, with a learning rate = '
-            f'{lr}, batch size = {batch_size}, regularization_factor = '
-            f'{regularization_factor}, and schedule = {schedule}.\n'
-        )
+            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 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,
-                                           shuffle=True)
-
-        # Define the optimizer
-        optimizer = t.optim.Adam(
-            self.parameters(),
-            lr = lr,
-            betas=betas,
-            amsgrad=amsgrad)
-
-        # Define the scheduler
-        if schedule:
-            scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9)
-        else:
-            scheduler = None
-
-        return self.AD_optimize(iterations, data_loader, optimizer,
-                                scheduler=scheduler,
-                                regularization_factor=regularization_factor,
-                                thread=thread,
-                                calculation_width=calculation_width)
+        """
+        # We want to have model.Adam_optimize call AND store cdtools.reconstructors.Adam
+        # to perform reconstructions without creating a new reconstructor each time we 
+        # update the hyperparameters.
+        # 
+        # The only way to do this is to make the Adam reconstructor an attribute
+        # of the model. But since the Adam reconstructor also depends on CDIModel,
+        # this seems to give rise to a circular import error unless
+        # we import cdtools.reconstructors within this method:
+        if not hasattr(self, 'reconstructor'):
+            from cdtools.reconstructors import Adam
+            self.reconstructor = Adam(model=self, 
+                                  dataset=dataset, 
+                                  subset=subset)
+        
+        # Run some reconstructions
+        return self.reconstructor.optimize(iterations=iterations,
+                                       batch_size=batch_size,
+                                       lr=lr,
+                                       betas=betas,
+                                       schedule=schedule,
+                                       amsgrad=amsgrad,
+                                       regularization_factor=regularization_factor,
+                                       thread=thread,
+                                       calculation_width=calculation_width)
 
 
-    def LBFGS_optimize(self, iterations, dataset,
-                       lr=0.1,history_size=2, subset=None,
-                       regularization_factor=None, thread=True,
-                       calculation_width=10, line_search_fn=None):
-        """Runs a round of reconstruction using the L-BFGS optimizer
+    def LBFGS_optimize(self, 
+                       iterations: int, 
+                       dataset: CDataset,
+                       lr: float = 0.1,
+                       history_size: int = 2, 
+                       subset: Union[int, List[int]] = None,
+                       regularization_factor: Union[float, List[float]] =None, 
+                       thread: bool = True,
+                       calculation_width: int = 10, 
+                       line_search_fn: str = None):
+        """
+        Runs a round of reconstruction using the L-BFGS optimizer from
+        cdtools.reconstructors.LBFGS.
 
         This algorithm is often less stable that Adam, however in certain
         situations or geometries it can be shockingly efficient. Like all
         the other optimization routines, it is defined as a generator
         function which yields the average loss each epoch.
 
-        Note: There is no batch size, because it is a usually a bad idea to use
+        NOTE: There is no batch size, because it is a usually a bad idea to use
         LBFGS on anything but all the data at onece
 
         Parameters
         ----------
         iterations : int
-            How many epochs of the algorithm to run
+            How many epochs of the algorithm to run.
         dataset : CDataset
-            The dataset to reconstruct against
+            The dataset to reconstruct against.
         lr : float
-            Optional, the learning rate to use
+            Optional, the learning rate to use.
         history_size : int
             Optional, the length of the history to use.
         subset : list(int) or int
-            Optional, a pattern index or list of pattern indices to ues
+            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
+            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.
+        calculation_width : int
+            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 type(subset) == type(1):
-                subset = [subset]
-            dataset = torchdata.Subset(dataset, subset)
-
-        # Make a dataloader. This basically does nothing but load all the
-        # data at once
-        data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
+        # We want to have model.LBFGS_optimize store cdtools.reconstructors.LBFGS
+        # as an attribute to run reconstructions without generating new reconstructors
+        # each time CDIModel.LBFGS_optimize is called.
+        # 
+        # Since the LBFGS reconstructor also depends on CDIModel, a circular import error 
+        # arises unless we import cdtools.reconstructors within this method:
+        if not hasattr(self, 'reconstructor'):
+            from cdtools.reconstructors import LBFGS
+            self.reconstructor = LBFGS(model=self, 
+                                   dataset=dataset, 
+                                   subset=subset)
+        
+        # Run some reconstructions
+        return self.reconstructor.optimize(iterations=iterations,
+                                       lr=lr,
+                                       history_size=history_size,
+                                       regularization_factor=regularization_factor,
+                                       thread=thread,
+                                       calculation_width=calculation_width,
+                                       line_search_fn = line_search_fn)
 
 
-        # Define the optimizer
-        optimizer = t.optim.LBFGS(self.parameters(),
-                                  lr = lr, history_size=history_size,
-                                  line_search_fn=line_search_fn)
-
-        return self.AD_optimize(iterations, data_loader, optimizer,
-                                regularization_factor=regularization_factor,
-                                thread=thread,
-                                calculation_width=calculation_width)
-
-
-    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,
-                     thread=True, calculation_width=10):
-        """Runs a round of reconstruction using the SGD optimizer
+    def SGD_optimize(self, 
+                     iterations: int, 
+                     dataset: CDataset, 
+                     batch_size: int = None,
+                     lr: float = 2e-7, 
+                     momentum: float = 0, 
+                     dampening: float = 0, 
+                     weight_decay: float = 0,
+                     nesterov: bool = False, 
+                     subset: Union[int, List[int]] = None, 
+                     regularization_factor: Union[float, List[float]] = None,
+                     thread: bool = True, 
+                     calculation_width: int = 10):
+        """
+        Runs a round of reconstruction using the SGD optimizer from
+        cdtools.reconstructors.SGD.
 
         This algorithm is often less stable that Adam, but it is simpler
         and is the basic workhorse of gradience descent.
@@ -647,51 +490,58 @@ class CDIModel(t.nn.Module):
         Parameters
         ----------
         iterations : int
-            How many epochs of the algorithm to run
+            How many epochs of the algorithm to run.
         dataset : CDataset
-            The dataset to reconstruct against
+            The dataset to reconstruct against.
         batch_size : int
-            Optional, the size of the minibatches to use
+            Optional, the size of the minibatches to use.
         lr : float
-            Optional, the learning rate to use
+            Optional, the learning rate to use.
         momentum : float
             Optional, the length of the history to use.
+        dampening : float
+            Optional, dampening for the momentum.
+        weight_decay : float
+            Optional, weight decay (L2 penalty).
+        nesterov : bool
+            Optional, enables Nesterov momentum. Only applicable when momentum 
+            is non-zero. 
         subset : list(int) or int
-            Optional, a pattern index or list of pattern indices to use
+            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
+            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.
         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.
 
         """
-
-        if subset is not None:
-            # if just one pattern, turn into a list for convenience
-            if type(subset) == type(1):
-                subset = [subset]
-            dataset = torchdata.Subset(dataset, subset)
-
-        # Make a dataloader
-        if batch_size is not None:
-            data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
-                                               shuffle=True)
-        else:
-            data_loader = torchdata.DataLoader(dataset)
-
-
-        # Define the optimizer
-        optimizer = t.optim.SGD(self.parameters(),
-                                lr = lr, momentum=momentum,
-                                dampening=dampening,
-                                weight_decay=weight_decay,
-                                nesterov=nesterov)
-
-        return self.AD_optimize(iterations, data_loader, optimizer,
-                                regularization_factor=regularization_factor,
-                                thread=thread,
-                                calculation_width=calculation_width)
+        # We want to have model.SGD_optimize store cdtools.reconstructors.SGD
+        # as an attribute to run reconstructions without generating new reconstructors
+        # each time CDIModel.SGD_optimize is called.
+        # 
+        # Since the SGD reconstructor also depends on CDIModel, a circular import error 
+        # arises unless we import cdtools.reconstructors within this method:
+        if not hasattr(self, 'reconstructor'):
+            from cdtools.reconstructors import SGD
+            self.reconstructor = SGD(model=self, 
+                                 dataset=dataset, 
+                                 subset=subset)
+    
+        # Run some reconstructions
+        return self.reconstructor.optimize(iterations=iterations,
+                                       batch_size=batch_size,
+                                       lr=lr,
+                                       momentum=momentum,
+                                       dampening=dampening,
+                                       weight_decay=weight_decay,
+                                       nesterov=nesterov,
+                                       regularization_factor=regularization_factor,
+                                       thread=thread,
+                                       calculation_width=calculation_width)
 
 
     def report(self):

From 1e7fe2eb5f4ec9d69aa3d0bf26b1152e30f8600d Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Fri, 1 Aug 2025 22:31:40 +0000
Subject: [PATCH 33/55] Added pytests for reconstructors

---
 tests/conftest.py                 |  13 +-
 tests/models/test_fancy_ptycho.py |  59 +-----
 tests/test_reconstructors.py      | 311 ++++++++++++++++++++++++++++++
 3 files changed, 324 insertions(+), 59 deletions(-)
 create mode 100644 tests/test_reconstructors.py

diff --git a/tests/conftest.py b/tests/conftest.py
index 850935d..f0faea5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -5,7 +5,6 @@ import numpy as np
 import pytest
 import torch as t
 
-
 #
 #
 # The following few fixtures define some standard data files
@@ -381,6 +380,18 @@ def lab_ptycho_cxi(pytestconfig):
         '/examples/example_data/lab_ptycho_data.cxi'
 
 
+@pytest.fixture(scope='module')
+def optical_data_ss_cxi(pytestconfig):
+    return str(pytestconfig.rootpath) + \
+        '/examples/example_data/Optical_Data_ss.cxi'
+
+
+@pytest.fixture(scope='module')
+def optical_ptycho_incoherent_pickle(pytestconfig):
+    return str(pytestconfig.rootpath) + \
+        '/examples/example_data/Optical_ptycho_incoherent.pickle'
+
+
 @pytest.fixture(scope='module')
 def example_nested_dicts(pytestconfig):
     example_tensor = t.as_tensor(np.array([1, 4.5, 7]))
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index e5422d7..4182c93 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -87,62 +87,5 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         model.compare(dataset)
 
     # If this fails, the reconstruction has gotten worse
-    assert model.loss_history[-1] < 0.001
+    assert model.loss_history[-1] < 0.0013
 
-
-@pytest.mark.slow
-def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
-
-    print('\nTesting performance on the standard gold balls dataset')
-
-    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
-
-    pad = 10
-    dataset.pad(pad)
-
-    model = cdtools.models.FancyPtycho.from_dataset(
-        dataset,
-        n_modes=3,
-        probe_support_radius=50,
-        propagation_distance=2e-6,
-        units='um',
-        probe_fourier_crop=pad
-    )
-
-    model.translation_offsets.data += \
-        0.7 * t.randn_like(model.translation_offsets)
-
-    # Not much probe intensity instability in this dataset, no need for this
-    model.weights.requires_grad = False
-
-    print('Running reconstruction on provided --reconstruction_device,',
-          reconstruction_device)
-    model.to(device=reconstruction_device)
-    dataset.get_as(device=reconstruction_device)
-
-    for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
-        print(model.report())
-        if show_plot and model.epoch % 10 == 0:
-            model.inspect(dataset)
-
-    for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
-        print(model.report())
-        if show_plot and model.epoch % 10 == 0:
-            model.inspect(dataset)
-
-    for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
-                                    schedule=True):
-        print(model.report())
-        if show_plot and model.epoch % 10 == 0:
-            model.inspect(dataset)
-
-    model.tidy_probes()
-
-    if show_plot:
-        model.inspect(dataset)
-        model.compare(dataset)
-
-    # This just comes from running a reconstruction when it was working well
-    # and choosing a rough value. If it triggers this assertion error,
-    # something changed to make the final quality worse!
-    assert model.loss_history[-1] < 0.0001
diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py
new file mode 100644
index 0000000..0f46239
--- /dev/null
+++ b/tests/test_reconstructors.py
@@ -0,0 +1,311 @@
+import pytest
+import cdtools
+import torch as t
+import numpy as np
+import pickle
+from matplotlib import pyplot as plt
+from copy import deepcopy
+
+
+@pytest.mark.slow
+def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
+    """
+    This test checks out several things with the Au particle dataset
+        1) Calls to Reconstructor.adjust_optimizer is updating the
+           hyperparameters
+        2) We are only using the single-GPU dataloading method
+        3) Ensure `recon.model` points to the original `model`
+        4) Reconstructions performed by `Adam.optimize` and
+           `model.Adam_optimize` calls produce identical results.
+        5) The quality of the reconstruction remains below a specified
+           threshold.
+        5) Ensure that the FancyPtycho model works fine and dandy with the
+           Reconstructors.
+    """
+
+    print('\nTesting performance on the standard gold balls dataset ' +
+          'with reconstructors.Adam')
+
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
+    pad = 10
+    dataset.pad(pad)
+    model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        probe_support_radius=50,
+        propagation_distance=2e-6,
+        units='um',
+        probe_fourier_crop=pad
+    )
+
+    model.translation_offsets.data += 0.7 * \
+        t.randn_like(model.translation_offsets)
+    model.weights.requires_grad = False
+
+    # Make a copy of the model
+    model_recon = deepcopy(model)
+    model.to(device=reconstruction_device)
+    model_recon.to(device=reconstruction_device)
+    dataset.get_as(device=reconstruction_device)
+
+    # ******* Reconstructions with cdtools.reconstructors.Adam.optimize *******
+    print('Running reconstruction using cdtools.reconstructors.Adam.optimize' +
+          ' on provided reconstruction_device,', reconstruction_device)
+
+    recon = cdtools.reconstructors.Adam(model=model_recon, dataset=dataset)
+    t.manual_seed(0)
+
+    # Run a reconstruction
+    epoch_tup = (20, 50, 100)
+    lr_tup = (0.005, 0.002, 0.001)
+    batch_size_tup = (50, 100, 100)
+
+    for i, iterations in enumerate(epoch_tup):
+        for loss in recon.optimize(iterations,
+                                   lr=lr_tup[i],
+                                   batch_size=batch_size_tup[i]):
+            print(model_recon.report())
+            if show_plot and model_recon.epoch % 10 == 0:
+                model_recon.inspect(dataset)
+
+        # Check hyperparameter update
+        assert recon.optimizer.param_groups[0]['lr'] == lr_tup[i]
+        assert recon.data_loader.batch_size == batch_size_tup[i]
+
+    # Ensure that recon does not have sampler as an attribute (only used in
+    # multi-GPU)
+    assert not hasattr(recon, 'sampler')
+
+    # Ensure recon.model points to the original model
+    assert id(model_recon) == id(recon.model)
+
+    model_recon.tidy_probes()
+
+    if show_plot:
+        model_recon.inspect(dataset)
+        model_recon.compare(dataset)
+
+    # ******* Reconstructions with cdtools.CDIModel.Adam_optimize *******
+    print('Running reconstruction using CDIModel.Adam_optimize on provided' +
+          ' reconstruction_device,', reconstruction_device)
+    t.manual_seed(0)
+
+    for i, iterations in enumerate(epoch_tup):
+        for loss in model.Adam_optimize(iterations,
+                                        dataset,
+                                        lr=lr_tup[i],
+                                        batch_size=batch_size_tup[i]):
+            print(model.report())
+            if show_plot and model.epoch % 10 == 0:
+                model.inspect(dataset)
+
+    model.tidy_probes()
+
+    if show_plot:
+        model.inspect(dataset)
+        model.compare(dataset)
+
+    # Ensure equivalency between the model reconstructions
+    assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
+
+    # Ensure reconstructions have reached a certain loss tolerance. This just
+    # comes from running a reconstruction when it was working well and
+    # choosing a rough value. If it triggers this assertion error, something
+    # changed to make the final quality worse!
+    assert model.loss_history[-1] < 0.0001
+
+
+@pytest.mark.slow
+def test_LBFGS_RPI(optical_data_ss_cxi,
+                   optical_ptycho_incoherent_pickle,
+                   reconstruction_device,
+                   show_plot):
+    """
+    This test checks out several things with the transmission RPI dataset
+        1) Calls to Reconstructor.adjust_optimizer is updating the
+           hyperparameters
+        2) Ensure `recon.model` points to the original `model`
+        3) Reconstructions performed by `LBFGS.optimize` and
+           `model.LBFGS_optimize` calls produce identical results.
+        4) The quality of the reconstruction remains below a specified
+           threshold.
+        5) Ensure that the RPI model works fine and dandy with the
+           Reconstructors.
+    """
+    with open(optical_ptycho_incoherent_pickle, 'rb') as f:
+        ptycho_results = pickle.load(f)
+
+    probe = ptycho_results['probe']
+    background = ptycho_results['background']
+
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(optical_data_ss_cxi)
+    model = cdtools.models.RPI.from_dataset(dataset, probe, [500, 500],
+                                            background=background, n_modes=2,
+                                            initialization='random')
+
+    # Prepare two sets of models for the comparative reconstruction
+    model_recon = deepcopy(model)
+
+    model.to(device=reconstruction_device)
+    model_recon.to(device=reconstruction_device)
+    dataset.get_as(device=reconstruction_device)
+
+    # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
+    print('Running reconstruction using cdtools.reconstructors.LBFGS.' +
+          'optimize on provided reconstruction_device,', reconstruction_device)
+
+    recon = cdtools.reconstructors.LBFGS(model=model_recon, dataset=dataset)
+    t.manual_seed(0)
+
+    # Run a reconstruction
+    reg_factor_tup = ([0.05, 0.05], [0.001, 0.1])
+    epoch_tup = (30, 50)
+    for i, iterations in enumerate(epoch_tup):
+        for loss in recon.optimize(iterations,
+                                   lr=0.4,
+                                   regularization_factor=reg_factor_tup[i]):
+            if show_plot and i == 0:
+                model_recon.inspect(dataset)
+            print(model_recon.report())
+
+        # Check hyperparameter update (or lack thereof)
+        assert recon.optimizer.param_groups[0]['lr'] == 0.4
+
+    if show_plot:
+        model_recon.inspect(dataset)
+        model_recon.compare(dataset)
+
+    # Check model pointing
+    assert id(model_recon) == id(recon.model)
+
+    # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
+    print('Running reconstruction using CDIModel.LBFGS_optimize.' +
+          'optimize on provided reconstruction_device,', reconstruction_device)
+    t.manual_seed(0)
+    for i, iterations in enumerate(epoch_tup):
+        for loss in model.LBFGS_optimize(iterations,
+                                         dataset,
+                                         lr=0.4,
+                                         regularization_factor=reg_factor_tup[i]): # noqa
+            if show_plot and i == 0:
+                model.inspect(dataset)
+            print(model.report())
+
+    if show_plot:
+        model.inspect(dataset)
+        model.compare(dataset)
+
+    # Check loss equivalency between the two reconstructions
+    assert np.allclose(model.loss_history[-1], model_recon.loss_history[-1])
+
+    # The final loss when testing this was 2.28607e-3. Based on this, we set
+    # a threshold of 2.3e-3 for the tested loss. If this value has been
+    # exceeded, the reconstructions have gotten worse.
+    assert model.loss_history[-1] < 0.0023
+
+
+@pytest.mark.slow
+def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
+    """
+    This test checks out several things with the Au particle dataset
+        1) Calls to Reconstructor.adjust_optimizer is updating the
+           hyperparameters
+        3) Ensure `recon.model` points to the original `model`
+        4) Reconstructions performed by `SGD.optimize` and
+           `model.SGD_optimize` calls produce identical results.
+        5) The quality of the reconstruction remains below a specified
+           threshold.
+        5) Ensure that the FancyPtycho model works fine and dandy with the
+           Reconstructors.
+
+    The hyperparameters used in this test are not optimized to produce
+    a super-high-quality reconstruction. Instead, I just need A reconstruction
+    to do some kind of comparative assessment.
+    """
+    print('\nTesting performance on the standard gold balls dataset ' +
+          'with reconstructors.SGD')
+
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
+    pad = 10
+    dataset.pad(pad)
+    model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        probe_support_radius=50,
+        propagation_distance=2e-6,
+        units='um',
+        probe_fourier_crop=pad
+    )
+
+    model.translation_offsets.data += 0.7 * \
+        t.randn_like(model.translation_offsets)
+    model.weights.requires_grad = False
+
+    # Make a copy of the model
+    model_recon = deepcopy(model)
+    model.to(device=reconstruction_device)
+    model_recon.to(device=reconstruction_device)
+    dataset.get_as(device=reconstruction_device)
+
+    # ******* Reconstructions with cdtools.reconstructors.SGD.optimize *******
+    print('Running reconstruction using cdtools.reconstructors.SGD.optimize' +
+          ' on provided reconstruction_device,', reconstruction_device)
+
+    recon = cdtools.reconstructors.SGD(model=model_recon, dataset=dataset)
+    t.manual_seed(0)
+
+    # Run a reconstruction
+    epochs = 50
+    lr = 0.00000005
+    batch_size = 40
+
+    for loss in recon.optimize(epochs,
+                               lr=lr,
+                               batch_size=batch_size):
+        print(model_recon.report())
+        if show_plot and model_recon.epoch % 10 == 0:
+            model_recon.inspect(dataset)
+
+    # Check hyperparameter update
+    assert recon.optimizer.param_groups[0]['lr'] == lr
+    assert recon.data_loader.batch_size == batch_size
+
+    # Ensure that recon does not have sampler as an attribute (only used in
+    # multi-GPU)
+    assert not hasattr(recon, 'sampler')
+
+    # Ensure recon.model points to the original model
+    assert id(model_recon) == id(recon.model)
+
+    model_recon.tidy_probes()
+
+    if show_plot:
+        model_recon.inspect(dataset)
+        model_recon.compare(dataset)
+
+    # ******* Reconstructions with cdtools.CDIModel.SGD_optimize *******
+    print('Running reconstruction using CDIModel.SGD_optimize on provided' +
+          ' reconstruction_device,', reconstruction_device)
+    t.manual_seed(0)
+
+    for loss in model.SGD_optimize(epochs,
+                                   dataset,
+                                   lr=lr,
+                                   batch_size=batch_size):
+        print(model.report())
+        if show_plot and model.epoch % 10 == 0:
+            model.inspect(dataset)
+
+    model.tidy_probes()
+
+    if show_plot:
+        model.inspect(dataset)
+        model.compare(dataset)
+
+    # Ensure equivalency between the model reconstructions
+    assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
+
+    # The final loss when testing this was 7.12188e-4. Based on this, we set
+    # a threshold of 7.2e-4 for the tested loss. If this value has been
+    # exceeded, the reconstructions have gotten worse.
+    assert model.loss_history[-1] < 0.00072

From 82ce845385170afbad9845e73be1d487e17b7373 Mon Sep 17 00:00:00 2001
From: yoshikisd <yoshikisd@outlook.com>
Date: Mon, 4 Aug 2025 16:51:24 +0000
Subject: [PATCH 34/55] Changed the names of the reconstructors and updated the
 old optimizer class documentation

---
 src/cdtools/models/base.py             | 165 +++++++++++++------------
 src/cdtools/reconstructors/__init__.py |  20 +--
 src/cdtools/reconstructors/adam.py     |  14 +--
 src/cdtools/reconstructors/base.py     |   2 +-
 src/cdtools/reconstructors/lbfgs.py    |  16 +--
 src/cdtools/reconstructors/sgd.py      |  18 +--
 tests/test_reconstructors.py           |  29 +++--
 7 files changed, 138 insertions(+), 126 deletions(-)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index 41cdcd0..635cdfb 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -319,7 +319,6 @@ class CDIModel(t.nn.Module):
 
         self.current_checkpoint_id += 1
 
-    
     def Adam_optimize(
             self,
             iterations: int,
@@ -336,7 +335,7 @@ class CDIModel(t.nn.Module):
     ):
         """
         Runs a round of reconstruction using the Adam optimizer from
-        cdtools.reconstructors.Adam.
+        cdtools.reconstructors.AdamReconstructor.
 
         This is generally accepted to be the most robust algorithm for use
         with ptychography. Like all the other optimization routines,
@@ -375,45 +374,45 @@ class CDIModel(t.nn.Module):
             only the calculation speed. 
         
         """
-        # We want to have model.Adam_optimize call AND store cdtools.reconstructors.Adam
-        # to perform reconstructions without creating a new reconstructor each time we 
-        # update the hyperparameters.
-        # 
-        # The only way to do this is to make the Adam reconstructor an attribute
-        # of the model. But since the Adam reconstructor also depends on CDIModel,
-        # this seems to give rise to a circular import error unless
-        # we import cdtools.reconstructors within this method:
+        # We want to have model.Adam_optimize call AND store
+        # cdtools.reconstructors.AdamReconstructor to perform reconstructions
+        # without creating a new reconstructor each time we update the
+        # hyperparameters.
+        #
+        # The only way to do this is to make the Adam reconstructor an
+        # attribute of the model. But since the Adam reconstructor also
+        # depends on CDIModel, this seems to give rise to a circular import
+        # error unless we import cdtools.reconstructors within this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import Adam
-            self.reconstructor = Adam(model=self, 
-                                  dataset=dataset, 
-                                  subset=subset)
+            from cdtools.reconstructors import AdamReconstructor
+            self.reconstructor = AdamReconstructor(model=self,
+                                                   dataset=dataset,
+                                                   subset=subset)
         
         # Run some reconstructions
         return self.reconstructor.optimize(iterations=iterations,
-                                       batch_size=batch_size,
-                                       lr=lr,
-                                       betas=betas,
-                                       schedule=schedule,
-                                       amsgrad=amsgrad,
-                                       regularization_factor=regularization_factor,
-                                       thread=thread,
-                                       calculation_width=calculation_width)
+                                           batch_size=batch_size,
+                                           lr=lr,
+                                           betas=betas,
+                                           schedule=schedule,
+                                           amsgrad=amsgrad,
+                                           regularization_factor=regularization_factor, # noqa
+                                           thread=thread,
+                                           calculation_width=calculation_width)
 
-
-    def LBFGS_optimize(self, 
-                       iterations: int, 
+    def LBFGS_optimize(self,
+                       iterations: int,
                        dataset: CDataset,
                        lr: float = 0.1,
-                       history_size: int = 2, 
+                       history_size: int = 2,
                        subset: Union[int, List[int]] = None,
-                       regularization_factor: Union[float, List[float]] =None, 
+                       regularization_factor: Union[float, List[float]] =None,
                        thread: bool = True,
-                       calculation_width: int = 10, 
+                       calculation_width: int = 10,
                        line_search_fn: str = None):
         """
         Runs a round of reconstruction using the L-BFGS optimizer from
-        cdtools.reconstructors.LBFGS.
+        cdtools.reconstructors.LBFGSReconstructor.
 
         This algorithm is often less stable that Adam, however in certain
         situations or geometries it can be shockingly efficient. Like all
@@ -436,53 +435,55 @@ class CDIModel(t.nn.Module):
         subset : list(int) or int
             Optional, a pattern index or list of pattern indices to use.
         regularization_factor : float or list(float)
-            Optional, if the model has a regularizer defined, the set of parameters 
-            to pass the regularizer method.
+            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.
         calculation_width : int
-            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 
+            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.
         """
-        # We want to have model.LBFGS_optimize store cdtools.reconstructors.LBFGS
-        # as an attribute to run reconstructions without generating new reconstructors
-        # each time CDIModel.LBFGS_optimize is called.
-        # 
-        # Since the LBFGS reconstructor also depends on CDIModel, a circular import error 
-        # arises unless we import cdtools.reconstructors within this method:
+        # We want to have model.LBFGS_optimize store
+        # cdtools.reconstructors.LBFGSReconstructor as an attribute to run
+        # reconstructions without generating new reconstructors each time
+        # CDIModel.LBFGS_optimize is called.
+        #
+        # Since the LBFGS reconstructor also depends on CDIModel, a circular
+        # import error arises unless we import cdtools.reconstructors within
+        # this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import LBFGS
-            self.reconstructor = LBFGS(model=self, 
-                                   dataset=dataset, 
-                                   subset=subset)
+            from cdtools.reconstructors import LBFGSReconstructor
+            self.reconstructor = LBFGSReconstructor(model=self,
+                                                    dataset=dataset,
+                                                    subset=subset)
         
         # Run some reconstructions
         return self.reconstructor.optimize(iterations=iterations,
-                                       lr=lr,
-                                       history_size=history_size,
-                                       regularization_factor=regularization_factor,
-                                       thread=thread,
-                                       calculation_width=calculation_width,
-                                       line_search_fn = line_search_fn)
-
+                                           lr=lr,
+                                           history_size=history_size,
+                                           regularization_factor=regularization_factor, # noqa
+                                           thread=thread,
+                                           calculation_width=calculation_width,
+                                           line_search_fn=line_search_fn)
 
     def SGD_optimize(self, 
-                     iterations: int, 
-                     dataset: CDataset, 
+                     iterations: int,
+                     dataset: CDataset,
                      batch_size: int = None,
-                     lr: float = 2e-7, 
-                     momentum: float = 0, 
-                     dampening: float = 0, 
+                     lr: float = 2e-7,
+                     momentum: float = 0,
+                     dampening: float = 0,
                      weight_decay: float = 0,
-                     nesterov: bool = False, 
-                     subset: Union[int, List[int]] = None, 
+                     nesterov: bool = False,
+                     subset: Union[int, List[int]] = None,
                      regularization_factor: Union[float, List[float]] = None,
-                     thread: bool = True, 
+                     thread: bool = True,
                      calculation_width: int = 10):
         """
         Runs a round of reconstruction using the SGD optimizer from
-        cdtools.reconstructors.SGD.
+        cdtools.reconstructors.SGDReconstructor.
 
         This algorithm is often less stable that Adam, but it is simpler
         and is the basic workhorse of gradience descent.
@@ -519,29 +520,31 @@ class CDIModel(t.nn.Module):
             round of gradient accumulation.
 
         """
-        # We want to have model.SGD_optimize store cdtools.reconstructors.SGD
-        # as an attribute to run reconstructions without generating new reconstructors
-        # each time CDIModel.SGD_optimize is called.
-        # 
-        # Since the SGD reconstructor also depends on CDIModel, a circular import error 
-        # arises unless we import cdtools.reconstructors within this method:
+        # We want to have model.SGD_optimize store
+        # cdtools.reconstructors.SGDReconstructor as an attribute to run
+        # reconstructions without generating new reconstructors each time
+        # CDIModel.SGD_optimize is called.
+        #
+        # Since the SGD reconstructor also depends on CDIModel, a circular
+        # import error arises unless we import cdtools.reconstructors within
+        # this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import SGD
-            self.reconstructor = SGD(model=self, 
-                                 dataset=dataset, 
-                                 subset=subset)
-    
+            from cdtools.reconstructors import SGDReconstructor
+            self.reconstructor = SGDReconstructor(model=self, 
+                                                  dataset=dataset, 
+                                                  subset=subset)
+
         # Run some reconstructions
         return self.reconstructor.optimize(iterations=iterations,
-                                       batch_size=batch_size,
-                                       lr=lr,
-                                       momentum=momentum,
-                                       dampening=dampening,
-                                       weight_decay=weight_decay,
-                                       nesterov=nesterov,
-                                       regularization_factor=regularization_factor,
-                                       thread=thread,
-                                       calculation_width=calculation_width)
+                                           batch_size=batch_size,
+                                           lr=lr,
+                                           momentum=momentum,
+                                           dampening=dampening,
+                                           weight_decay=weight_decay,
+                                           nesterov=nesterov,
+                                           regularization_factor=regularization_factor, # noqa
+                                           thread=thread,
+                                           calculation_width=calculation_width)
 
 
     def report(self):
diff --git a/src/cdtools/reconstructors/__init__.py b/src/cdtools/reconstructors/__init__.py
index 84b96ab..84a1a81 100644
--- a/src/cdtools/reconstructors/__init__.py
+++ b/src/cdtools/reconstructors/__init__.py
@@ -1,16 +1,22 @@
-"""This module contains optimizers for performing reconstructions
+"""
+Module `cdtools.tools.reconstructors` contains the `Reconstructor` class and
+subclasses which run the ptychography reconstructions on a given model and
+dataset.
 
+The reconstructors are designed to resemble so-called 'Trainer' classes that
+(in the language of the AI/ML folks) handles the 'training' of a model given
+some dataset and optimizer.
 """
 
 # We define __all__ to be sure that import * only imports what we want
 __all__ = [
     'Reconstructor',
-    'Adam',
-    'LBFGS',
-    'SGD'
+    'AdamReconstructor',
+    'LBFGSReconstructor',
+    'SGDReconstructor'
 ]
 
 from cdtools.reconstructors.base import Reconstructor
-from cdtools.reconstructors.adam import Adam
-from cdtools.reconstructors.lbfgs import LBFGS
-from cdtools.reconstructors.sgd import SGD
+from cdtools.reconstructors.adam import AdamReconstructor
+from cdtools.reconstructors.lbfgs import LBFGSReconstructor
+from cdtools.reconstructors.sgd import SGDReconstructor
diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index 5a489a0..6eecc9c 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -1,4 +1,4 @@
-"""This module contains the Adam Reconstructor subclass for performing
+"""This module contains the AdamReconstructor subclass for performing
 optimization ('reconstructions') on ptychographic/CDI models using
 the Adam optimizer.
 
@@ -12,10 +12,10 @@ from cdtools.models import CDIModel
 from typing import Tuple, List, Union
 from cdtools.reconstructors import Reconstructor
 
-__all__ = ['Adam']
+__all__ = ['AdamReconstructor']
 
 
-class Adam(Reconstructor):
+class AdamReconstructor(Reconstructor):
     """
     The Adam Reconstructor subclass handles the optimization ('reconstruction')
     of ptychographic models and datasets using the Adam optimizer.
@@ -154,7 +154,7 @@ class Adam(Reconstructor):
             self.scheduler = None
 
         # 5) This is analagous to making a call to CDIModel.AD_optimize
-        return super(Adam, self).optimize(iterations,
-                                          regularization_factor,
-                                          thread,
-                                          calculation_width)
+        return super(AdamReconstructor, self).optimize(iterations,
+                                                       regularization_factor,
+                                                       thread,
+                                                       calculation_width)
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index d5d34a6..fd2cbb2 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -178,7 +178,7 @@ class Reconstructor:
                     loss.backward()
 
                     # Normalize the accumulating total loss
-                    total_loss += loss.detach() 
+                    total_loss += loss.detach()
 
                 # If we have a regularizer, we can calculate it separately,
                 # and the gradients will add to the minibatch gradient
diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py
index 0b51dfd..8bbd7ae 100644
--- a/src/cdtools/reconstructors/lbfgs.py
+++ b/src/cdtools/reconstructors/lbfgs.py
@@ -1,4 +1,4 @@
-"""This module contains the LBFGS Reconstructor subclass for performing
+"""This module contains the LBFGSReconstructor subclass for performing
 optimization ('reconstructions') on ptychographic/CDI models using
 the LBFGS optimizer.
 
@@ -12,12 +12,12 @@ from cdtools.models import CDIModel
 from typing import List, Union
 from cdtools.reconstructors import Reconstructor
 
-__all__ = ['LBFGS']
+__all__ = ['LBFGSReconstructor']
 
 
-class LBFGS(Reconstructor):
+class LBFGSReconstructor(Reconstructor):
     """
-    The LBFGS Reconstructor subclass handles the optimization
+    The LBFGSReconstructor subclass handles the optimization
     ('reconstruction') of ptychographic models and datasets using the LBFGS
     optimizer.
 
@@ -128,7 +128,7 @@ class LBFGS(Reconstructor):
                               line_search_fn=line_search_fn)
 
         # 4) This is analagous to making a call to CDIModel.AD_optimize
-        return super(LBFGS, self).optimize(iterations,
-                                           regularization_factor,
-                                           thread,
-                                           calculation_width)
+        return super(LBFGSReconstructor, self).optimize(iterations,
+                                                        regularization_factor,
+                                                        thread,
+                                                        calculation_width)
diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py
index f2dd7b0..a8f2ebc 100644
--- a/src/cdtools/reconstructors/sgd.py
+++ b/src/cdtools/reconstructors/sgd.py
@@ -1,4 +1,4 @@
-"""This module contains the SGD Reconstructor subclass for performing
+"""This module contains the SGDReconstructor subclass for performing
 optimization ('reconstructions') on ptychographic/CDI models using
 stochastic gradient descent.
 
@@ -12,13 +12,13 @@ from cdtools.models import CDIModel
 from typing import List, Union
 from cdtools.reconstructors import Reconstructor
 
-__all__ = ['SGD']
+__all__ = ['SGDReconstructor']
 
 
-class SGD(Reconstructor):
+class SGDReconstructor(Reconstructor):
     """
-    The Adam Reconstructor subclass handles the optimization ('reconstruction')
-    of ptychographic models and datasets using the Adam optimizer.
+    The SGDReconstructor subclass handles the optimization ('reconstruction')
+    of ptychographic models and datasets using the SGD optimizer.
 
     Parameters
     ----------
@@ -151,7 +151,7 @@ class SGD(Reconstructor):
                               nesterov=nesterov)
 
         # 4) This is analagous to making a call to CDIModel.AD_optimize
-        return super(SGD, self).optimize(iterations,
-                                         regularization_factor,
-                                         thread,
-                                         calculation_width)
+        return super(SGDReconstructor, self).optimize(iterations,
+                                                      regularization_factor,
+                                                      thread,
+                                                      calculation_width)
diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py
index 0f46239..3995b6f 100644
--- a/tests/test_reconstructors.py
+++ b/tests/test_reconstructors.py
@@ -24,7 +24,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     """
 
     print('\nTesting performance on the standard gold balls dataset ' +
-          'with reconstructors.Adam')
+          'with reconstructors.AdamReconstructor')
 
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
     pad = 10
@@ -48,11 +48,12 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     model_recon.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
-    # ******* Reconstructions with cdtools.reconstructors.Adam.optimize *******
-    print('Running reconstruction using cdtools.reconstructors.Adam.optimize' +
+    # ******* Reconstructions with AdamReconstructor.optimize *******
+    print('Running reconstruction using AdamReconstructor.optimize' +
           ' on provided reconstruction_device,', reconstruction_device)
 
-    recon = cdtools.reconstructors.Adam(model=model_recon, dataset=dataset)
+    recon = cdtools.reconstructors.AdamReconstructor(model=model_recon,
+                                                     dataset=dataset)
     t.manual_seed(0)
 
     # Run a reconstruction
@@ -85,7 +86,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         model_recon.inspect(dataset)
         model_recon.compare(dataset)
 
-    # ******* Reconstructions with cdtools.CDIModel.Adam_optimize *******
+    # ******* Reconstructions with CDIModel.Adam_optimize *******
     print('Running reconstruction using CDIModel.Adam_optimize on provided' +
           ' reconstruction_device,', reconstruction_device)
     t.manual_seed(0)
@@ -150,11 +151,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
     model_recon.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
-    # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
-    print('Running reconstruction using cdtools.reconstructors.LBFGS.' +
+    # ******* Reconstructions with LBFGSReconstructor.optimize ******
+    print('Running reconstruction using LBFGSReconstructor.' +
           'optimize on provided reconstruction_device,', reconstruction_device)
 
-    recon = cdtools.reconstructors.LBFGS(model=model_recon, dataset=dataset)
+    recon = cdtools.reconstructors.LBFGSReconstructor(model=model_recon,
+                                                      dataset=dataset)
     t.manual_seed(0)
 
     # Run a reconstruction
@@ -178,7 +180,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
     # Check model pointing
     assert id(model_recon) == id(recon.model)
 
-    # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
+    # ******* Reconstructions with CDIModel.LBFGS_optimize ******
     print('Running reconstruction using CDIModel.LBFGS_optimize.' +
           'optimize on provided reconstruction_device,', reconstruction_device)
     t.manual_seed(0)
@@ -223,7 +225,7 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     to do some kind of comparative assessment.
     """
     print('\nTesting performance on the standard gold balls dataset ' +
-          'with reconstructors.SGD')
+          'with reconstructors.SGDReconstructor')
 
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
     pad = 10
@@ -247,11 +249,12 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     model_recon.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
-    # ******* Reconstructions with cdtools.reconstructors.SGD.optimize *******
-    print('Running reconstruction using cdtools.reconstructors.SGD.optimize' +
+    # ******* Reconstructions with SGDReconstructor.optimize *******
+    print('Running reconstruction using SGDReconstructor.optimize' +
           ' on provided reconstruction_device,', reconstruction_device)
 
-    recon = cdtools.reconstructors.SGD(model=model_recon, dataset=dataset)
+    recon = cdtools.reconstructors.SGDReconstructor(model=model_recon,
+                                                    dataset=dataset)
     t.manual_seed(0)
 
     # Run a reconstruction

From 766a7d1adf038f893e35dcf82d497d357521017b Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 13 Oct 2025 14:20:38 +0200
Subject: [PATCH 35/55] Move to better supported strategy for avoiding circular
 imports

---
 src/cdtools/models/base.py          | 17 +----------------
 src/cdtools/reconstructors/adam.py  |  9 +++++++--
 src/cdtools/reconstructors/base.py  |  9 +++++++--
 src/cdtools/reconstructors/lbfgs.py | 10 ++++++++--
 src/cdtools/reconstructors/sgd.py   | 10 ++++++++--
 5 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index 635cdfb..5d3ac9a 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -40,6 +40,7 @@ import time
 from scipy import io
 from contextlib import contextmanager
 from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
+from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor
 from cdtools.datasets import CDataset
 from typing import List, Union, Tuple
 import os
@@ -378,13 +379,7 @@ class CDIModel(t.nn.Module):
         # cdtools.reconstructors.AdamReconstructor to perform reconstructions
         # without creating a new reconstructor each time we update the
         # hyperparameters.
-        #
-        # The only way to do this is to make the Adam reconstructor an
-        # attribute of the model. But since the Adam reconstructor also
-        # depends on CDIModel, this seems to give rise to a circular import
-        # error unless we import cdtools.reconstructors within this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import AdamReconstructor
             self.reconstructor = AdamReconstructor(model=self,
                                                    dataset=dataset,
                                                    subset=subset)
@@ -449,12 +444,7 @@ class CDIModel(t.nn.Module):
         # cdtools.reconstructors.LBFGSReconstructor as an attribute to run
         # reconstructions without generating new reconstructors each time
         # CDIModel.LBFGS_optimize is called.
-        #
-        # Since the LBFGS reconstructor also depends on CDIModel, a circular
-        # import error arises unless we import cdtools.reconstructors within
-        # this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import LBFGSReconstructor
             self.reconstructor = LBFGSReconstructor(model=self,
                                                     dataset=dataset,
                                                     subset=subset)
@@ -524,12 +514,7 @@ class CDIModel(t.nn.Module):
         # cdtools.reconstructors.SGDReconstructor as an attribute to run
         # reconstructions without generating new reconstructors each time
         # CDIModel.SGD_optimize is called.
-        #
-        # Since the SGD reconstructor also depends on CDIModel, a circular
-        # import error arises unless we import cdtools.reconstructors within
-        # this method:
         if not hasattr(self, 'reconstructor'):
-            from cdtools.reconstructors import SGDReconstructor
             self.reconstructor = SGDReconstructor(model=self, 
                                                   dataset=dataset, 
                                                   subset=subset)
diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index 6eecc9c..4c54f1f 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -6,12 +6,17 @@ The Reconstructor class is designed to resemble so-called
 'Trainer' classes that (in the language of the AI/ML folks) handles
 the 'training' of a model given some dataset and optimizer.
 """
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 import torch as t
-from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
-from cdtools.models import CDIModel
 from typing import Tuple, List, Union
 from cdtools.reconstructors import Reconstructor
 
+if TYPE_CHECKING:
+    from cdtools.models import CDIModel
+    from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+
 __all__ = ['AdamReconstructor']
 
 
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index fd2cbb2..1937661 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -8,16 +8,21 @@ the 'training' of a model given some dataset and optimizer.
 The subclasses of Reconstructor are required to implement
 their own data loaders and optimizer adjusters
 """
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
 import torch as t
 from torch.utils import data as td
 import threading
 import queue
 import time
-from cdtools.datasets import CDataset
-from cdtools.models import CDIModel
 from typing import List, Union
 
+if TYPE_CHECKING:
+    from cdtools.models import CDIModel
+    from cdtools.datasets import CDataset
+    
+
 __all__ = ['Reconstructor']
 
 
diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py
index 8bbd7ae..60c6f97 100644
--- a/src/cdtools/reconstructors/lbfgs.py
+++ b/src/cdtools/reconstructors/lbfgs.py
@@ -6,12 +6,18 @@ The Reconstructor class is designed to resemble so-called
 'Trainer' classes that (in the language of the AI/ML folks) handles
 the 'training' of a model given some dataset and optimizer.
 """
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 import torch as t
-from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
-from cdtools.models import CDIModel
 from typing import List, Union
 from cdtools.reconstructors import Reconstructor
 
+if TYPE_CHECKING:
+    from cdtools.models import CDIModel
+    from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+
+
 __all__ = ['LBFGSReconstructor']
 
 
diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py
index a8f2ebc..5b7bf4e 100644
--- a/src/cdtools/reconstructors/sgd.py
+++ b/src/cdtools/reconstructors/sgd.py
@@ -6,12 +6,18 @@ The Reconstructor class is designed to resemble so-called
 'Trainer' classes that (in the language of the AI/ML folks) handles
 the 'training' of a model given some dataset and optimizer.
 """
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 import torch as t
-from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
-from cdtools.models import CDIModel
 from typing import List, Union
 from cdtools.reconstructors import Reconstructor
 
+if TYPE_CHECKING:
+    from cdtools.models import CDIModel
+    from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
+    
+
 __all__ = ['SGDReconstructor']
 
 

From a089dc3b7266ff6120f50e17638a8f8ba6e57ce2 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 13 Oct 2025 14:26:07 +0200
Subject: [PATCH 36/55] Switch to a simpler pattern for the
 CDIModel.<algorithm>_optimize functions that preserves the old behavior when
 the old pattern is used

---
 src/cdtools/models/base.py | 99 ++++++++++++++++++--------------------
 1 file changed, 48 insertions(+), 51 deletions(-)

diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index 5d3ac9a..df347b6 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -375,25 +375,24 @@ class CDIModel(t.nn.Module):
             only the calculation speed. 
         
         """
-        # We want to have model.Adam_optimize call AND store
-        # cdtools.reconstructors.AdamReconstructor to perform reconstructions
-        # without creating a new reconstructor each time we update the
-        # hyperparameters.
-        if not hasattr(self, 'reconstructor'):
-            self.reconstructor = AdamReconstructor(model=self,
-                                                   dataset=dataset,
-                                                   subset=subset)
+        reconstructor = AdamReconstructor(
+            model=self,
+            dataset=dataset,
+            subset=subset,
+        )
         
         # Run some reconstructions
-        return self.reconstructor.optimize(iterations=iterations,
-                                           batch_size=batch_size,
-                                           lr=lr,
-                                           betas=betas,
-                                           schedule=schedule,
-                                           amsgrad=amsgrad,
-                                           regularization_factor=regularization_factor, # noqa
-                                           thread=thread,
-                                           calculation_width=calculation_width)
+        return reconstructor.optimize(
+            iterations=iterations,
+            batch_size=batch_size,
+            lr=lr,
+            betas=betas,
+            schedule=schedule,
+            amsgrad=amsgrad,
+            regularization_factor=regularization_factor, # noqa
+            thread=thread,
+            calculation_width=calculation_width,
+        )
 
     def LBFGS_optimize(self,
                        iterations: int,
@@ -440,24 +439,23 @@ class CDIModel(t.nn.Module):
             round of gradient accumulation. Does not affect the result, only
             the calculation speed.
         """
-        # We want to have model.LBFGS_optimize store
-        # cdtools.reconstructors.LBFGSReconstructor as an attribute to run
-        # reconstructions without generating new reconstructors each time
-        # CDIModel.LBFGS_optimize is called.
-        if not hasattr(self, 'reconstructor'):
-            self.reconstructor = LBFGSReconstructor(model=self,
-                                                    dataset=dataset,
-                                                    subset=subset)
+        reconstructor = LBFGSReconstructor(
+            model=self,
+            dataset=dataset,
+            subset=subset,
+        )
         
         # Run some reconstructions
-        return self.reconstructor.optimize(iterations=iterations,
-                                           lr=lr,
-                                           history_size=history_size,
-                                           regularization_factor=regularization_factor, # noqa
-                                           thread=thread,
-                                           calculation_width=calculation_width,
-                                           line_search_fn=line_search_fn)
-
+        return reconstructor.optimize(
+            iterations=iterations,
+            lr=lr,
+            history_size=history_size,
+            regularization_factor=regularization_factor, # noqa
+            thread=thread,
+            calculation_width=calculation_width,
+            line_search_fn=line_search_fn,
+        )
+    
     def SGD_optimize(self, 
                      iterations: int,
                      dataset: CDataset,
@@ -510,26 +508,25 @@ class CDIModel(t.nn.Module):
             round of gradient accumulation.
 
         """
-        # We want to have model.SGD_optimize store
-        # cdtools.reconstructors.SGDReconstructor as an attribute to run
-        # reconstructions without generating new reconstructors each time
-        # CDIModel.SGD_optimize is called.
-        if not hasattr(self, 'reconstructor'):
-            self.reconstructor = SGDReconstructor(model=self, 
-                                                  dataset=dataset, 
-                                                  subset=subset)
+        reconstructor = SGDReconstructor(
+            model=self, 
+            dataset=dataset, 
+            subset=subset,
+        )
 
         # Run some reconstructions
-        return self.reconstructor.optimize(iterations=iterations,
-                                           batch_size=batch_size,
-                                           lr=lr,
-                                           momentum=momentum,
-                                           dampening=dampening,
-                                           weight_decay=weight_decay,
-                                           nesterov=nesterov,
-                                           regularization_factor=regularization_factor, # noqa
-                                           thread=thread,
-                                           calculation_width=calculation_width)
+        return reconstructor.optimize(
+            iterations=iterations,
+            batch_size=batch_size,
+            lr=lr,
+            momentum=momentum,
+            dampening=dampening,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            regularization_factor=regularization_factor, # noqa
+            thread=thread,
+            calculation_width=calculation_width,
+        )
 
 
     def report(self):

From aa684f27eb447ae7efd56bed0652943c682cca6c Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 13 Oct 2025 19:20:42 +0200
Subject: [PATCH 37/55] Change the pattern for Reconstructor so that the
 optimizer is defined at object creation, and move the dataloader creation
 logic to the base optimize() function as it was reused in all subclasses

---
 src/cdtools/reconstructors/adam.py  | 45 ++++++++++------
 src/cdtools/reconstructors/base.py  | 83 +++++++++++++++++++++++------
 src/cdtools/reconstructors/lbfgs.py | 33 ++++++------
 src/cdtools/reconstructors/sgd.py   | 39 +++++++-------
 4 files changed, 133 insertions(+), 67 deletions(-)

diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index 4c54f1f..ac11ee5 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -51,10 +51,12 @@ class AdamReconstructor(Reconstructor):
                  dataset: Ptycho2DDataset,
                  subset: List[int] = None):
 
-        super().__init__(model, dataset, subset)
-
         # Define the optimizer for use in this subclass
-        self.optimizer = t.optim.Adam(self.model.parameters())
+        optimizer = t.optim.Adam(model.parameters())
+
+        super().__init__(model, dataset, optimizer, subset=subset)
+
+
 
     def adjust_optimizer(self,
                          lr: int = 0.005,
@@ -79,11 +81,13 @@ class AdamReconstructor(Reconstructor):
             param_group['betas'] = betas
             param_group['amsgrad'] = amsgrad
 
+        
     def optimize(self,
                  iterations: int,
                  batch_size: int = 15,
                  lr: float = 0.005,
                  betas: Tuple[float] = (0.9, 0.999),
+                 custom_data_loader = None,
                  schedule: bool = False,
                  amsgrad: bool = False,
                  regularization_factor: Union[float, List[float]] = None,
@@ -99,6 +103,12 @@ class AdamReconstructor(Reconstructor):
         (formerly `CDIModel.AD_optimize`) to run a round of reconstruction
         once the dataloader and optimizer hyperparameters have been
         set up.
+        
+        The `batch_size` parameter sets the batch size for the default
+        dataloader. If a custom data loader is desired, it can be passed
+        in to the `custom_data_loader` argument, which will override the
+        `batch_size` and `shuffle` parameters
+
 
         Parameters
         ----------
@@ -115,6 +125,9 @@ class AdamReconstructor(Reconstructor):
         schedule : bool
             Optional, create a learning rate scheduler
             (torch.optim.lr_scheduler._LRScheduler).
+        custom_data_loader : t.utils.data.DataLoader
+            Optional, a custom DataLoader to use. If set, will override
+            batch_size.
         amsgrad : bool
             Optional, whether to use the AMSGrad variant of this algorithm.
         regularization_factor : float or list(float)
@@ -138,18 +151,13 @@ class AdamReconstructor(Reconstructor):
             f'{regularization_factor}, and schedule = {schedule}.\n'
         )
 
-        # 1) The subset statement is contained in Reconstructor.__init__
-
-        # 2) Set up / re-initialize the data laoder
-        self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
-
-        # 3) The optimizer is created in self.__init__, but the
-        #    hyperparameters need to be set up with self.adjust_optimizer
+        # The optimizer is created in self.__init__, but the
+        # hyperparameters need to be set up with self.adjust_optimizer
         self.adjust_optimizer(lr=lr,
                               betas=betas,
                               amsgrad=amsgrad)
 
-        # 4) Set up the scheduler
+        # Set up the scheduler
         if schedule:
             self.scheduler = \
                 t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
@@ -158,8 +166,13 @@ class AdamReconstructor(Reconstructor):
         else:
             self.scheduler = None
 
-        # 5) This is analagous to making a call to CDIModel.AD_optimize
-        return super(AdamReconstructor, self).optimize(iterations,
-                                                       regularization_factor,
-                                                       thread,
-                                                       calculation_width)
+        # Now, we run the optimize routine defined in the base class
+        return super(AdamReconstructor, self).optimize(
+            iterations,
+            batch_size=batch_size,
+            custom_data_loader=custom_data_loader,
+            regularization_factor=regularization_factor,
+            thread=thread,
+            calculation_width=calculation_width,
+            shuffle=shuffle,
+        )
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index 1937661..bae992b 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -40,6 +40,8 @@ class Reconstructor:
         Model for CDI/ptychography reconstruction
     dataset: CDataset
         The dataset to reconstruct against
+    optimizer: torch.optim.Optimizer
+        The optimizer to use for the reconstruction
     subset : list(int) or int
         Optional, a pattern index or list of pattern indices to use
 
@@ -55,31 +57,32 @@ class Reconstructor:
     def __init__(self,
                  model: CDIModel,
                  dataset: CDataset,
+                 optimizer: t.optim.Optimizer,
                  subset: Union[int, List[int]] = None):
+        
         # Store parameters as attributes of Reconstructor
-        self.subset = subset
-
-        # Initialize attributes that must be defined by the subclasses
-        self.optimizer = None
-        self.scheduler = None
-        self.data_loader = None
-
-        # Store the original model
         self.model = model
+        self.optimizer = optimizer
 
-        # Store the dataset
+        # Store the dataset, clipping it to a subset if needed
         if subset is not None:
             # if subset is just one pattern, turn into a list for convenience
             if isinstance(subset, int):
                 subset = [subset]
             dataset = td.Subset(dataset, subset)
+
         self.dataset = dataset
 
+        # Initialize attributes that must be defined by the subclasses
+        self.scheduler = None
+        self.data_loader = None
+
+
     def setup_dataloader(self,
                          batch_size: int = None,
                          shuffle: bool = True):
         """
-        Sets up / re-initializes the dataloader.
+        Sets up or re-initializes the dataloader.
 
         Parameters
         ----------
@@ -96,6 +99,7 @@ class Reconstructor:
         else:
             self.data_loader = td.Dataloader(self.dataset)
 
+            
     def adjust_optimizer(self, **kwargs):
         """
         Change hyperparameters for the utilized optimizer.
@@ -105,7 +109,8 @@ class Reconstructor:
         """
         raise NotImplementedError()
 
-    def _run_epoch(self,
+    
+    def run_epoch(self,
                    stop_event: threading.Event = None,
                    regularization_factor: Union[float, List[float]] = None,
                    calculation_width: int = 10):
@@ -133,6 +138,18 @@ class Reconstructor:
             diffraction pattern intensity
         """
 
+        # Setting this as an explicit catch makes me feel more comfortable
+        # exposing it as a public method. This way a user won't be confused
+        # if they try to use this directly
+        if self.data_loader is None:
+            raise RuntimeError(
+                'No data loader was defined. Please run '
+                'Reconstructor.setup_dataloader() before running '
+                'Reconstructor.run_epoch(), or use Reconstructor.optimize(), '
+                'which does it automatically.'
+            )
+        
+        
         # Initialize some tracking variables
         normalization = 0
         loss = 0
@@ -212,9 +229,12 @@ class Reconstructor:
 
     def optimize(self,
                  iterations: int,
+                 batch_size: int = 1,
+                 custom_data_loader = None,
                  regularization_factor: Union[float, List[float]] = None,
                  thread: bool = True,
-                 calculation_width: int = 10):
+                 calculation_width: int = 10,
+                 shuffle=True):
         """
         Runs a round of reconstruction using the provided optimizer
 
@@ -233,10 +253,25 @@ class Reconstructor:
         the plots. This behavior can be turned off by setting the keyword
         argument 'thread' to False.
 
+        The `batch_size` parameter sets the batch size for the default
+        dataloader. If a custom data loader is desired, it can be passed
+        in to the `custom_data_loader` argument, which will override the
+        `batch_size` and `shuffle` parameters
+
+        Please see `AdamReconstructor.optimize()` for an example of how to
+        override this function when designing a subclass
+
         Parameters
         ----------
         iterations : int
             How many epochs of the algorithm to run.
+        batch_size : int
+            Optional, the batch size to use. Default is 1. This is typically
+            overridden by subclasses with an appropriate default for the
+            specific optimizer.
+        custom_data_loader : torch.utils.data.DataLoader
+            Optional, a custom DataLoader to use. Will override batch_size
+            if set.
         regularization_factor : float or list(float)
             Optional, if the model has a regularizer defined, the set of
             parameters to pass the regularizer method.
@@ -247,6 +282,10 @@ class Reconstructor:
             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.
+        shuffle : bool
+            Optional, enable/disable shuffling of the dataset. This option
+            is intended for diagnostic purposes and should be left as True.
+
 
         Yields
         ------
@@ -255,6 +294,11 @@ class Reconstructor:
             diffraction pattern intensity.
         """
 
+        if custom_data_loader is None:
+            self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
+        else:
+            self.data_loader = custom_data_loader
+
         # We store the current optimizer as a model parameter so that
         # it can be saved and loaded for checkpointing
         self.current_optimizer = self.optimizer
@@ -270,8 +314,10 @@ class Reconstructor:
                         yield float('nan')
                     continue
 
-                yield self._run_epoch(regularization_factor=regularization_factor, # noqa
-                                      calculation_width=calculation_width)
+                yield self.run_epoch(
+                    regularization_factor=regularization_factor, # noqa
+                    calculation_width=calculation_width,
+                )
 
         # But if we do want to thread, it's annoying:
         else:
@@ -282,9 +328,12 @@ class Reconstructor:
             def target():
                 try:
                     result_queue.put(
-                        self._run_epoch(stop_event=stop_event,
-                                        regularization_factor=regularization_factor, # noqa
-                                        calculation_width=calculation_width))
+                        self.run_epoch(
+                            stop_event=stop_event,
+                            regularization_factor=regularization_factor, # noqa
+                            calculation_width=calculation_width,
+                        )
+                    )
                 except Exception as e:
                     # If something bad happens, put the exception into the
                     # result queue
diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py
index 60c6f97..8907782 100644
--- a/src/cdtools/reconstructors/lbfgs.py
+++ b/src/cdtools/reconstructors/lbfgs.py
@@ -53,10 +53,16 @@ class LBFGSReconstructor(Reconstructor):
                  dataset: Ptycho2DDataset,
                  subset: List[int] = None):
 
-        super().__init__(model, dataset, subset)
-
         # Define the optimizer for use in this subclass
-        self.optimizer = t.optim.LBFGS(self.model.parameters())
+        optimizer = t.optim.LBFGS(model.parameters())
+        
+        super().__init__(
+            model,
+            dataset,
+            optimizer,
+            subset=subset,
+        )
+
 
     def adjust_optimizer(self,
                          lr: int = 0.005,
@@ -121,20 +127,17 @@ class LBFGSReconstructor(Reconstructor):
             round of gradient accumulation. Does not affect the result, only
             the calculation speed.
         """
-        # 1) The subset statement is contained in Reconstructor.__init__
 
-        # 2) Set up / re-initialize the data loader. For LBFGS, we load
-        #    all the data at once.
-        self.setup_dataloader(batch_size=len(self.dataset))
-
-        # 3) The optimizer is created in self.__init__, but the
-        #    hyperparameters need to be set up with self.adjust_optimizer
+        # The optimizer is created in self.__init__, but the
+        # hyperparameters need to be set up with self.adjust_optimizer
         self.adjust_optimizer(lr=lr,
                               history_size=history_size,
                               line_search_fn=line_search_fn)
 
-        # 4) This is analagous to making a call to CDIModel.AD_optimize
-        return super(LBFGSReconstructor, self).optimize(iterations,
-                                                        regularization_factor,
-                                                        thread,
-                                                        calculation_width)
+        # Now, we run the optimize routine defined in the base class
+        return super(LBFGSReconstructor, self).optimize(
+            iterations,
+            batch_size=len(self.dataset),
+            regularization_factor=regularization_factor,
+            thread=thread,
+            calculation_width=calculation_width)
diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py
index 5b7bf4e..cb1e26b 100644
--- a/src/cdtools/reconstructors/sgd.py
+++ b/src/cdtools/reconstructors/sgd.py
@@ -49,11 +49,17 @@ class SGDReconstructor(Reconstructor):
                  dataset: Ptycho2DDataset,
                  subset: List[int] = None):
 
-        super().__init__(model, dataset, subset)
-
         # Define the optimizer for use in this subclass
-        self.optimizer = t.optim.SGD(self.model.parameters())
+        optimizer = t.optim.SGD(model.parameters())
 
+        super().__init__(
+            model,
+            dataset,
+            optimizer,
+            subset=subset,
+        )
+
+        
     def adjust_optimizer(self,
                          lr: int = 0.005,
                          momentum: float = 0,
@@ -88,7 +94,7 @@ class SGDReconstructor(Reconstructor):
 
     def optimize(self,
                  iterations: int,
-                 batch_size: int = None,
+                 batch_size: int = 15,
                  lr: float = 2e-7,
                  momentum: float = 0,
                  dampening: float = 0,
@@ -139,25 +145,20 @@ class SGDReconstructor(Reconstructor):
             Optional, enable/disable shuffling of the dataset. This option
             is intended for diagnostic purposes and should be left as True.
         """
-        # 1) The subset statement is contained in Reconstructor.__init__
 
-        # 2) Set up / re-initialize the data laoder
-        if batch_size is not None:
-            self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
-        else:
-            # Use default torch dataloader parameters
-            self.setup_dataloader(batch_size=1, shuffle=False)
-
-        # 3) The optimizer is created in self.__init__, but the
-        #    hyperparameters need to be set up with self.adjust_optimizer
+        # The optimizer is created in self.__init__, but the
+        # hyperparameters need to be set up with self.adjust_optimizer
         self.adjust_optimizer(lr=lr,
                               momentum=momentum,
                               dampening=dampening,
                               weight_decay=weight_decay,
                               nesterov=nesterov)
 
-        # 4) This is analagous to making a call to CDIModel.AD_optimize
-        return super(SGDReconstructor, self).optimize(iterations,
-                                                      regularization_factor,
-                                                      thread,
-                                                      calculation_width)
+        # Now, we run the optimize routine defined in the base class
+        return super(SGDReconstructor, self).optimize(
+            iterations,
+            batch_size=batch_size,
+            regularization_factor=regularization_factor,
+            thread=thread,
+            calculation_width=calculation_width,
+        )

From 3c8de5dc192072fdc5e0753d00183e3c31cca41f Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 16 Oct 2025 14:16:12 +0200
Subject: [PATCH 38/55] Update the example codes

---
 examples/fancy_ptycho.py     | 20 ++++++++++++++++----
 examples/gold_ball_ptycho.py | 11 +++++++----
 examples/gold_ball_split.py  |  9 ++++++---
 examples/simple_ptycho.py    | 11 ++++++++++-
 4 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py
index 1cf5855..067f78d 100644
--- a/examples/fancy_ptycho.py
+++ b/examples/fancy_ptycho.py
@@ -19,19 +19,31 @@ device = 'cuda'
 model.to(device=device)
 dataset.get_as(device=device)
 
+# Now, to do the reconstruction we will use the more flexible pattern of
+# creating an explicit reconstructor. This is what is used behind the
+# scenes by the convenience functions model.<algorithm>_optimize. For a long
+# reconstruction script with multiple steps, it is better to create the
+# reconstructor explicitly.
+#
+# The reconstructor will store the model and dataset and create an appropriate
+# optimizer. This allows the optimizer to persist, along with e.g. estimates
+# of the moments of individual parameters between loops
+recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
+
 # The learning rate parameter sets the alpha for Adam.
 # The beta parameters are (0.9, 0.999) by default
 # The batch size sets the minibatch size
-for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
+for loss in recon.optimize(50, lr=0.02, batch_size=10):
     print(model.report())
     # Plotting is expensive, so we only do it every tenth epoch
     if model.epoch % 10 == 0:
         model.inspect(dataset)
 
 # It's common to chain several different reconstruction loops. Here, we
-# started with an aggressive refinement to find the probe, and now we
-# polish the reconstruction with a lower learning rate and larger minibatch
-for loss in model.Adam_optimize(50, dataset,  lr=0.005, batch_size=50):
+# started with an aggressive refinement to find the probe in the previous
+# loop, and now we polish the reconstruction with a lower learning rate
+# and larger minibatch
+for loss in recon.optimize(50, lr=0.005, batch_size=50):
     print(model.report())
     if model.epoch % 10 == 0:
         model.inspect(dataset)
diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py
index 4747666..4971975 100644
--- a/examples/gold_ball_ptycho.py
+++ b/examples/gold_ball_ptycho.py
@@ -29,6 +29,7 @@ model = cdtools.models.FancyPtycho.from_dataset(
     probe_fourier_crop=pad 
 )
 
+
 # This is a trick that my grandmother taught me, to combat the raster grid
 # pathology: we randomze the our initial guess of the probe positions.
 # The units here are pixels in the object array.
@@ -42,17 +43,20 @@ device = 'cuda'
 model.to(device=device)
 dataset.get_as(device=device)
 
+# Create the reconstructor
+recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
+
 # This will save out the intermediate results if an exception is thrown
 # during the reconstruction
 with model.save_on_exception(
         'example_reconstructions/gold_balls_earlyexit.h5', dataset):
     
-    for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
+    for loss in recon.optimize(20, lr=0.005, batch_size=50):
         print(model.report())
         if model.epoch % 10 == 0:
             model.inspect(dataset)
 
-    for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
+    for loss in recon.optimize(50, lr=0.002, batch_size=100):
         print(model.report())
         if model.epoch % 10 == 0:
             model.inspect(dataset)
@@ -64,8 +68,7 @@ with model.save_on_exception(
 
     # Setting schedule=True automatically lowers the learning rate if
     # the loss fails to improve after 10 epochs
-    for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
-                                    schedule=True):
+    for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True):
         print(model.report())
         if model.epoch % 10 == 0:
             model.inspect(dataset)
diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py
index 8624b9f..9fc5b08 100644
--- a/examples/gold_ball_split.py
+++ b/examples/gold_ball_split.py
@@ -36,15 +36,18 @@ for label, dataset in zip(labels, datasets):
     model.to(device=device)
     dataset.get_as(device=device)
 
+    # Create the reconstructor
+    recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
+
     # For batched reconstructions like this, there's no need to live-plot
     # the progress
-    for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
+    for loss in recon.optimize(20, lr=0.005, batch_size=50):
         print(model.report())
 
-    for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
+    for loss in recon.optimize(50, lr=0.002, batch_size=100):
         print(model.report())
 
-    for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
+    for loss in recon.optimize(100, lr=0.001, batch_size=100,
                                     schedule=True):
         print(model.report())
 
diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py
index 726af94..217d96b 100644
--- a/examples/simple_ptycho.py
+++ b/examples/simple_ptycho.py
@@ -1,3 +1,12 @@
+"""
+Runs a very simple reconstruction using the SimplePtycho model, which was
+designed to be an easy introduction to show how the models are made and used.
+
+For a more realistic example of how to use cdtools for real-world data,
+look at fancy_ptycho.py and gold_ball_ptycho.py, both of which use the
+more powerful FancyPtycho model and include more information on how to
+correct for common sources of error.
+"""
 import cdtools
 from matplotlib import pyplot as plt
 
@@ -13,7 +22,7 @@ device = 'cuda'
 model.to(device=device)
 dataset.get_as(device=device)
 
-# We run the actual reconstruction
+# We run the reconstruction
 for loss in model.Adam_optimize(100, dataset, batch_size=10):
     # We print a quick report of the optimization status
     print(model.report())

From 8328b852d393a164aa89abb2660ae242fdfd21a7 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 16 Oct 2025 14:43:04 +0200
Subject: [PATCH 39/55] Update the docs to include the reconstructors class and
 discuss their use in the examples

---
 docs/source/examples.rst           | 10 ++++++++--
 docs/source/index.rst              |  1 +
 docs/source/reconstructors.rst     |  5 +++++
 examples/fancy_ptycho.py           | 14 +++++---------
 src/cdtools/reconstructors/base.py | 18 ++++++++++--------
 5 files changed, 29 insertions(+), 19 deletions(-)
 create mode 100644 docs/source/reconstructors.rst

diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index 517f1b6..a3ef41e 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -31,9 +31,9 @@ When reading this script, note the basic workflow. After the data is loaded, a m
 
 Next, the model is moved to the GPU using the :code:`model.to` function. Any device understood by :code:`torch.Tensor.to` can be specified here. The next line is a bit more subtle - the dataset is told to move patterns to the GPU before passing them to the model using the :code:`dataset.get_as` function. This function does not move the stored patterns to the GPU. If there is sufficient GPU memory, the patterns can also be pre-moved to the GPU using :code:`dataset.to`, but the speedup is empirically quite small.
 
-Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at every epoch, to allow some monitoring code to be run.
+Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run.
 
-Finally, the results can be studied using :code:`model.inspect(dataet)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
+Finally, the results can be studied using :code:`model.inspect(dataset)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
 
 
 Fancy Ptycho
@@ -63,6 +63,12 @@ By default, FancyPtycho will also optimize over the following model parameters,
 
 These corrections can be turned off (on) by calling :code:`model.<parameter>.requires_grad = False #(True)`.
 
+Note as well two other changes that are made in this script, when compared to `simple_ptycho.py`. First, a `Reconstructor` object is explicitly created, in this case an `AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to `Reconstructor.optimize()`.
+
+We use this pattern, instead of the simpler call to `model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object.
+
+In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. `model.LBFGS_optimize()`).
+
 
 Gold Ball Ptycho
 ----------------
diff --git a/docs/source/index.rst b/docs/source/index.rst
index c6fcb8d..b8cd288 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -9,6 +9,7 @@
    general
    datasets
    models
+   reconstructors
    tools/index
    indices_tables
    
diff --git a/docs/source/reconstructors.rst b/docs/source/reconstructors.rst
new file mode 100644
index 0000000..2e78a74
--- /dev/null
+++ b/docs/source/reconstructors.rst
@@ -0,0 +1,5 @@
+Reconstructors
+==============
+
+.. automodule:: cdtools.reconstructors
+   :members:
diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py
index 067f78d..ebee68f 100644
--- a/examples/fancy_ptycho.py
+++ b/examples/fancy_ptycho.py
@@ -19,15 +19,11 @@ device = 'cuda'
 model.to(device=device)
 dataset.get_as(device=device)
 
-# Now, to do the reconstruction we will use the more flexible pattern of
-# creating an explicit reconstructor. This is what is used behind the
-# scenes by the convenience functions model.<algorithm>_optimize. For a long
-# reconstruction script with multiple steps, it is better to create the
-# reconstructor explicitly.
-#
-# The reconstructor will store the model and dataset and create an appropriate
-# optimizer. This allows the optimizer to persist, along with e.g. estimates
-# of the moments of individual parameters between loops
+# For this script, we use a slightly different pattern where we explicitly
+# create a `Reconstructor` class to orchestrate the reconstruction. The
+# reconstructor will store the model and dataset and create an appropriate
+# optimizer. This allows the optimizer to persist between loops, along with
+# e.g. estimates of the moments of individual parameters
 recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
 
 # The learning rate parameter sets the alpha for Adam.
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index bae992b..5177b79 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -45,14 +45,16 @@ class Reconstructor:
     subset : list(int) or int
         Optional, a pattern index or list of pattern indices to use
 
-    Important attributes:
-    - **model** -- Always points to the core model used.
-    - **optimizer** -- A `torch.optim.Optimizer` that must be defined when
-        initializing the Reconstructor subclass.
-    - **scheduler** -- A `torch.optim.lr_scheduler` that may be defined during
-        the `optimize` method.
-    - **data_loader** -- A torch.utils.data.DataLoader that is defined by
-        calling the `setup_dataloader` method.
+    Attributes
+    ----------
+    model : CDIModel
+        Points to the core model used.
+    optimizer : torch.optim.Optimizer
+        Must be defined when initializing the Reconstructor subclass.
+    scheduler : torch.optim.lr_scheduler, optional
+        May be defined during the ``optimize`` method.
+    data_loader : torch.utils.data.DataLoader
+        Defined by calling the ``setup_dataloader`` method.
     """
     def __init__(self,
                  model: CDIModel,

From f022b5e2c4ba90f9876d842b5b7407aece1cf2f5 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Thu, 16 Oct 2025 15:58:28 +0200
Subject: [PATCH 40/55] Get the tests passing again after the change of the
 optimization functions in the model classes

---
 tests/models/test_fancy_ptycho.py |  9 +++++++--
 tests/test_reconstructors.py      | 16 +++++++++-------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 4182c93..8bc4d87 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -56,8 +56,8 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         dataset,
         n_modes=3,
         oversampling=2,
-        exponentiate_obj=True,
         dm_rank=2,
+        exponentiate_obj=True,
         probe_support_radius=120,
         propagation_distance=5e-3,
         units='mm',
@@ -70,7 +70,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
     model.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
 
-    for loss in model.Adam_optimize(70, dataset, lr=0.02, batch_size=10):
+    for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
@@ -79,6 +79,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
+            
+    for loss in model.Adam_optimize(25, dataset, lr=0.001, batch_size=50):
+        print(model.report())
+        if show_plot and model.epoch % 10 == 0:
+            model.inspect(dataset)
 
     model.tidy_probes()
 
diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py
index 3995b6f..2e30828 100644
--- a/tests/test_reconstructors.py
+++ b/tests/test_reconstructors.py
@@ -91,7 +91,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
           ' reconstruction_device,', reconstruction_device)
     t.manual_seed(0)
 
-    for i, iterations in enumerate(epoch_tup):
+    # We only need to test the first loop to ensure it's identical
+    for i, iterations in enumerate(epoch_tup[:1]):
         for loss in model.Adam_optimize(iterations,
                                         dataset,
                                         lr=lr_tup[i],
@@ -106,14 +107,15 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         model.inspect(dataset)
         model.compare(dataset)
 
-    # Ensure equivalency between the model reconstructions
-    assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
+    # Ensure equivalency between the model reconstructions during the first
+    # pass, where they should be identical
+    assert np.allclose(model_recon.loss_history[:epoch_tup[0]], model.loss_history[:epoch_tup[0]])
 
     # Ensure reconstructions have reached a certain loss tolerance. This just
     # comes from running a reconstruction when it was working well and
     # choosing a rough value. If it triggers this assertion error, something
     # changed to make the final quality worse!
-    assert model.loss_history[-1] < 0.0001
+    assert model_recon.loss_history[-1] < 0.0001
 
 
 @pytest.mark.slow
@@ -184,7 +186,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
     print('Running reconstruction using CDIModel.LBFGS_optimize.' +
           'optimize on provided reconstruction_device,', reconstruction_device)
     t.manual_seed(0)
-    for i, iterations in enumerate(epoch_tup):
+    for i, iterations in enumerate(epoch_tup[:1]):
         for loss in model.LBFGS_optimize(iterations,
                                          dataset,
                                          lr=0.4,
@@ -198,12 +200,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
         model.compare(dataset)
 
     # Check loss equivalency between the two reconstructions
-    assert np.allclose(model.loss_history[-1], model_recon.loss_history[-1])
+    assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]])
 
     # The final loss when testing this was 2.28607e-3. Based on this, we set
     # a threshold of 2.3e-3 for the tested loss. If this value has been
     # exceeded, the reconstructions have gotten worse.
-    assert model.loss_history[-1] < 0.0023
+    assert model_recon.loss_history[-1] < 0.0023
 
 
 @pytest.mark.slow

From 9d633074d8fa509087a48f26668ec3aa7f70d60c Mon Sep 17 00:00:00 2001
From: Abe Levitan <allevitan@gmail.com>
Date: Fri, 17 Oct 2025 18:37:38 -0300
Subject: [PATCH 41/55] Apply suggestions from code review

Co-authored-by: Dayne Yoshiki Sasaki <37006268+yoshikisd@users.noreply.github.com>
---
 src/cdtools/reconstructors/adam.py | 2 +-
 src/cdtools/reconstructors/base.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py
index ac11ee5..e298bdb 100644
--- a/src/cdtools/reconstructors/adam.py
+++ b/src/cdtools/reconstructors/adam.py
@@ -87,7 +87,7 @@ class AdamReconstructor(Reconstructor):
                  batch_size: int = 15,
                  lr: float = 0.005,
                  betas: Tuple[float] = (0.9, 0.999),
-                 custom_data_loader = None,
+                 custom_data_loader: t.utils.data.DataLoader = None,
                  schedule: bool = False,
                  amsgrad: bool = False,
                  regularization_factor: Union[float, List[float]] = None,
diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index 5177b79..cf3de8c 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -232,7 +232,7 @@ class Reconstructor:
     def optimize(self,
                  iterations: int,
                  batch_size: int = 1,
-                 custom_data_loader = None,
+                 custom_data_loader: torch.utils.data.DataLoader = None,
                  regularization_factor: Union[float, List[float]] = None,
                  thread: bool = True,
                  calculation_width: int = 10,

From 7d390157199b90fefddea691a6d5eaba737ba9c9 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Fri, 17 Oct 2025 23:41:08 +0200
Subject: [PATCH 42/55] response to dayne's review

---
 tests/test_reconstructors.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py
index 2e30828..132a8c2 100644
--- a/tests/test_reconstructors.py
+++ b/tests/test_reconstructors.py
@@ -16,7 +16,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         2) We are only using the single-GPU dataloading method
         3) Ensure `recon.model` points to the original `model`
         4) Reconstructions performed by `Adam.optimize` and
-           `model.Adam_optimize` calls produce identical results.
+           `model.Adam_optimize` calls produce identical results when
+           run over one round of optimization.
         5) The quality of the reconstruction remains below a specified
            threshold.
         5) Ensure that the FancyPtycho model works fine and dandy with the
@@ -129,7 +130,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
            hyperparameters
         2) Ensure `recon.model` points to the original `model`
         3) Reconstructions performed by `LBFGS.optimize` and
-           `model.LBFGS_optimize` calls produce identical results.
+           `model.LBFGS_optimize` calls produce identical results when
+           run over one round of reconstruction.
         4) The quality of the reconstruction remains below a specified
            threshold.
         5) Ensure that the RPI model works fine and dandy with the
@@ -216,7 +218,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
            hyperparameters
         3) Ensure `recon.model` points to the original `model`
         4) Reconstructions performed by `SGD.optimize` and
-           `model.SGD_optimize` calls produce identical results.
+           `model.SGD_optimize` calls produce identical results
+           when run over one round of reconstruction.
         5) The quality of the reconstruction remains below a specified
            threshold.
         5) Ensure that the FancyPtycho model works fine and dandy with the

From 3eec9c0cec31775ec3d5c2ea9a04e7af47b5def5 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Sun, 19 Oct 2025 12:45:19 +0200
Subject: [PATCH 43/55] Add the original positions to the fancy_ptycho position
 plotting

---
 src/cdtools/models/fancy_ptycho.py     | 24 +++++++++++++++++++++++-
 src/cdtools/tools/plotting/plotting.py | 24 +++++++++++++++++++-----
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 6220299..fa4757f 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -798,6 +798,28 @@ class FancyPtycho(CDIModel):
             **kwargs),
 
         
+    def plot_translations_and_originals(self, fig, dataset):
+        """Only used to make a plot for the plot list."""
+        p.plot_translations(
+            dataset.translations,
+            fig=fig,
+            units=self.units,
+            label='original translations',
+            color='#CCCCCC',
+            marker='o',
+        )
+        p.plot_translations(
+            self.corrected_translations(dataset),
+            fig=fig,
+            units=self.units,
+            clear_fig=False,
+            label='refined translations',
+            color='k',
+            marker='.'
+        )
+        plt.legend()
+        
+        
     plot_list = [
         ('',
          lambda self, fig, dataset: self.plot_wavefront_variation(
@@ -895,7 +917,7 @@ class FancyPtycho(CDIModel):
          lambda self: self.exponentiate_obj),
 
         ('Corrected Translations',
-         lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
+         lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset)),
         ('Background',
          lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
         ('Quantum Efficiency Mask',
diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py
index abd8f84..f2357f7 100644
--- a/src/cdtools/tools/plotting/plotting.py
+++ b/src/cdtools/tools/plotting/plotting.py
@@ -522,7 +522,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
                       units=units, show_cbar=False, **kwargs)
 
 
-def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
+def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, clear_fig=True, label=None, color=None, marker='.', **kwargs):
     """Plots a set of probe translations in a nicely formatted way
 
     Parameters
@@ -537,6 +537,14 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
         Whether to plot lines indicating the path taken
     invert_xaxis : bool
         Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
+    clear_fig : bool
+        Default is True. Whether to clear the figure before plotting.
+    label : str
+        Default is None. A label to give the plotted markers for a legend.
+    color : str
+        Default is None. The color to plot the markers in. By default, will follow the matplotlib color cycle.
+    color : str
+        Default is '.'. The marker style to plot with.
     \\**kwargs
         All other args are passed to fig.add_subplot(111, \\**kwargs)
 
@@ -554,18 +562,24 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
-        plt.gcf().clear()
+        if clear_fig:
+            plt.gcf().clear()
 
     if isinstance(translations, t.Tensor):
         translations = translations.detach().cpu().numpy()
 
     translations = translations * factor
-    plt.plot(translations[:,0], translations[:,1],'k.')
+
+    linestyle = '-' if lines else 'None'
+    linewidth = 1 if lines else 0
+    plt.plot(translations[:,0], translations[:,1],
+             marker=marker, linestyle=linestyle,
+             label=label, color=color, 
+             linewidth=linewidth)
+    
     if invert_xaxis:
         plt.gca().invert_xaxis()
         
-    if lines:
-        plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')
 

From 86e85e18530ea71b71fee6a2da696fffb6e95a7c Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Sun, 19 Oct 2025 13:08:16 +0200
Subject: [PATCH 44/55] Fix a bug where the x-axis was double flipped

---
 src/cdtools/tools/plotting/plotting.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py
index f2357f7..0f570ea 100644
--- a/src/cdtools/tools/plotting/plotting.py
+++ b/src/cdtools/tools/plotting/plotting.py
@@ -578,7 +578,11 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver
              linewidth=linewidth)
     
     if invert_xaxis:
-        plt.gca().invert_xaxis()
+        ax = plt.gca()
+        x_min, x_max = ax.get_xlim()
+        # Protect against flipping twice if plotting on top of existing graph
+        if x_min <= x_max:
+            ax.invert_xaxis()
         
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')

From e38ebd737277ba565506d654b0ceecaca0c267e0 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Fri, 24 Oct 2025 15:11:46 +0200
Subject: [PATCH 45/55] Fix a high priority bug with the masking system which
 was introduced in the switch to reconstructors classes

---
 src/cdtools/reconstructors/base.py | 2 +-
 tests/models/test_fancy_ptycho.py  | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py
index cf3de8c..1666814 100644
--- a/src/cdtools/reconstructors/base.py
+++ b/src/cdtools/reconstructors/base.py
@@ -190,7 +190,7 @@ class Reconstructor:
                     sim_patterns = self.model.forward(*inp)
 
                     # Calculate the loss
-                    if hasattr(self, 'mask'):
+                    if hasattr(self.model, 'mask'):
                         loss = self.model.loss(pats,
                                                sim_patterns,
                                                mask=self.model.mask)
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 8bc4d87..02893f9 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -52,6 +52,10 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
     print('\nTesting performance on the standard transmission ptycho dataset')
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
 
+    # Test the masking system
+    dataset.mask[110:115,65:70] = 0
+    dataset.patterns[...,~dataset.mask] = t.max(dataset.patterns)
+    
     model = cdtools.models.FancyPtycho.from_dataset(
         dataset,
         n_modes=3,

From 008fef6244b7990c8ce5364bf816cc473891a3e5 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Tue, 28 Oct 2025 16:06:21 +0100
Subject: [PATCH 46/55] Add cdtools.__version__ and set it up so that there is
 a single source of truth for the version number. Still not done with
 setuptools_scm, but one step at a time)

---
 setup.py                | 12 +++++++++++-
 src/cdtools/__init__.py |  2 ++
 src/cdtools/_version.py |  1 +
 3 files changed, 14 insertions(+), 1 deletion(-)
 create mode 100644 src/cdtools/_version.py

diff --git a/setup.py b/setup.py
index 4e62c0f..ae710d7 100644
--- a/setup.py
+++ b/setup.py
@@ -1,11 +1,21 @@
 import setuptools
+import os
+import re
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
+# read version from src/cdtools/_version.py
+version_file = os.path.join("src/cdtools", "_version.py")
+with open(version_file) as f:
+    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M)
+if not version_match:
+    raise RuntimeError("Unable to find version string.")
+version = version_match.group(1)
+    
 setuptools.setup(
     name="cdtools-py",
-    version="0.3.0",
+    version=version,
     python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
     author="Abe Levitan",
     author_email="abraham.levitan@psi.ch",
diff --git a/src/cdtools/__init__.py b/src/cdtools/__init__.py
index 9132209..19dc4c6 100644
--- a/src/cdtools/__init__.py
+++ b/src/cdtools/__init__.py
@@ -6,6 +6,8 @@ warnings.filterwarnings("ignore",
 
 __all__ = ['tools', 'datasets', 'models', 'reconstructors']
 
+from ._version import __version__
+
 from cdtools import tools
 from cdtools import datasets
 from cdtools import models
diff --git a/src/cdtools/_version.py b/src/cdtools/_version.py
new file mode 100644
index 0000000..9163035
--- /dev/null
+++ b/src/cdtools/_version.py
@@ -0,0 +1 @@
+__version__ = "0.3.1.dev"

From f79836bfc26225a92075ff1d85e22211436ca758 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Wed, 29 Oct 2025 21:40:12 -0700
Subject: [PATCH 47/55] Add tests for version existence and semantic format

---
 tests/test_version.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)
 create mode 100644 tests/test_version.py

diff --git a/tests/test_version.py b/tests/test_version.py
new file mode 100644
index 0000000..5b6a778
--- /dev/null
+++ b/tests/test_version.py
@@ -0,0 +1,18 @@
+from cdtools import __version__
+import re
+
+
+def test_version_exists():
+    """Test that version is defined and not empty."""
+    assert __version__
+    assert isinstance(__version__, str)
+    assert len(__version__) > 0
+
+
+def test_version_format():
+    """Test that version follows semantic versioning format."""
+    # Basic semantic versioning pattern (X.Y.Z with optional pre-release)
+    pattern = r"^\d+\.\d+\.\d+(?:[-.]?(?:alpha|beta|rc|dev)\d*)?$"
+    assert re.match(
+        pattern, __version__
+    ), f"Version '{__version__}' doesn't follow semantic versioning"

From be28903be7e91e0819a93b7fe57474cbfdbcd47c Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Nov 2025 12:23:09 -0800
Subject: [PATCH 48/55] refactor: CI installation process to use 'uv' for
 faster dependency management

---
 .github/workflows/main.yml    | 21 ++++++++-------
 .github/workflows/publish.yml | 12 +++++----
 pyproject.toml                | 50 +++++++++++++++++++++++++++++++++-
 requirements.txt              | 11 --------
 setup.py                      | 51 -----------------------------------
 5 files changed, 67 insertions(+), 78 deletions(-)
 delete mode 100644 requirements.txt
 delete mode 100644 setup.py

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 5cc6d9b..84f4a00 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -22,14 +22,15 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install dependencies
+      - name: Install uv
         run: |
-          pip install --upgrade pip
-          pip install -r requirements.txt
-          pip install -e . --no-deps
+          pip install uv
+      - name: Install dependencies with uv
+        run: |
+          uv pip install ."[tests]"
 
       - name: Run tests
-        run: pytest
+        run: python -m pytest
 
   build-docs:
     if: github.event_name == 'push' && github.ref == 'refs/heads/master'
@@ -52,12 +53,12 @@ jobs:
         with:
           python-version: '3.9'
 
-      - name: Install dependencies
+      - name: Install uv
         run: |
-          pip install --upgrade pip
-          pip install -r requirements.txt
-          pip install sphinx sphinx_rtd_theme sphinx-argparse
-          pip install -e . --no-deps
+          pip install uv
+      - name: Install dependencies with uv
+        run: |
+          uv pip install ."[docs]""
 
       - name: Build docs
         working-directory: docs
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 8ff3518..dbb1a1a 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -19,12 +19,14 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: "3.x"
-    - name: Install dependencies
+    - name: Install uv
       run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel
-    - name: Build package (setup.py)
+        pip install uv
+    - name: Install project with uv
       run: |
-        python setup.py sdist bdist_wheel
+        uv pip install -e .
+    - name: Build package with uv
+      run: |
+        uv build
     - name: Publish package distributions to PyPI
       uses: pypa/gh-action-pypi-publish@release/v1
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index b7fd46a..703942e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,51 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cdtools-py"
+description = "Tools for coherent diffractive imaging and ptychography"
+readme = "README.md"
+requires-python = ">=3.8"
+license = { file = "LICENSE.txt" }
+authors = [
+	{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" },
+	{ name = "Madelyn Cain" },
+	{ name = "Anastasiia Kutakh" }
+]
+maintainers = [
+	{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" }
+]
+keywords = ["ptychography", "CDI", "imaging", "torch", "differentiable"]
+classifiers = [
+	"Programming Language :: Python :: 3",
+	"Operating System :: OS Independent",
+	"License :: OSI Approved :: MIT License"
+]
+urls = { "Homepage" = "https://github.com/cdtools-developers/cdtools", "Documentation" = "https://cdtools-developers.github.io/cdtools/" }
+dependencies = [
+	"numpy>=1.0",
+	"scipy>=1.0",
+	"matplotlib>=2.0",
+	"torch>=2.3.0",
+	"h5py>=2.1",
+	"python-dateutil",
+]
+dynamic = ["version"]
+
+[tool.setuptools.dynamic]
+version = {attr = "cdtools._version.__version__"}
+
+[project.optional-dependencies]
+tests = [
+	"pytest",
+	"pooch"
+]
+docs = [
+	"sphinx>=4.3.0",
+	"sphinx-argparse",
+	"sphinx_rtd_theme>=0.5.1"
+]
+
 [tool.ruff]
-# Decrease the maximum line length to 79 characters.
 line-length = 79
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 1c584c8..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,11 +0,0 @@
-numpy>=1.0
-scipy>=1.0
-matplotlib>=2.0  # 2.0 introduces better colormaps which are used by default
-torch>=1.9.0 #1.9.0 implements support for autograd on indexed complex tensors
-h5py>=2.1
-python-dateutil
-pytest
-pooch
-sphinx>=4.3.0 # Fixes a bug with bulleted lists
-sphinx-argparse
-sphinx_rtd_theme>=0.5.1 # Fixes a bug with bulleted lists
diff --git a/setup.py b/setup.py
deleted file mode 100644
index ae710d7..0000000
--- a/setup.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import setuptools
-import os
-import re
-
-with open("README.md", "r") as fh:
-    long_description = fh.read()
-
-# read version from src/cdtools/_version.py
-version_file = os.path.join("src/cdtools", "_version.py")
-with open(version_file) as f:
-    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M)
-if not version_match:
-    raise RuntimeError("Unable to find version string.")
-version = version_match.group(1)
-    
-setuptools.setup(
-    name="cdtools-py",
-    version=version,
-    python_requires='>3.8', # recommended minimum version for pytorch 2.3.0
-    author="Abe Levitan",
-    author_email="abraham.levitan@psi.ch",
-    description="Tools for coherent diffractive imaging and ptychography",
-    long_description=long_description,
-    long_description_content_type="text/markdown",
-    url="https://github.com/cdtools-developers/cdtools",
-    install_requires=[
-        "numpy>=1.0",
-        "scipy>=1.0",
-        "matplotlib>=2.0", # 2.0 has better colormaps which are used by default
-        "python-dateutil",
-        "torch>=2.3.0", #2.3.0 is the earliest release for which L-BFGS works directly on complex-valued leaf tensors
-        "h5py>=2.1"],
-    extras_require={
-        'tests': [
-            "pytest",
-            "pooch",
-        ],
-        'docs': [
-            "sphinx>=4.3.0",
-            "sphinx-argparse",
-            "sphinx_rtd_theme>=0.5.1"
-        ]
-    },
-    package_dir={"": "src"},
-    packages=setuptools.find_packages("src"),
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "Operating System :: OS Independent",
-    ],
-)
-

From 728cf9d59029457901b2f2aa9f41dc35c0584b50 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Nov 2025 12:23:25 -0800
Subject: [PATCH 49/55] update readme and installation docs

---
 README.md                    | 45 ++++++++++++++++++++
 docs/source/installation.rst | 80 +++++++++++++++++++++++-------------
 2 files changed, 97 insertions(+), 28 deletions(-)

diff --git a/README.md b/README.md
index 76371bf..71f81e6 100644
--- a/README.md
+++ b/README.md
@@ -26,12 +26,57 @@ model.compare(dataset) # See how the simulated and measured patterns compare
 plt.show()
 ```
 
+# Installation
+
+CDTools can be installed in several ways depending on your needs. For most users, installation from pypi is recommended. For developers or those who want the latest features, installation from source is available.
+
+## Installation from pypi
+
 CDTools can be installed via pip as the [cdtools-py](https://pypi.org/project/cdtools-py/) package on [PyPI](https://pypi.org/):
 
 ```bash
 $ pip install cdtools-py
 ```
 
+or using [uv](https://github.com/astral-sh/uv):
+
+```bash
+$ uv pip install cdtools-py
+```
+
+## Installation from Source
+
+For development or to access the latest features, CDTools can be installed directly from source:
+
+
+```bash
+$ git clone https://github.com/cdtools-developers/cdtools.git
+$ cd cdtools
+$ pip install -e .
+```
+
+
+or using [uv](https://github.com/astral-sh/uv):
+
+```bash
+$ git clone https://github.com/cdtools-developers/cdtools.git
+$ cd cdtools
+$ uv pip install -e .
+```
+
+## Installing for Contributors (with tests and docs dependencies)
+
+If you want to run the test suite or build the documentation, install with the extra dependencies:
+
+```bash
+$ pip install -e ."[tests,docs]"
+```
+or with uv:
+```bash
+$ uv pip install -e ."[tests,docs]"
+```
+
+
 Further documentation is found [here](https://cdtools-developers.github.io/cdtools/).
 
 
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index cf86477..486588f 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -15,6 +15,14 @@ To install from `PyPI`_, run:
 	  
    $ pip install cdtools-py
 
+or you can use `uv`_ for a faster installation:
+
+.. _`uv`: https://github.com/astral-sh/uv
+
+.. code:: bash
+
+   $ uv pip install cdtools-py
+
 Pytorch, a major dependence of CDTools, often needs to be installed with a specific CUDA version for machine compatability. If you run into issues with pytorch, consider first installing pytorch into your environment using the instructions on `the pytorch site`_.
 
 .. _`the pytorch site`: https://pytorch.org/get-started/locally/
@@ -31,18 +39,51 @@ The source code for CDTools is hosted on `Github`_.
 .. _`Github`: https://github.com/cdtools-developers/cdtools
 
 
-Step 2: Install Dependencies
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-The major dependency for CDTools is pytorch version 2.3.0 or greater. Because the details of the pytorch installation can vary depending on platform and GPU availability, it is recommended that you first install pytorch using the instructions on `the pytorch site`_. The remaining dependencies can be installed by running the following command from the top level directory of the git repository:
+To download the source code, you can either clone the repository using git:
 
 .. code:: bash
-	  
-   $ pip install -r requirements.txt
 
-Note that several optional dependencies used for testing and documentation will also be installed. The full set of dependencies and minimum requirements are listed below. CDTools is reguarly tested with the latest versions of these packages and with python 3.8 through 3.12.
+   $ git clone https://github.com/cdtools-developers/cdtools.git
 
-Required dependencies:
+or you can download a zip file of the repository from the `releases page`_.
+
+.. _`releases page`: https://github.com/cdtools-developers/cdtools/releases
+
+
+Step 2: Install
+^^^^^^^^^^^^^^^
+
+Move to the directory where you downloaded the source code. It is recommended that you create a new python virtual environment to install CDTools into.
+
+Installation using pip and uv. Editable mode is recommended for development purposes and is added with the `-e` flag.
+
+.. code:: bash
+
+   $ pip install -e .
+
+or using uv:
+
+.. code:: bash
+
+   $ uv pip install -e .
+
+
+To install the required test and documentation dependencies as well, use:
+
+.. code:: bash
+
+   $ pip install -e ."[tests,docs]"
+
+or using uv:
+
+.. code:: bash
+
+   $ uv pip install -e ."[tests,docs]"
+
+CDTools is reguarly tested with the latest versions of these packages and with python 3.8 through 3.12.
+
+
+Required dependencies (see pyproject.toml for all details):
 
    * `numpy <http://www.numpy.org>`_ >= 1.0
    * `scipy <http://www.scipy.org>`_ >= 1.0
@@ -63,30 +104,13 @@ Optional dependencies for building docs:
    * `sphinx_rtd_theme <https://sphinx-rtd-theme.readthedocs.io/en/stable/>`_ >= 0.5.1
 
 
-Step 3: Install
-^^^^^^^^^^^^^^^
-
-To install CDTools, run the following command from the top level directory of the git repository:
-
-.. code:: bash
-	  
-   $ pip install -e . --no-deps
-
-   
-This will install CDTools in developer mode, so that changes to the code will propagate to the installed version immediately. This is best if you plan to actively develop CDTools. If you simply need a custom environment, you can also install CDTools in standard mode using:
-
-.. code:: bash
-	  
-   $ pip install . --no-deps
-   
-  
-Step 4: Run The Tests
-^^^^^^^^^^^^^^^^^^^^^
+Optional step 4: Run The Tests
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 To ensure that the installation has worked correctly, it is recommended that you run the unit tests. Execute the following command from the top level directory of the git repository:
 
 .. code:: bash
 
-   $ pytest
+   $ python -m pytest
 
 If any tests fail, make sure that you have all the noted dependencies properly installed. If you do, and things still aren't working, `open an issue on the github page <https://github.com/cdtools-developers/cdtools/issues>`_ and we'll get to the bottom of it.

From 18c794b5db8d03e4d6bc0e5f1c7c108d76bfb427 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Nov 2025 12:37:44 -0800
Subject: [PATCH 50/55] fix: add virtual environment creation step in CI github
 workflows

---
 .github/workflows/main.yml    | 12 +++++++++++-
 .github/workflows/publish.yml |  5 +++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 84f4a00..3316f7d 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -25,6 +25,11 @@ jobs:
       - name: Install uv
         run: |
           pip install uv
+
+      - name: Create virtual environment
+        run: |
+          uv venv
+
       - name: Install dependencies with uv
         run: |
           uv pip install ."[tests]"
@@ -56,9 +61,14 @@ jobs:
       - name: Install uv
         run: |
           pip install uv
+      
+      - name: Create virtual environment
+        run: |
+          uv venv
+      
       - name: Install dependencies with uv
         run: |
-          uv pip install ."[docs]""
+          uv pip install ."[docs]"
 
       - name: Build docs
         working-directory: docs
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index dbb1a1a..9beed2b 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -22,6 +22,11 @@ jobs:
     - name: Install uv
       run: |
         pip install uv
+    
+    - name : Create virtual environment
+      run: |
+        uv venv
+
     - name: Install project with uv
       run: |
         uv pip install -e .

From e94499a6a4eecc5d289a8787ec334cf3954abf85 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Nov 2025 12:42:32 -0800
Subject: [PATCH 51/55] fix: update pytest command to use 'uv' for consistency
 in CI workflow

---
 .github/workflows/main.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 3316f7d..672014e 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -35,7 +35,8 @@ jobs:
           uv pip install ."[tests]"
 
       - name: Run tests
-        run: python -m pytest
+        run: |
+          uv run pytest
 
   build-docs:
     if: github.event_name == 'push' && github.ref == 'refs/heads/master'

From 5a3e154e8df6181bc8cb9904b0c9b189a537da69 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Nov 2025 12:50:59 -0800
Subject: [PATCH 52/55] add Python 3.13 to the CI workflow matrix

---
 .github/workflows/main.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 672014e..db280c7 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
     continue-on-error: true
 
     steps:

From 11ed3df35e72bb8844c6bbfe3b66c8af8b76c1d2 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Tue, 11 Nov 2025 14:33:55 -0600
Subject: [PATCH 53/55] ending python 3.8 support, adding python 3.14 to
 automatic test

---
 .github/workflows/main.yml   | 2 +-
 docs/source/installation.rst | 4 ++--
 pyproject.toml               | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index db280c7..f748b51 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
+        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14']
     continue-on-error: true
 
     steps:
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index 486588f..11c3aad 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -1,7 +1,7 @@
 Installation
 ============
 
-CDTools supports python >=3.8 and can be installed via pip as the the `cdtools-py`_ package on `PyPI`_. If you plan to contribute to the code or need a custom environment, installation from source is also possible.
+CDTools supports python >=3.9 and can be installed via pip as the the `cdtools-py`_ package on `PyPI`_. If you plan to contribute to the code or need a custom environment, installation from source is also possible.
 
 .. _`cdtools-py`: https://pypi.org/project/cdtools-py/
 .. _`PyPI`: https://pypi.org/
@@ -80,7 +80,7 @@ or using uv:
 
    $ uv pip install -e ."[tests,docs]"
 
-CDTools is reguarly tested with the latest versions of these packages and with python 3.8 through 3.12.
+CDTools is reguarly tested with the latest versions of these packages and with python 3.9 through 3.14.
 
 
 Required dependencies (see pyproject.toml for all details):
diff --git a/pyproject.toml b/pyproject.toml
index 703942e..11ccebc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "cdtools-py"
 description = "Tools for coherent diffractive imaging and ptychography"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 license = { file = "LICENSE.txt" }
 authors = [
 	{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" },

From 3c9e9eb1c7fd58a493dace13c156d70c46efa568 Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Wed, 12 Nov 2025 08:59:37 -0600
Subject: [PATCH 54/55] adding pure pip install pipeline to CI

---
 .github/workflows/main.yml | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index f748b51..a24af2b 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -7,7 +7,7 @@ on:
     branches: [ master ]
 
 jobs:
-  test:
+  test-uv-pip-install:
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -37,6 +37,28 @@ jobs:
       - name: Run tests
         run: |
           uv run pytest
+  
+  test-pip-install:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.14']
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies with pip
+        run: |
+          pip install ."[tests]"
+
+      - name: Run tests
+        run: |
+          pytest
 
   build-docs:
     if: github.event_name == 'push' && github.ref == 'refs/heads/master'

From 3f46caf757daba51e39ec03193cc1b0b6af0db0c Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Wed, 12 Nov 2025 09:04:53 -0600
Subject: [PATCH 55/55] update authors list in pyproject.toml and move further
 doc link

---
 README.md      | 6 ++----
 pyproject.toml | 2 ++
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 71f81e6..742b472 100644
--- a/README.md
+++ b/README.md
@@ -26,6 +26,8 @@ model.compare(dataset) # See how the simulated and measured patterns compare
 plt.show()
 ```
 
+Further documentation is found [here](https://cdtools-developers.github.io/cdtools/).
+
 # Installation
 
 CDTools can be installed in several ways depending on your needs. For most users, installation from pypi is recommended. For developers or those who want the latest features, installation from source is available.
@@ -76,10 +78,6 @@ or with uv:
 $ uv pip install -e ."[tests,docs]"
 ```
 
-
-Further documentation is found [here](https://cdtools-developers.github.io/cdtools/).
-
-
 CDTools was developed in the [photon scattering lab](https://scattering.mit.edu/) at MIT, and further development took place within the [computational x-ray imaging group](https://www.psi.ch/en/cxi) at PSI. The code is distributed under an MIT (a.k.a. Expat) license. If you would like to publish any work that uses CDTools, please contact [Abe Levitan](mailto:abraham.levitan@psi.ch).
 
 Have a wonderful day!
diff --git a/pyproject.toml b/pyproject.toml
index 11ccebc..9ba1438 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,6 +10,8 @@ requires-python = ">=3.9"
 license = { file = "LICENSE.txt" }
 authors = [
 	{ name = "Abe Levitan", email = "abraham.levitan@psi.ch" },
+	{ name = "Dayne Y. Sasaki" },
+	{ name = "Damian Guenzing" },
 	{ name = "Madelyn Cain" },
 	{ name = "Anastasiia Kutakh" }
 ]