From b3fad62bbbf4df666a44f970f1f4120b26657562 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Sun, 19 Apr 2026 06:34:39 -0700
Subject: [PATCH 01/59] update to pretrain object plotting functions

---
 .../diffractive_imaging/object_models.py      | 55 ++++++++++++++++++-
 .../diffractive_imaging/ptychography_lite.py  |  2 +
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 1c6fdab90..29c12fbe8 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -22,6 +22,7 @@
     validate_tensor,
 )
 from quantem.core.visualization import show_2d
+from quantem.core.visualization.custom_normalizations import CustomNormalization
 from quantem.diffractive_imaging.constraints import BaseConstraints
 from quantem.diffractive_imaging.ptycho_utils import sum_patches
 
@@ -1056,6 +1057,7 @@ def pretrain(
         apply_constraints: bool = False,
         show: bool = True,
         device: str | None = None,  # allow overwriting of device
+        normalize_object_plotting: bool = True,
     ):
         if device is not None:
             self.to(device)
@@ -1091,6 +1093,7 @@ def pretrain(
             loss_fn=loss_fn,
             apply_constraints=apply_constraints,
             show=show,
+            normalize_object_plotting=normalize_object_plotting,
         )
         self._set_pretrained_weights(self.model)
 
@@ -1100,6 +1103,7 @@ def _pretrain(
         loss_fn: Callable,
         apply_constraints: bool = False,
         show: bool = False,
+        normalize_object_plotting: bool = True,
     ):
         """Pretrain the DIP model."""
         if self.pretrain_target is None:
@@ -1151,9 +1155,16 @@ def _pretrain(
             pbar.set_description(f"Iter {a0 + 1}/{num_iters}, Loss: {loss.item():.3e}, ")
 
         if show:
-            self.visualize_pretrain(output)
+            self.visualize_pretrain(
+                output,
+                normalize_object_plotting=normalize_object_plotting,
+            )
 
-    def visualize_pretrain(self, pred_obj: torch.Tensor):
+    def visualize_pretrain(
+        self,
+        pred_obj: torch.Tensor,
+        normalize_object_plotting: bool = True,
+    ):
         import matplotlib.gridspec as gridspec
 
         fig = plt.figure(figsize=(12, 6))
@@ -1190,6 +1201,31 @@ def visualize_pretrain(self, pred_obj: torch.Tensor):
         if target is None:
             raise ValueError("Model has not been pre-trained")
         if n_bot == 4:
+            norm_angle = None
+            norm_abs = None
+            if normalize_object_plotting:
+                target_mean_angle = target.mean(0).angle().cpu().detach().numpy()
+                target_mean_abs = target.mean(0).abs().cpu().detach().numpy()
+
+                target_norm_angle = CustomNormalization(
+                    interval_type="quantile",
+                    data=target_mean_angle,
+                )
+                norm_angle = {
+                    "interval_type": "manual",
+                    "vmin": target_norm_angle.vmin,
+                    "vmax": target_norm_angle.vmax,
+                }
+
+                target_norm_abs = CustomNormalization(
+                    interval_type="quantile",
+                    data=target_mean_abs,
+                )
+                norm_abs = {
+                    "interval_type": "manual",
+                    "vmin": target_norm_abs.vmin,
+                    "vmax": target_norm_abs.vmax,
+                }
             show_2d(
                 [
                     pred_obj.mean(0).angle().cpu().detach().numpy(),
@@ -1206,8 +1242,22 @@ def visualize_pretrain(self, pred_obj: torch.Tensor):
                 ],
                 cmap="magma",
                 cbar=True,
+                norm=[norm_angle, norm_angle, norm_abs, norm_abs],
             )
         else:
+            norm = None
+            if normalize_object_plotting:
+                target_mean = target.mean(0).cpu().detach().numpy()
+                target_norm = CustomNormalization(
+                    interval_type="quantile",
+                    data=target_mean,
+                )
+                norm = {
+                    "interval_type": "manual",
+                    "vmin": target_norm.vmin,
+                    "vmax": target_norm.vmax,
+                }
+
             show_2d(
                 [
                     pred_obj.mean(0).cpu().detach().numpy(),
@@ -1217,6 +1267,7 @@ def visualize_pretrain(self, pred_obj: torch.Tensor):
                 title=[f"Pred obj ({self.obj_type})", f"Target obj ({self.obj_type})"],
                 cmap="magma",
                 cbar=True,
+                norm=norm,
             )
         plt.suptitle(
             f"Final loss: {self._pretrain_losses[-1]:.3e} | Iters: {len(self._pretrain_losses)}",
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index e3837b834..0a79c4820 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -268,6 +268,7 @@ def from_ptycholite(
         pretrain_lr: float = 1e-3,
         pretrain_probe: bool = True,
         pretrain_object: bool = True,
+        normalize_object_plotting: bool = True,
         # model settings
         cnn_num_layers: int = 3,
         # logging/device
@@ -323,6 +324,7 @@ def from_ptycholite(
                     },
                     apply_constraints=False,
                     device=config.get("device"),
+                    normalize_object_plotting=normalize_object_plotting,
                 )
             if pretrain_probe:
                 probe_model.pretrain(

From 24f2e3dff8950ac4ebe657645047b4cd58a0d51c Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 22 Apr 2026 14:45:14 -0700
Subject: [PATCH 02/59] fixing type hints

---
 src/quantem/core/utils/array_funcs.py                | 8 ++++----
 src/quantem/diffractive_imaging/object_models.py     | 2 +-
 src/quantem/diffractive_imaging/ptycho_utils.py      | 8 ++++----
 src/quantem/diffractive_imaging/ptychography_lite.py | 2 +-
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/quantem/core/utils/array_funcs.py b/src/quantem/core/utils/array_funcs.py
index 44ddb26c9..f0436764e 100644
--- a/src/quantem/core/utils/array_funcs.py
+++ b/src/quantem/core/utils/array_funcs.py
@@ -10,7 +10,7 @@
 from quantem.core import config
 
 if TYPE_CHECKING:
-    import cupy as cp
+    import cupy as cp  # type: ignore
     import torch
 else:
     if config.get("has_cupy"):
@@ -299,10 +299,10 @@ def fftshift(a: ArrayLike, axes: tuple[int, ...] | int | None = None) -> ArrayLi
 
 
 @overload
-def as_type(a: np.ndarray, dtype: "type|str|torch.dtype") -> np.ndarray: ...
+def as_type(a: np.ndarray, dtype: "type|str|torch.dtype|np.dtype") -> np.ndarray: ...
 @overload
-def as_type(a: "torch.Tensor", dtype: "type|str|torch.dtype") -> "torch.Tensor": ...
-def as_type(a: ArrayLike, dtype: "type|str|torch.dtype") -> ArrayLike:
+def as_type(a: "torch.Tensor", dtype: "type|str|torch.dtype|np.dtype") -> "torch.Tensor": ...
+def as_type(a: ArrayLike, dtype: "type|str|torch.dtype|np.dtype") -> ArrayLike:
     """Cast the array to a specified type."""
     if config.get("has_torch"):
         if isinstance(a, torch.Tensor):
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 29c12fbe8..27854d86d 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1242,7 +1242,7 @@ def visualize_pretrain(
                 ],
                 cmap="magma",
                 cbar=True,
-                norm=[norm_angle, norm_angle, norm_abs, norm_abs],
+                norm=[norm_angle, norm_angle, norm_abs, norm_abs], # type:ignore 
             )
         else:
             norm = None
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index 113761172..023963ffc 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -136,7 +136,7 @@ def fourier_shift_expand(
     if af.is_complex(array):
         return shifted_array
     else:
-        return shifted_array.real
+        return shifted_array.real # type:ignore ## will be numeric so this should be safe 
 
 
 @overload
@@ -144,20 +144,20 @@ def fourier_translation_operator(
     positions: np.ndarray,
     shape: tuple,
     expand_dim: bool = True,
-    dtype: "str|torch.dtype|None" = None,
+    dtype: "str|torch.dtype|np.dtype|None" = None,
 ) -> np.ndarray: ...
 @overload
 def fourier_translation_operator(
     positions: "torch.Tensor",
     shape: tuple,
     expand_dim: bool = True,
-    dtype: "str|torch.dtype|None" = None,
+    dtype: "str|torch.dtype|np.dtype|None" = None,
 ) -> "torch.Tensor": ...
 def fourier_translation_operator(
     positions: ArrayLike,
     shape: tuple,
     expand_dim: bool = True,
-    dtype: "str|torch.dtype|None" = None,
+    dtype: "str|torch.dtype|np.dtype|None" = None,
 ) -> ArrayLike:
     """Returns phase ramp for fourier-shifting array of shape `shape`."""
     nr, nc = shape[-2:]
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 0a79c4820..28e48e605 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -113,7 +113,7 @@ def from_dataset(
             probe_params.update(polar_parameters)
 
         if middle_focus:
-            if num_slices > 1:
+            if num_slices > 1 and obj_model.slice_thicknesses is not None:
                 half_thickness = obj_model.slice_thicknesses.sum() / 2
                 if "C10" in probe_params and probe_params["C10"] is not None:
                     probe_params["C10"] -= half_thickness

From 2cd6d12bf2c9c9b4aef52254c84d885365733cce Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Thu, 23 Apr 2026 08:50:36 -0700
Subject: [PATCH 03/59] adding final activation

---
 src/quantem/diffractive_imaging/ptychography_lite.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 28e48e605..2c0e3e32f 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -1,9 +1,10 @@
 import os
 from pathlib import Path
-from typing import Any, Literal, Self, Sequence
+from typing import Any, Callable, Literal, Self, Sequence
 
 import numpy as np
 import torch
+import torch.nn as nn
 
 from quantem.core import config
 from quantem.core.datastructures import Dataset4dstem
@@ -271,6 +272,7 @@ def from_ptycholite(
         normalize_object_plotting: bool = True,
         # model settings
         cnn_num_layers: int = 3,
+        final_activation: str | Callable = nn.Identity(),
         # logging/device
         log_dir: os.PathLike | str | None = None,
         log_prefix: str = "",
@@ -287,6 +289,7 @@ def from_ptycholite(
             out_channels=ptycholite.obj_model.num_slices,
             num_layers=cnn_num_layers,
             dtype=torch.complex64 if ptycholite.obj_model.obj_type == "complex" else torch.float32,
+            final_activation=final_activation,
         )
 
         obj_model = ObjectDIP.from_pixelated(

From e7b2ce023fabe87ff5a78a89f1cce4d24896222f Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Wed, 6 May 2026 09:40:10 -0700
Subject: [PATCH 04/59] position correction plotting and functionality for
 ptycholite

---
 .../diffractive_imaging/ptychography_lite.py  |  69 ++++++-
 .../ptychography_visualizations.py            | 187 ++++++++++++++++++
 2 files changed, 252 insertions(+), 4 deletions(-)

diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 2c0e3e32f..af9421c2d 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -91,6 +91,11 @@ def from_dataset(
                 f"dset must be Dataset4dstem or PtychographyDatasetRaster, got {type(dset)}"
             )
 
+        dset_model.learn_scan_positions = False
+        dset_model.learn_descan = False
+        dset_model.scan_positions_px.requires_grad_(False)
+        dset_model.descan_shifts.requires_grad_(False)
+
         if not dset_model.preprocessed:
             dset_model.preprocess(com_fit_function="constant")
 
@@ -167,6 +172,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         lr_obj: float = 5e-3,
         learn_probe: bool = True,
         lr_probe: float = 1e-3,
+        lr_scan_positions: float = 0.0,
         batch_size: int | None = None,
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
@@ -177,8 +183,26 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose
-
-        if new_optimizers or reset or self.num_iters == 0:
+        lr_scan_positions = float(lr_scan_positions)
+        if lr_scan_positions < 0:
+            raise ValueError(f"lr_scan_positions must be non-negative, got {lr_scan_positions}")
+
+        learn_scan_positions = lr_scan_positions > 0
+        self.dset.learn_scan_positions = learn_scan_positions
+        self.dset.learn_descan = False
+        self.dset.scan_positions_px.requires_grad_(learn_scan_positions)
+        self.dset.descan_shifts.requires_grad_(False)
+
+        needs_dataset_optimizer = learn_scan_positions
+        if not needs_dataset_optimizer and "dataset" in self.optimizers:
+            self.remove_optimizer("dataset")
+
+        if (
+            new_optimizers
+            or reset
+            or self.num_iters == 0
+            or (needs_dataset_optimizer and "dataset" not in self.optimizers)
+        ):
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -200,6 +224,15 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
                     "name": scheduler_type,
                     "factor": scheduler_factor,
                 }
+            if needs_dataset_optimizer:
+                opt_params["dataset"] = {
+                    "name": "adamw",
+                    "lr": lr_scan_positions,
+                }
+                scheduler_params["dataset"] = {
+                    "name": scheduler_type,
+                    "factor": scheduler_factor,
+                }
         else:
             opt_params = None
             scheduler_params = None
@@ -381,6 +414,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         lr_obj: float = 1e-3,
         learn_probe: bool = True,
         lr_probe: float = 1e-3,
+        lr_scan_positions: float = 0.0,
         batch_size: int | None = None,
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
@@ -391,8 +425,26 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose
-
-        if new_optimizers or reset or self.num_iters == 0:
+        lr_scan_positions = float(lr_scan_positions)
+        if lr_scan_positions < 0:
+            raise ValueError(f"lr_scan_positions must be non-negative, got {lr_scan_positions}")
+
+        learn_scan_positions = lr_scan_positions > 0
+        self.dset.learn_scan_positions = learn_scan_positions
+        self.dset.learn_descan = False
+        self.dset.scan_positions_px.requires_grad_(learn_scan_positions)
+        self.dset.descan_shifts.requires_grad_(False)
+
+        needs_dataset_optimizer = learn_scan_positions
+        if not needs_dataset_optimizer and "dataset" in self.optimizers:
+            self.remove_optimizer("dataset")
+
+        if (
+            new_optimizers
+            or reset
+            or self.num_iters == 0
+            or (needs_dataset_optimizer and "dataset" not in self.optimizers)
+        ):
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -414,6 +466,15 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
                     "name": scheduler_type,
                     "factor": scheduler_factor,
                 }
+            if needs_dataset_optimizer:
+                opt_params["dataset"] = {
+                    "name": "adamw",
+                    "lr": lr_scan_positions,
+                }
+                scheduler_params["dataset"] = {
+                    "name": scheduler_type,
+                    "factor": scheduler_factor,
+                }
         else:
             opt_params = None
             scheduler_params = None
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index 8cf2fa49d..41e769d74 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -4,10 +4,12 @@
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 from scipy.signal.windows import tukey
 
 from quantem.core import config
 from quantem.core.visualization import show_2d
+from quantem.diffractive_imaging.ptycho_utils import AffineTransform
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase, Snapshot
 
 
@@ -928,6 +930,191 @@ def show_scan_positions(
             )
         plt.show()
 
+    def show_updated_scan_positions(
+        self,
+        scan_positions_px: np.ndarray | None = None,
+        initial_scan_positions_px: np.ndarray | None = None,
+        scale_arrows: float = 1.0,
+        plot_arrow_freq: int | None = None,
+        plot_cropped_rotated_fov: bool = True,
+        cbar: bool = True,
+        verbose: bool = True,
+        return_fig: bool = False,
+        **kwargs,
+    ):
+        r"""Show changes to scan positions during ptychographic reconstruction.
+
+        The default plot compares the current learned scan positions against the
+        initial scan positions from preprocessing. Stored positions are in object
+        pixels; this visualization converts them to Å for the axis labels and
+        colorbar.
+
+        Parameters
+        ----------
+        scan_positions_px : np.ndarray | None, optional
+            Updated scan positions in object pixels. If None, uses
+            ``self.dset.scan_positions_px``.
+        initial_scan_positions_px : np.ndarray | None, optional
+            Reference scan positions in object pixels. If None, uses
+            ``self.dset.initial_scan_positions_px``.
+        scale_arrows : float, optional
+            Scaling factor applied to displacement vectors before plotting.
+            Default is 1.
+        plot_arrow_freq : int | None, optional
+            If provided, plot every Nth row and column of a raster scan grid.
+        plot_cropped_rotated_fov : bool, optional
+            If True, plot positions in the same cropped/rotated FOV convention as
+            ``show_obj``. Default is True.
+        cbar : bool, optional
+            Whether to show a colorbar for displacement magnitudes. Default is True.
+        verbose : bool, optional
+            Whether to print summary displacement statistics. Default is True.
+        return_fig : bool, optional
+            If True, return ``(fig, ax)`` instead of calling ``plt.show()``.
+        **kwargs
+            Additional keyword arguments passed to ``matplotlib.axes.Axes.quiver``.
+            ``figsize``, ``figax``, ``cmap``, and ``title`` are consumed by this
+            method.
+        """
+
+        if scan_positions_px is None:
+            scan_positions_px = self.dset.scan_positions_px
+        if initial_scan_positions_px is None:
+            initial_scan_positions_px = self.dset.initial_scan_positions_px
+
+        scan_positions_px = self._to_numpy(scan_positions_px).astype(float, copy=False)
+        initial_scan_positions_px = self._to_numpy(initial_scan_positions_px).astype(
+            float, copy=False
+        )
+
+        if scan_positions_px.ndim == 3:
+            scan_positions_px = scan_positions_px.mean(axis=0)
+        if initial_scan_positions_px.ndim == 3:
+            initial_scan_positions_px = initial_scan_positions_px.mean(axis=0)
+
+        if scan_positions_px.ndim != 2 or scan_positions_px.shape[-1] != 2:
+            raise ValueError(
+                "scan_positions_px must have shape (num_positions, 2), "
+                f"got {scan_positions_px.shape}"
+            )
+        if initial_scan_positions_px.ndim != 2 or initial_scan_positions_px.shape[-1] != 2:
+            raise ValueError(
+                "initial_scan_positions_px must have shape (num_positions, 2), "
+                f"got {initial_scan_positions_px.shape}"
+            )
+        if scan_positions_px.shape != initial_scan_positions_px.shape:
+            raise ValueError(
+                "scan_positions_px and initial_scan_positions_px must have the same shape, "
+                f"got {scan_positions_px.shape} and {initial_scan_positions_px.shape}"
+            )
+
+        sampling = np.asarray(self.sampling)
+        scan_positions = scan_positions_px * sampling
+        initial_scan_positions = initial_scan_positions_px * sampling
+
+        if plot_cropped_rotated_fov:
+            angle = (
+                self.dset.com_rotation_rad
+                if self.dset.com_transpose
+                else -self.dset.com_rotation_rad
+            )
+            tf = AffineTransform(angle=angle)
+            origin = initial_scan_positions.mean(axis=0)
+            initial_scan_positions = tf(initial_scan_positions, origin=origin)
+            scan_positions = tf(scan_positions, origin=origin)
+
+            obj_shape = self.obj_cropped.shape[-2:]
+            center_shift = initial_scan_positions.mean(axis=0) - (
+                np.array(obj_shape) / 2 * sampling
+            )
+            initial_scan_positions -= center_shift
+            scan_positions -= center_shift
+        else:
+            obj_shape = self.obj_shape_full[-2:]
+
+        if plot_arrow_freq is not None:
+            freq = int(plot_arrow_freq)
+            if freq <= 0:
+                raise ValueError(f"plot_arrow_freq must be positive, got {plot_arrow_freq}")
+
+            rshape = tuple(self.dset.gpts) + (2,)
+            if np.prod(self.dset.gpts) != initial_scan_positions.shape[0]:
+                raise ValueError(
+                    "plot_arrow_freq only supports full raster scan grids with "
+                    f"{np.prod(self.dset.gpts)} positions, got {initial_scan_positions.shape[0]}"
+                )
+
+            initial_scan_positions = initial_scan_positions.reshape(rshape)[
+                ::freq, ::freq
+            ].reshape(-1, 2)
+            scan_positions = scan_positions.reshape(rshape)[::freq, ::freq].reshape(-1, 2)
+
+        deltas = scan_positions - initial_scan_positions
+        norms = np.linalg.norm(deltas, axis=1)
+
+        if verbose:
+            print(
+                "Updated scan position shifts: "
+                f"mean={norms.mean():.4g} Å, "
+                f"rms={np.sqrt(np.mean(norms**2)):.4g} Å, "
+                f"max={norms.max():.4g} Å"
+            )
+
+        extent = [
+            0,
+            sampling[1] * obj_shape[1],
+            sampling[0] * obj_shape[0],
+            0,
+        ]
+
+        figsize = kwargs.pop("figsize", (4, 4))
+        figax = kwargs.pop("figax", None)
+        cmap = kwargs.pop("cmap", "Reds")
+        title = kwargs.pop("title", "Updated probe positions")
+
+        if figax is None:
+            fig, ax = plt.subplots(figsize=figsize)
+        else:
+            fig, ax = figax
+
+        quiver_kwargs = {
+            "angles": "xy",
+            "scale_units": "xy",
+            "scale": 1,
+            "cmap": cmap,
+        }
+        quiver_kwargs.update(kwargs)
+
+        im = ax.quiver(
+            initial_scan_positions[:, 1],
+            initial_scan_positions[:, 0],
+            deltas[:, 1] * scale_arrows,
+            deltas[:, 0] * scale_arrows,
+            norms,
+            **quiver_kwargs,
+        )
+
+        if cbar:
+            divider = make_axes_locatable(ax)
+            ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
+            fig.add_axes(ax_cb)
+            cb = fig.colorbar(im, cax=ax_cb)
+            cb.set_label("Δ [Å]", rotation=0, ha="left", va="bottom")
+            cb.ax.yaxis.set_label_coords(0.5, 1.01)
+
+        ax.set_ylabel("x [Å]")
+        ax.set_xlabel("y [Å]")
+        ax.set_xlim((extent[0], extent[1]))
+        ax.set_ylim((extent[2], extent[3]))
+        ax.set_aspect("equal")
+        ax.set_title(title)
+
+        if return_fig:
+            return fig, ax
+
+        plt.show()
+        return None
+
     def show_fourier_probe_and_amplitudes(
         self,
         probe: np.ndarray | None = None,

From 0c5fe6eae96c974bae9984c7b3429678e2554e1b Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Wed, 6 May 2026 09:57:22 -0700
Subject: [PATCH 05/59] small bug fix

---
 .../diffractive_imaging/ptychography_lite.py  | 26 +++++++++----------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index af9421c2d..3080e9f45 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -187,7 +187,11 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if lr_scan_positions < 0:
             raise ValueError(f"lr_scan_positions must be non-negative, got {lr_scan_positions}")
 
-        learn_scan_positions = lr_scan_positions > 0
+        setup_new_optimizers = new_optimizers or reset or self.num_iters == 0
+        has_dataset_optimizer = "dataset" in self.optimizers
+        learn_scan_positions = lr_scan_positions > 0 or (
+            not setup_new_optimizers and has_dataset_optimizer
+        )
         self.dset.learn_scan_positions = learn_scan_positions
         self.dset.learn_descan = False
         self.dset.scan_positions_px.requires_grad_(learn_scan_positions)
@@ -197,12 +201,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if not needs_dataset_optimizer and "dataset" in self.optimizers:
             self.remove_optimizer("dataset")
 
-        if (
-            new_optimizers
-            or reset
-            or self.num_iters == 0
-            or (needs_dataset_optimizer and "dataset" not in self.optimizers)
-        ):
+        if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -429,7 +428,11 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if lr_scan_positions < 0:
             raise ValueError(f"lr_scan_positions must be non-negative, got {lr_scan_positions}")
 
-        learn_scan_positions = lr_scan_positions > 0
+        setup_new_optimizers = new_optimizers or reset or self.num_iters == 0
+        has_dataset_optimizer = "dataset" in self.optimizers
+        learn_scan_positions = lr_scan_positions > 0 or (
+            not setup_new_optimizers and has_dataset_optimizer
+        )
         self.dset.learn_scan_positions = learn_scan_positions
         self.dset.learn_descan = False
         self.dset.scan_positions_px.requires_grad_(learn_scan_positions)
@@ -439,12 +442,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if not needs_dataset_optimizer and "dataset" in self.optimizers:
             self.remove_optimizer("dataset")
 
-        if (
-            new_optimizers
-            or reset
-            or self.num_iters == 0
-            or (needs_dataset_optimizer and "dataset" not in self.optimizers)
-        ):
+        if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
             opt_params = {
                 "object": {
                     "name": "adamw",

From b504b76b5330160b1798d5f748c0258785aaa91b Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Wed, 6 May 2026 10:17:03 -0700
Subject: [PATCH 06/59] one more tiny bug fix

---
 src/quantem/diffractive_imaging/ptychography_lite.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 3080e9f45..3cdb4a02f 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -95,6 +95,7 @@ def from_dataset(
         dset_model.learn_descan = False
         dset_model.scan_positions_px.requires_grad_(False)
         dset_model.descan_shifts.requires_grad_(False)
+        dset_model.remove_optimizer()
 
         if not dset_model.preprocessed:
             dset_model.preprocess(com_fit_function="constant")

From 5c539ba9ca704e99314486ebb6c07dcdf2d2ff5b Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Mon, 18 May 2026 17:44:32 -0700
Subject: [PATCH 07/59] multiprocessing working for ptycho multi gpu single
 node, linter errors to fix

---
 src/quantem/core/ml/dist_utils.py             |  64 +++++
 .../diffractive_imaging/dataset_models.py     |  54 ++++-
 .../diffractive_imaging/ptychography.py       | 218 ++++++++++++++++--
 .../diffractive_imaging/ptychography_base.py  |  13 ++
 .../diffractive_imaging/ptychography_lite.py  |   2 +-
 5 files changed, 327 insertions(+), 24 deletions(-)
 create mode 100644 src/quantem/core/ml/dist_utils.py

diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
new file mode 100644
index 000000000..136a9c536
--- /dev/null
+++ b/src/quantem/core/ml/dist_utils.py
@@ -0,0 +1,64 @@
+"""
+Standalone distributed training utilities for ptychography.
+
+These are kept separate from ddp.py (which imports tomography types) so they
+can be used by diffractive_imaging without circular imports.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.distributed as dist
+
+
+def is_distributed_launch() -> bool:
+    """True when launched via torchrun / torch.distributed.launch (RANK env var is set)."""
+    return "RANK" in os.environ
+
+
+def init_process_group(
+    rank: int,
+    world_size: int,
+    backend: str = "nccl",
+    master_addr: str = "127.0.0.1",
+    master_port: str = "29500",
+) -> None:
+    """Initialize the distributed process group from within an mp.spawn worker."""
+    os.environ["MASTER_ADDR"] = master_addr
+    os.environ["MASTER_PORT"] = master_port
+    dist.init_process_group(
+        backend=backend,
+        rank=rank,
+        world_size=world_size,
+    )
+    if backend == "nccl":
+        torch.cuda.set_device(rank)
+
+
+def get_rank() -> int:
+    """Return the current process rank (0 if not in a distributed context)."""
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_rank()
+    return 0
+
+
+def get_world_size() -> int:
+    """Return the world size (1 if not in a distributed context)."""
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_world_size()
+    return 1
+
+
+def all_reduce_params(*params: torch.Tensor, op: dist.ReduceOp = dist.ReduceOp.AVG) -> None:
+    """Average the .grad tensors of the given parameters across all ranks in-place."""
+    for p in params:
+        if p.grad is not None:
+            _ = dist.all_reduce(p.grad, op=op)  # type: ignore[arg-type]
+
+
+def broadcast_params(*params: torch.Tensor, src: int = 0) -> None:
+    """Broadcast .data of each parameter from rank src to all other ranks."""
+    for p in params:
+        _ = dist.broadcast(p.data, src=src)
\ No newline at end of file
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 4a4b9d5bb..5eb5552c9 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from math import ceil
 from pathlib import Path
 from typing import Any, Literal, Self
 
@@ -403,6 +404,8 @@ def roi_shape(self) -> np.ndarray:
 
     @property
     def num_gpts(self) -> int:
+        if hasattr(self, "_local_num_gpts") and self._local_num_gpts is not None:
+            return self._local_num_gpts
         return int(self.dset.shape[0])
 
     @property
@@ -536,7 +539,7 @@ def _set_patch_indices(self, obj_padding_px: np.ndarray | tuple) -> None:
             patch_indices_list.append(patch_indices_chunk)
 
         self._patch_indices = torch.cat(patch_indices_list, dim=0)
-        self._last_patch_positions_px = self.scan_positions_px.clone()
+        self._last_patch_positions_px = self.scan_positions_px.detach().clone()
 
     def patch_indices_need_update(self) -> bool:
         """
@@ -550,6 +553,55 @@ def reset(self) -> None:
         self.descan_shifts = self.initial_descan_shifts.clone().to(self.device)
         self.scan_positions_px = self.initial_scan_positions_px.clone().to(self.device)
 
+    def shard(self, rank: int, world_size: int) -> None:
+        """Partition diffraction data across DDP ranks (call after preprocess, before to(device)).
+
+        Each rank retains a contiguous slice [start:end] of the N scan positions. The object
+        and probe are not touched — they remain full-size and are replicated on every GPU.
+        After sharding, num_gpts returns the local shard size.
+        """
+        if not self._preprocessed:
+            raise RuntimeError("shard() must be called after preprocess()")
+        n_total = int(self.dset.shape[0])
+        shard_size = ceil(n_total / world_size)
+        start = rank * shard_size
+        end = min(start + shard_size, n_total)
+
+        self._shard_start: int = start
+        self._shard_end: int = end
+        self._global_num_gpts: int = n_total
+        self._local_num_gpts: int = end - start
+
+        sl = slice(start, end)
+
+        # Slice preprocessed amplitude/intensity data (plain attributes, not buffers)
+        for attr in (
+            "_amplitudes",
+            "_centered_amplitudes",
+            "_centered_intensities",
+            "_intensities",
+        ):
+            if hasattr(self, attr):
+                setattr(self, attr, getattr(self, attr)[sl].clone())
+
+        # Re-register buffers with rank-local slices
+        self.register_buffer("_patch_indices", self._patch_indices[sl].clone())
+        self.register_buffer("_last_patch_positions_px", self._last_patch_positions_px[sl].clone())
+        # _targets will be rebuilt from the sliced amplitudes in _set_targets(); reset it here
+        self.register_buffer("_targets", self._targets[sl].clone())
+
+        # Replace learnable parameters with rank-local slices
+        self._scan_positions_px = nn.Parameter(
+            self._scan_positions_px.data[sl].clone(),
+            requires_grad=self.learn_scan_positions,
+        )
+        self._descan_shifts = nn.Parameter(
+            self._descan_shifts.data[sl].clone(),
+            requires_grad=self.learn_descan,
+        )
+        self._initial_scan_positions_px = self._initial_scan_positions_px[sl].clone()
+        self._initial_descan_shifts = self._initial_descan_shifts[sl].clone()
+
     # endregion --- class methods ---
 
 
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 199e8588a..dd7dfcd7b 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -1,9 +1,10 @@
 import contextlib
 import copy
 import gc
+import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Self, Sequence, cast
+from typing import TYPE_CHECKING, Any, Literal, Self, Sequence, cast
 from warnings import warn
 
 import numpy as np
@@ -11,6 +12,10 @@
 
 from quantem.core import config
 from quantem.core.io.serialize import load as autoserialize_load
+from quantem.core.ml.dist_utils import (
+    init_process_group,
+    is_distributed_launch,
+)
 from quantem.diffractive_imaging.dataset_models import DatasetModelType
 from quantem.diffractive_imaging.detector_models import DetectorModelType
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
@@ -23,9 +28,53 @@
 
 if TYPE_CHECKING:
     import torch
+    import torch.distributed as dist
+    import torch.multiprocessing as mp
 else:
     if config.get("has_torch"):
         import torch
+        import torch.distributed as dist
+        import torch.multiprocessing as mp
+
+
+def _ddp_ptycho_worker(
+    rank: int,
+    world_size: int,
+    ptycho_path: str,
+    devices: list[int],
+    recon_kwargs: dict[str, Any],
+    result_path: str,
+) -> None:
+    """Module-level worker for mp.start_processes — must live at module scope to be picklable.
+
+    Receives a file path rather than the Ptychography object directly so that no
+    large tensors cross the process boundary via pickle (which triggers PyTorch's
+    shared-memory tensor mechanism and fails in some Linux environments).
+    """
+    device_id = devices[rank]
+    init_process_group(rank, world_size, backend="nccl" if torch.cuda.is_available() else "gloo")
+
+    ptycho = torch.load(ptycho_path, map_location="cpu", weights_only=False)
+    ptycho.dset.shard(rank, world_size)
+    ptycho.to(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
+
+    if dist.is_available() and dist.is_initialized():
+        ptycho._broadcast_parameters(src=0)
+
+    ptycho._reconstruct_inner(**recon_kwargs, _dist_rank=rank, _dist_world_size=world_size)
+
+    if rank == 0:
+        torch.save(
+            {
+                "obj": ptycho.obj_model._obj.data.cpu(),
+                "probe": ptycho.probe_model._probe.data.cpu(),
+                "iter_losses": ptycho._iter_losses,
+                "iter_val_losses": ptycho._iter_val_losses,
+            },
+            result_path,
+        )
+
+    dist.destroy_process_group()
 
 
 class Ptychography(PtychographyOpt, PtychographyVisualizations, PtychographyBase):
@@ -153,23 +202,98 @@ def reconstruct(
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
-        device: Literal["cpu", "gpu"] | None = None,
+        device: Literal["cpu", "gpu"] | int | list[int] | None = None,
         autograd: bool = True,
         loss_type: Literal[
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ] = "l2_amplitude",
     ) -> Self:
-        """
-        reason for having a single reconstruct() is so that updating things like constraints
-        or recon_types only happens in one place, reason for having separate reoconstruction_
-        methods would be to simplify the flags for this and not have to include all
+        """Run iterative ptychography reconstruction.
+
+        ``device`` accepts:
+          - ``None``             — keep current device
+          - ``"cpu"`` / ``"gpu"`` — existing string form
+          - ``int``              — specific GPU index, e.g. ``device=2`` → cuda:2
+          - ``list[int]``        — multi-GPU, e.g. ``device=[0,1,2,3]``
 
+        Multi-GPU (``device`` is a list) launches worker processes via ``mp.spawn`` when called
+        from a notebook, or uses the existing distributed process group when launched with
+        ``torchrun``. Only autograd mode is supported for multi-GPU in this release.
         """
-        # TODO maybe make an "process args" method that handles things like:
-        # mode, store_iterations, device,
         self._check_preprocessed()
-        if device is not None:
-            self.to(device)
+
+        # Route to multi-GPU path when a list of device IDs is given
+        if isinstance(device, list):
+            if not autograd:
+                raise ValueError("Multi-GPU reconstruction requires autograd=True.")
+            if not is_distributed_launch():
+                return self._spawn_reconstruct(
+                    devices=device,
+                    num_iters=num_iters,
+                    reset=reset,
+                    optimizer_params=optimizer_params,
+                    scheduler_params=scheduler_params,
+                    constraints=constraints,
+                    batch_size=batch_size,
+                    store_snapshots=store_snapshots,
+                    store_snapshots_every=store_snapshots_every,
+                    autograd=autograd,
+                    loss_type=loss_type,
+                )
+            # torchrun: fall through — process group already initialised externally
+
+        # Handle torchrun distributed launch (RANK env var present)
+        if is_distributed_launch():
+            rank = int(os.environ["RANK"])
+            world_size = int(os.environ["WORLD_SIZE"])
+            if not torch.distributed.is_initialized():
+                torch.distributed.init_process_group(
+                    backend="nccl" if torch.cuda.is_available() else "gloo",
+                    init_method="env://",
+                )
+            local_rank = int(os.environ.get("LOCAL_RANK", rank))
+            dev = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
+            self.to(dev)
+            self.dset.shard(rank, world_size)
+            self._broadcast_parameters(src=0)
+        else:
+            rank, world_size = 0, 1
+            if device is not None and not isinstance(device, list):
+                self.to(device)
+
+        return self._reconstruct_inner(
+            num_iters=num_iters,
+            reset=reset,
+            optimizer_params=optimizer_params,
+            scheduler_params=scheduler_params,
+            constraints=constraints,
+            batch_size=batch_size,
+            store_snapshots=store_snapshots,
+            store_snapshots_every=store_snapshots_every,
+            autograd=autograd,
+            loss_type=loss_type,
+            _dist_rank=rank,
+            _dist_world_size=world_size,
+        )
+
+    def _reconstruct_inner(
+        self,
+        num_iters: int = 0,
+        reset: bool = False,
+        optimizer_params: dict | None = None,
+        scheduler_params: dict | None = None,
+        constraints: dict = {},
+        batch_size: int | None = None,
+        store_snapshots: bool | None = None,
+        store_snapshots_every: int | None = None,
+        autograd: bool = True,
+        loss_type: Literal[
+            "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
+        ] = "l2_amplitude",
+        _dist_rank: int = 0,
+        _dist_world_size: int = 1,
+    ) -> Self:
+        """Core reconstruction loop. Called by reconstruct() for all launch modes."""
         self.batch_size = batch_size
         self.store_snapshot_every = store_snapshots_every
         if store_snapshots_every is not None and store_snapshots is None:
@@ -203,7 +327,7 @@ def reconstruct(
             val_ratio=self.val_ratio,
             val_mode=self.val_mode,
         )
-        pbar = tqdm(range(num_iters), disable=not self.verbose)
+        pbar = tqdm(range(num_iters), disable=not self.verbose or _dist_rank != 0)
 
         for a0 in pbar:
             consistency_loss = 0.0
@@ -240,6 +364,8 @@ def reconstruct(
                     patch_indices,
                     targets,
                 )
+                if _dist_world_size > 1:
+                    self._all_reduce_gradients()
                 self.step_optimizers()
                 consistency_loss += batch_consistency_loss.item()
                 total_loss += batch_loss.item()
@@ -248,6 +374,14 @@ def reconstruct(
             total_loss = total_loss / num_batches
             consistency_loss = consistency_loss / num_batches
 
+            # Average loss across ranks so rank-0 reports the global mean
+            if _dist_world_size > 1:
+                loss_t = torch.tensor(
+                    [total_loss, consistency_loss], device=self.device, dtype=torch.float64
+                )
+                dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
+                total_loss, consistency_loss = loss_t[0].item(), loss_t[1].item()
+
             # Validation pass (no gradient, no optimizer steps)
             val_loss = None
             if batcher.has_validation:
@@ -271,17 +405,19 @@ def reconstruct(
                         val_batches += 1
                 if val_batches > 0:
                     val_loss = val_consistency_loss / val_batches
-                    self._iter_val_losses.append(val_loss)
+                    if _dist_rank == 0:
+                        self._iter_val_losses.append(val_loss)
 
-            self._record_iter(total_loss)  # TODO record val loss as well
+            if _dist_rank == 0:
+                self._record_iter(total_loss)  # TODO record val loss as well
 
             # Step schedulers with current loss
             self.step_schedulers(total_loss)
 
-            if self.store_snapshots and (a0 % self.store_snapshot_every) == 0:
+            if _dist_rank == 0 and self.store_snapshots and (a0 % self.store_snapshot_every) == 0:
                 self._store_current_iter_snapshot()
 
-            if self.logger is not None:
+            if _dist_rank == 0 and self.logger is not None:
                 self.logger.log_iter(
                     self.obj_model,
                     self.probe_model,
@@ -292,21 +428,59 @@ def reconstruct(
                     self._get_current_lrs(),
                 )
 
-            if val_loss is not None:
-                pbar.set_description(
-                    f"Iter {a0 + 1}/{num_iters}, Loss: {total_loss:.3e}, Val: {val_loss:.3e}"
-                )
-            else:
-                pbar.set_description(f"Iter {a0 + 1}/{num_iters}, Loss: {total_loss:.3e}")
+            if _dist_rank == 0:
+                if val_loss is not None:
+                    pbar.set_description(
+                        f"Iter {a0 + 1}/{num_iters}, Loss: {total_loss:.3e}, Val: {val_loss:.3e}"
+                    )
+                else:
+                    pbar.set_description(f"Iter {a0 + 1}/{num_iters}, Loss: {total_loss:.3e}")
 
         gc.collect()
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         if hasattr(torch, "mps") and torch.backends.mps.is_available():
             torch.mps.empty_cache()
         gc.collect()
 
         return self
 
+    def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
+        """Notebook multi-GPU: spawn one worker process per device via forkserver.
+
+        State is saved to a temp file so that no tensors cross the process boundary
+        via pickle.  PyTorch's ForkingPickler automatically moves all CPU tensors to
+        shared memory when pickling for multiprocessing, which fails on some Linux
+        systems (EINVAL from ftruncate).  Passing only a file path (a plain string)
+        avoids that mechanism entirely.
+        """
+        self.to("cpu")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir_path = Path(tmpdir)
+            ptycho_path = str(tmpdir_path / "ptycho_state.pt")
+            result_path = str(tmpdir_path / "result.pt")
+
+            torch.save(self, ptycho_path, pickle_protocol=4)
+
+            # forkserver: workers fork from a clean pre-started server (no inherited
+            # CUDA, no Jupyter FDs).  Only plain Python scalars/strings cross the
+            # process boundary, so tensor pickling is never triggered.
+            mp.start_processes(
+                _ddp_ptycho_worker,
+                args=(len(devices), ptycho_path, devices, recon_kwargs, result_path),
+                nprocs=len(devices),
+                join=True,
+                start_method="forkserver",
+            )
+            result = torch.load(result_path, map_location="cpu", weights_only=False)
+
+        self.obj_model._obj.data.copy_(result["obj"])
+        self.probe_model._probe.data.copy_(result["probe"])
+        self._iter_losses.extend(result["iter_losses"])
+        self._iter_val_losses.extend(result["iter_val_losses"])
+        return self
+
     def _get_current_lrs(self) -> dict[str, float]:
         return {
             param_name: optimizer.param_groups[0]["lr"]
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 665819d9f..90f8e14ec 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -4,9 +4,11 @@
 import numpy as np
 import scipy.ndimage as ndi
 import torch
+import torch.distributed as dist
 
 from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
+from quantem.core.ml.dist_utils import all_reduce_params
 from quantem.core.utils.rng import RNGMixin
 from quantem.core.utils.utils import (
     electron_wavelength_angstrom,
@@ -876,6 +878,17 @@ def get_probe_intensities(
             intensities = np.abs(probe) ** 2
             return intensities.sum(axis=(-2, -1)) / intensities.sum()
 
+    def _broadcast_parameters(self, src: int = 0) -> None:
+        """Broadcast obj and probe data from rank src to all other ranks."""
+        dist.broadcast(self.obj_model._obj.data, src=src)
+        dist.broadcast(self.probe_model._probe.data, src=src)
+
+    def _all_reduce_gradients(self) -> None:
+        """Average obj.grad and probe.grad across all ranks (call after backward, before step)."""
+        params = [p for p in [self.obj_model._obj, self.probe_model._probe] if p.grad is not None]
+        if params:
+            all_reduce_params(*params)
+
     def to(self, device: str | int | torch.device):
         dev, _id = config.validate_device(device)
         if dev != self.device:
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 3cdb4a02f..9efec88f3 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -180,7 +180,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         new_optimizers: bool = False,  # not sure what the default should be
         constraints: dict = {},  # TODO add constraints flags
         store_iterations_every: int | None = None,
-        device: Literal["cpu", "gpu"] | None = None,
+        device: "Literal['cpu', 'gpu'] | int | list[int] | None" = None,
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose

From beb1dd96533deeebefcac74a6b2a860a574e97ef Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Tue, 19 May 2026 17:54:23 -0700
Subject: [PATCH 08/59] return fig options

---
 .../ptychography_visualizations.py            | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index 41e769d74..3470a5f36 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -356,6 +356,7 @@ def show_obj_slices(
         interval_type: Literal["quantile", "manual"] = "quantile",
         interval_scaling: Literal["each", "all"] = "each",
         max_width: int = 4,
+        return_fig: bool = False,
         **kwargs,
     ):
         """
@@ -372,13 +373,15 @@ def show_obj_slices(
             The interval scaling to use for the colorbar, by default "each"
         max_width: int, optional
             The maximum width of the object slices, by default 4
+        return_fig: bool, optional
+            If True, return ``(fig, axs)`` for saving or further customization.
         **kwargs: dict, optional
             Additional arguments passed to show_2d
 
         Returns
         -------
-        None
-            The object slices are shown in a new figure
+        tuple | None
+            ``(fig, axs)`` if return_fig is True, otherwise None.
         """
         if obj is None:
             obj = self.obj_cropped
@@ -426,14 +429,18 @@ def show_obj_slices(
         else:
             raise ValueError(f"Unknown interval type: {interval_type}")
 
-        show_2d(
+        fig, axs = show_2d(
             objs,
             title=titles,
             cmap=config.get("viz.phase_cmap"),
             norm=norm,
             cbar=cbar,
             scalebar=scalebars,
+            **kwargs,
         )
+        if return_fig:
+            return fig, axs
+        return None
 
     def plot_losses(self, figax: tuple | None = None, plot_lrs: bool = True):
         """
@@ -536,17 +543,20 @@ def plot_losses(self, figax: tuple | None = None, plot_lrs: bool = True):
             plt.tight_layout()
             plt.show()
 
-    def visualize(self, cbar: bool = True):
+    def visualize(self, cbar: bool = True, return_fig: bool = False):
         """
         Plot losses and show object and probe.
         Parameters
         ----------
         cbar: bool, optional
             Whether to show a colorbar, by default True
+        return_fig: bool, optional
+            If True, return ``(fig, axs)`` instead of calling ``plt.show()``.
 
         Returns
         -------
-        None
+        tuple | None
+            ``(fig, axs)`` if return_fig is True, otherwise None.
         """
         fig = plt.figure(figsize=(12, 6))
         gs = gridspec.GridSpec(2, 1, height_ratios=[1, 2], hspace=0.3)
@@ -562,7 +572,10 @@ def visualize(self, cbar: bool = True):
             fontsize=14,
             y=0.95,
         )
+        if return_fig:
+            return fig, (ax_top, axs_bot)
         plt.show()
+        return None
 
     def show_iters(
         self,

From 0e7396aa1b3501c8cc4013a1bc89bb90c38c41e4 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 21 May 2026 15:39:10 -0700
Subject: [PATCH 09/59] getting working for DGP

---
 .../diffractive_imaging/dataset_models.py     |  1 -
 .../diffractive_imaging/ptychography.py       | 22 +++++++++++----
 .../diffractive_imaging/ptychography_base.py  | 28 ++++++++++++++-----
 3 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index b02e1f683..5f7799bfa 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -569,7 +569,6 @@ def shard(self, rank: int, world_size: int) -> None:
 
         self._shard_start: int = start
         self._shard_end: int = end
-        self._global_num_gpts: int = n_total
         self._local_num_gpts: int = end - start
 
         sl = slice(start, end)
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index dd7dfcd7b..31fac8f68 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -66,8 +66,8 @@ def _ddp_ptycho_worker(
     if rank == 0:
         torch.save(
             {
-                "obj": ptycho.obj_model._obj.data.cpu(),
-                "probe": ptycho.probe_model._probe.data.cpu(),
+                "obj_state": {k: v.cpu() for k, v in ptycho.obj_model.state_dict().items()},
+                "probe_state": {k: v.cpu() for k, v in ptycho.probe_model.state_dict().items()},
                 "iter_losses": ptycho._iter_losses,
                 "iter_val_losses": ptycho._iter_val_losses,
             },
@@ -320,6 +320,17 @@ def _reconstruct_inner(
 
         self.dset._set_targets(loss_type)
         self.compute_propagator_arrays()  # required to avoid issue if stopped learning probe tilt
+
+        # Compute the global scan count once — needed to keep loss scale consistent across world sizes.
+        # In single-GPU mode num_gpts is already global; in distributed mode each rank holds a shard,
+        # so we sum across ranks to get the exact total.
+        if _dist_world_size > 1 and dist.is_available() and dist.is_initialized():
+            n_local = torch.tensor(self.dset.num_gpts, device=self.device, dtype=torch.long)
+            dist.all_reduce(n_local, op=dist.ReduceOp.SUM)
+            global_n = int(n_local.item())
+        else:
+            global_n = self.dset.num_gpts
+
         batcher = SimpleBatcher(
             self.dset.num_gpts,
             self.batch_size,
@@ -350,6 +361,7 @@ def _reconstruct_inner(
                     pred_intensities,
                     batch_indices,
                     loss_type=loss_type,
+                    global_n=global_n,
                 )
 
                 batch_soft_constraint_loss = self._soft_constraints()
@@ -399,7 +411,7 @@ def _reconstruct_inner(
                         )
                         pred_intensities = self.detector_model.forward(overlap)
                         batch_val_loss, _ = self.error_estimate(
-                            pred_intensities, batch_indices, loss_type=loss_type
+                            pred_intensities, batch_indices, loss_type=loss_type, global_n=global_n
                         )
                         val_consistency_loss += batch_val_loss.item()
                         val_batches += 1
@@ -475,8 +487,8 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
             )
             result = torch.load(result_path, map_location="cpu", weights_only=False)
 
-        self.obj_model._obj.data.copy_(result["obj"])
-        self.probe_model._probe.data.copy_(result["probe"])
+        self.obj_model.load_state_dict(result["obj_state"])
+        self.probe_model.load_state_dict(result["probe_state"])
         self._iter_losses.extend(result["iter_losses"])
         self._iter_val_losses.extend(result["iter_val_losses"])
         return self
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 90f8e14ec..b1685afa3 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -879,13 +879,25 @@ def get_probe_intensities(
             return intensities.sum(axis=(-2, -1)) / intensities.sum()
 
     def _broadcast_parameters(self, src: int = 0) -> None:
-        """Broadcast obj and probe data from rank src to all other ranks."""
-        dist.broadcast(self.obj_model._obj.data, src=src)
-        dist.broadcast(self.probe_model._probe.data, src=src)
+        """Broadcast obj and probe parameters from rank src to all other ranks.
+
+        Uses .parameters() so it works for both pixelated and DIP/INR models.
+        """
+        for p in self.obj_model.parameters():
+            dist.broadcast(p.data, src=src)
+        for p in self.probe_model.parameters():
+            dist.broadcast(p.data, src=src)
 
     def _all_reduce_gradients(self) -> None:
-        """Average obj.grad and probe.grad across all ranks (call after backward, before step)."""
-        params = [p for p in [self.obj_model._obj, self.probe_model._probe] if p.grad is not None]
+        """Average obj and probe gradients across all ranks (call after backward, before step).
+
+        Uses .parameters() so it works for both pixelated and DIP/INR models.
+        """
+        params = [
+            p
+            for p in list(self.obj_model.parameters()) + list(self.probe_model.parameters())
+            if p.grad is not None
+        ]
         if params:
             all_reduce_params(*params)
 
@@ -927,6 +939,7 @@ def error_estimate(
         loss_type: Literal[
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ] = "l2_amplitude",
+        global_n: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         targets = self.dset.targets[batch_indices]
         if "amplitude" in loss_type:
@@ -935,10 +948,11 @@ def error_estimate(
             preds = pred_intensities
 
         diff = preds * self.dset.detector_mask - targets * self.dset.detector_mask
+        n = global_n if global_n is not None else self.dset.num_gpts
         if "l1" in loss_type:
-            error = torch.sum(torch.abs(diff)) / (diff.shape[0] / self.dset.num_gpts)
+            error = torch.sum(torch.abs(diff)) / (diff.shape[0] / n)
         elif "l2" in loss_type:
-            error = torch.sum(torch.abs(diff) ** 2) / (diff.shape[0] / self.dset.num_gpts)
+            error = torch.sum(torch.abs(diff) ** 2) / (diff.shape[0] / n)
         elif loss_type == "poisson":
             error = torch.sum(preds - targets * torch.log(preds + 1e-6))
         else:

From f5c0c737f717c6052e31fed890fba9cedf17f9f3 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 21 May 2026 17:17:26 -0700
Subject: [PATCH 10/59] cleaning up multi GPU DGP fixing bugs in lr and opt
 persistence

---
 .../diffractive_imaging/ptychography.py       | 66 ++++++++++++++++++-
 .../diffractive_imaging/ptychography_base.py  |  3 +
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 31fac8f68..cd7fb695d 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -64,12 +64,20 @@ def _ddp_ptycho_worker(
     ptycho._reconstruct_inner(**recon_kwargs, _dist_rank=rank, _dist_world_size=world_size)
 
     if rank == 0:
+        obj_opt = ptycho.optimizers.get("object")
+        probe_opt = ptycho.optimizers.get("probe")
         torch.save(
             {
                 "obj_state": {k: v.cpu() for k, v in ptycho.obj_model.state_dict().items()},
                 "probe_state": {k: v.cpu() for k, v in ptycho.probe_model.state_dict().items()},
+                "obj_optimizer_params": ptycho.obj_model._optimizer_params,
+                "probe_optimizer_params": ptycho.probe_model._optimizer_params,
+                "obj_optimizer_state": obj_opt.state_dict() if obj_opt is not None else None,
+                "probe_optimizer_state": probe_opt.state_dict() if probe_opt is not None else None,
                 "iter_losses": ptycho._iter_losses,
                 "iter_val_losses": ptycho._iter_val_losses,
+                "iter_lrs": ptycho._iter_lrs,
+                "iter_recon_types": ptycho._iter_recon_types,
             },
             result_path,
         )
@@ -466,6 +474,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
         systems (EINVAL from ftruncate).  Passing only a file path (a plain string)
         avoids that mechanism entirely.
         """
+        restore_device = f"cuda:{devices[0]}" if torch.cuda.is_available() else "cpu"
         self.to("cpu")
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -487,10 +496,63 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
             )
             result = torch.load(result_path, map_location="cpu", weights_only=False)
 
+        # --- model weights ---
         self.obj_model.load_state_dict(result["obj_state"])
         self.probe_model.load_state_dict(result["probe_state"])
-        self._iter_losses.extend(result["iter_losses"])
-        self._iter_val_losses.extend(result["iter_val_losses"])
+        self.to(restore_device)
+
+        # --- restore optimizer params (worker may have set/changed them) so that future
+        #     spawns (e.g. reset=True without optimizer_params) can re-init the optimizer ---
+        for model, key in (
+            (self.obj_model, "obj_optimizer_params"),
+            (self.probe_model, "probe_optimizer_params"),
+        ):
+            saved = result.get(key)
+            if saved is not None:
+                model._optimizer_params = saved
+
+        # Re-create optimizers on the restored device (main process never ran set_optimizers).
+        # set_optimizers() skips models whose _optimizer_params is NoneOptimizer.
+        self.set_optimizers()
+
+        # --- optimizer states (params and device must be set before loading) ---
+        for name, key in (("object", "obj_optimizer_state"), ("probe", "probe_optimizer_state")):
+            opt_state = result.get(key)
+            opt = self.optimizers.get(name)
+            if opt_state is not None and opt is not None:
+                opt.load_state_dict(opt_state)
+                # State tensors were saved on CPU; move them to restore_device
+                for state in opt.state.values():
+                    for k, v in state.items():
+                        if isinstance(v, torch.Tensor):
+                            state[k] = v.to(restore_device)
+
+        # --- iteration tracking ---
+        # When reset=True the worker ran reset_recon() internally, so its lists start from 0.
+        # When reset=False the worker inherited existing history, so its lists are [old...new...].
+        # n_before lets us take only the genuinely new tail in both cases.
+        n_before = len(self._iter_losses)
+        is_reset = recon_kwargs.get("reset", False)
+
+        if is_reset:
+            self._iter_losses.clear()
+            self._iter_val_losses.clear()
+            self._iter_lrs.clear()
+            self._iter_recon_types.clear()
+            self._iter_losses.extend(result["iter_losses"])
+            self._iter_val_losses.extend(result["iter_val_losses"])
+            for k, v in result.get("iter_lrs", {}).items():
+                self._iter_lrs[k] = list(v)
+            self._iter_recon_types.extend(result.get("iter_recon_types", []))
+        else:
+            self._iter_losses.extend(result["iter_losses"][n_before:])
+            self._iter_val_losses.extend(result["iter_val_losses"][n_before:])
+            for k, v in result.get("iter_lrs", {}).items():
+                if k not in self._iter_lrs:
+                    self._iter_lrs[k] = []
+                self._iter_lrs[k].extend(list(v)[n_before:])
+            self._iter_recon_types.extend(result.get("iter_recon_types", [])[n_before:])
+
         return self
 
     def _get_current_lrs(self) -> dict[str, float]:
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index b1685afa3..8fa3a7b7d 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -905,6 +905,9 @@ def to(self, device: str | int | torch.device):
         dev, _id = config.validate_device(device)
         if dev != self.device:
             self._device = dev
+        # Sync each sub-model's own device tracker so their reset() uses the correct device
+        self.obj_model.device = dev
+        self.probe_model.device = dev
         self.obj_model.to(dev)
         self.probe_model.to(dev)
         self.dset.to(dev)

From 34a09ce08cb9bf32d5a30e389ce355e5dc04531d Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 21 May 2026 18:46:59 -0700
Subject: [PATCH 11/59] consistent devices

---
 .../diffractive_imaging/object_models.py      |  5 ++-
 .../diffractive_imaging/probe_models.py       |  5 ++-
 .../diffractive_imaging/ptychography.py       | 44 +++++++++++--------
 .../diffractive_imaging/ptychography_base.py  | 14 +++---
 4 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 27854d86d..786e953f1 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1056,11 +1056,12 @@ def pretrain(
         loss_fn: Callable | str = "l2",
         apply_constraints: bool = False,
         show: bool = True,
-        device: str | None = None,  # allow overwriting of device
+        device: str | int | None = None,
         normalize_object_plotting: bool = True,
     ):
         if device is not None:
-            self.to(device)
+            dev, _ = config.validate_device(device)
+            self.to(dev)
 
         if optimizer_params is not None:
             self.set_optimizer(optimizer_params)
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index f47ea1f77..b110a3350 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -1364,10 +1364,11 @@ def pretrain(
         loss_fn: Callable | str = "l2",
         apply_constraints: bool = False,
         show: bool = True,
-        device: str | None = None,  # allow overwriting of device
+        device: str | int | None = None,
     ):
         if device is not None:
-            self.to(device)
+            dev, _ = config.validate_device(device)
+            self.to(dev)
 
         if optimizer_params is not None:
             self.set_optimizer(optimizer_params)
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index cd7fb695d..2bfaefe92 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -230,25 +230,26 @@ def reconstruct(
         """
         self._check_preprocessed()
 
-        # Route to multi-GPU path when a list of device IDs is given
-        if isinstance(device, list):
+        # Determine effective device list: explicit arg takes priority, else fall back to stored.
+        devices_to_use = device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
+
+        # Route to multi-GPU path
+        if isinstance(devices_to_use, list) and not is_distributed_launch():
             if not autograd:
                 raise ValueError("Multi-GPU reconstruction requires autograd=True.")
-            if not is_distributed_launch():
-                return self._spawn_reconstruct(
-                    devices=device,
-                    num_iters=num_iters,
-                    reset=reset,
-                    optimizer_params=optimizer_params,
-                    scheduler_params=scheduler_params,
-                    constraints=constraints,
-                    batch_size=batch_size,
-                    store_snapshots=store_snapshots,
-                    store_snapshots_every=store_snapshots_every,
-                    autograd=autograd,
-                    loss_type=loss_type,
-                )
-            # torchrun: fall through — process group already initialised externally
+            return self._spawn_reconstruct(
+                devices=devices_to_use,
+                num_iters=num_iters,
+                reset=reset,
+                optimizer_params=optimizer_params,
+                scheduler_params=scheduler_params,
+                constraints=constraints,
+                batch_size=batch_size,
+                store_snapshots=store_snapshots,
+                store_snapshots_every=store_snapshots_every,
+                autograd=autograd,
+                loss_type=loss_type,
+            )
 
         # Handle torchrun distributed launch (RANK env var present)
         if is_distributed_launch():
@@ -302,7 +303,8 @@ def _reconstruct_inner(
         _dist_world_size: int = 1,
     ) -> Self:
         """Core reconstruction loop. Called by reconstruct() for all launch modes."""
-        self.batch_size = batch_size
+        if batch_size is not None:
+            self.batch_size = batch_size
         self.store_snapshot_every = store_snapshots_every
         if store_snapshots_every is not None and store_snapshots is None:
             self.store_snapshots = True
@@ -475,6 +477,11 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
         avoids that mechanism entirely.
         """
         restore_device = f"cuda:{devices[0]}" if torch.cuda.is_available() else "cpu"
+        # Persist batch_size on the main process so it carries into the saved file and
+        # is remembered on future calls that omit batch_size.
+        bs = recon_kwargs.get("batch_size")
+        if bs is not None:
+            self.batch_size = bs
         self.to("cpu")
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -553,6 +560,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
                 self._iter_lrs[k].extend(list(v)[n_before:])
             self._iter_recon_types.extend(result.get("iter_recon_types", [])[n_before:])
 
+        self._multi_gpu_devices = devices
         return self
 
     def _get_current_lrs(self) -> dict[str, float]:
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 8fa3a7b7d..183e4cb1e 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -101,6 +101,7 @@ def __init__(  # TODO prevent direct instantiation
         self.rng = rng
 
         # initializing default attributes
+        self._multi_gpu_devices: list[int] | None = None
         self._preprocessed: bool = False
         self._obj_padding_force_power2_level: int = 3
         self._store_snapshots: bool = False
@@ -597,12 +598,13 @@ def logger(self, logger: LoggerPtychography | None):
     # region --- implicit class properties ---
 
     @property
-    def device(self) -> str:
-        """This should be of form 'cuda:X' or 'cpu', as defined by quantem.config"""
+    def device(self) -> str | list[int]:
+        """Returns the active device: 'cuda:X'/'cpu' for single-GPU, or [gpu_ids] for multi-GPU."""
+        if getattr(self, "_multi_gpu_devices", None) is not None:
+            return self._multi_gpu_devices
         if hasattr(self, "_device"):
             return self._device
-        else:
-            return config.get("device")
+        return config.get("device")
 
     @device.setter
     def device(self, device: str | int | None):
@@ -903,8 +905,8 @@ def _all_reduce_gradients(self) -> None:
 
     def to(self, device: str | int | torch.device):
         dev, _id = config.validate_device(device)
-        if dev != self.device:
-            self._device = dev
+        self._device = dev
+        self._multi_gpu_devices = None
         # Sync each sub-model's own device tracker so their reset() uses the correct device
         self.obj_model.device = dev
         self.probe_model.device = dev

From f1417fe32c26cc942faf18f35b5ffcd32d069958 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 21 May 2026 19:40:14 -0700
Subject: [PATCH 12/59] fixing linter errors

---
 src/quantem/core/ml/dist_utils.py             |  7 +--
 .../diffractive_imaging/ptychography.py       | 54 +++++++++++--------
 .../diffractive_imaging/ptychography_base.py  | 39 +++++++++-----
 .../diffractive_imaging/ptychography_lite.py  | 30 +++++------
 .../diffractive_imaging/ptychography_opt.py   | 26 ++++-----
 5 files changed, 91 insertions(+), 65 deletions(-)

diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
index 136a9c536..0aeee4d0a 100644
--- a/src/quantem/core/ml/dist_utils.py
+++ b/src/quantem/core/ml/dist_utils.py
@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import os
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -51,14 +52,14 @@ def get_world_size() -> int:
     return 1
 
 
-def all_reduce_params(*params: torch.Tensor, op: dist.ReduceOp = dist.ReduceOp.AVG) -> None:
+def all_reduce_params(*params: torch.Tensor, op: Any = dist.ReduceOp.AVG) -> None:
     """Average the .grad tensors of the given parameters across all ranks in-place."""
     for p in params:
         if p.grad is not None:
-            _ = dist.all_reduce(p.grad, op=op)  # type: ignore[arg-type]
+            _ = dist.all_reduce(p.grad, op=op)
 
 
 def broadcast_params(*params: torch.Tensor, src: int = 0) -> None:
     """Broadcast .data of each parameter from rank src to all other ranks."""
     for p in params:
-        _ = dist.broadcast(p.data, src=src)
\ No newline at end of file
+        _ = dist.broadcast(p.data, src=src)
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 2bfaefe92..8b1e3e543 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -85,11 +85,14 @@ def _ddp_ptycho_worker(
     dist.destroy_process_group()
 
 
-class Ptychography(PtychographyOpt, PtychographyVisualizations, PtychographyBase):
+class Ptychography(PtychographyOpt, PtychographyVisualizations, PtychographyBase):  # pyright: ignore[reportUnsafeMultipleInheritance]
     """
     A class for performing phase retrieval using the Ptychography algorithm.
     """
 
+    _autograd: bool = True
+    _dataset_metadata: "dict[str, Any] | None" = None
+
     @classmethod
     def from_models(
         cls,
@@ -181,7 +184,7 @@ def _reset_iter_constraints(self) -> None:
 
     def _soft_constraints(self) -> torch.Tensor:
         """Calculate soft constraints by calling apply_soft_constraints on each model."""
-        total_loss = torch.tensor(0, device=self.device, dtype=self._dtype_real)
+        total_loss = torch.tensor(0, device=self._single_device, dtype=self._dtype_real)
 
         obj_loss = self.obj_model.apply_soft_constraints(
             self.obj_model.obj, mask=self.obj_model.mask
@@ -204,9 +207,9 @@ def reconstruct(
         self,
         num_iters: int = 0,
         reset: bool = False,
-        optimizer_params: dict | None = None,
-        scheduler_params: dict | None = None,
-        constraints: dict = {},
+        optimizer_params: dict[str, Any] | None = None,
+        scheduler_params: dict[str, Any] | None = None,
+        constraints: dict[str, Any] = {},
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
@@ -231,7 +234,9 @@ def reconstruct(
         self._check_preprocessed()
 
         # Determine effective device list: explicit arg takes priority, else fall back to stored.
-        devices_to_use = device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
+        devices_to_use = (
+            device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
+        )
 
         # Route to multi-GPU path
         if isinstance(devices_to_use, list) and not is_distributed_launch():
@@ -289,9 +294,9 @@ def _reconstruct_inner(
         self,
         num_iters: int = 0,
         reset: bool = False,
-        optimizer_params: dict | None = None,
-        scheduler_params: dict | None = None,
-        constraints: dict = {},
+        optimizer_params: dict[str, Any] | None = None,
+        scheduler_params: dict[str, Any] | None = None,
+        constraints: dict[str, Any] = {},
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
@@ -335,8 +340,10 @@ def _reconstruct_inner(
         # In single-GPU mode num_gpts is already global; in distributed mode each rank holds a shard,
         # so we sum across ranks to get the exact total.
         if _dist_world_size > 1 and dist.is_available() and dist.is_initialized():
-            n_local = torch.tensor(self.dset.num_gpts, device=self.device, dtype=torch.long)
-            dist.all_reduce(n_local, op=dist.ReduceOp.SUM)
+            n_local = torch.tensor(
+                self.dset.num_gpts, device=self._single_device, dtype=torch.long
+            )
+            _ = dist.all_reduce(n_local, op=dist.ReduceOp.SUM)  # type: ignore[call-overload]
             global_n = int(n_local.item())
         else:
             global_n = self.dset.num_gpts
@@ -399,7 +406,7 @@ def _reconstruct_inner(
             # Average loss across ranks so rank-0 reports the global mean
             if _dist_world_size > 1:
                 loss_t = torch.tensor(
-                    [total_loss, consistency_loss], device=self.device, dtype=torch.float64
+                    [total_loss, consistency_loss], device=self._single_device, dtype=torch.float64
                 )
                 dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
                 total_loss, consistency_loss = loss_t[0].item(), loss_t[1].item()
@@ -494,7 +501,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
             # forkserver: workers fork from a clean pre-started server (no inherited
             # CUDA, no Jupyter FDs).  Only plain Python scalars/strings cross the
             # process boundary, so tensor pickling is never triggered.
-            mp.start_processes(
+            mp.start_processes( # type: ignore 
                 _ddp_ptycho_worker,
                 args=(len(devices), ptycho_path, devices, recon_kwargs, result_path),
                 nprocs=len(devices),
@@ -585,12 +592,14 @@ def backward(
             # scaling pixelated ad gradients to closer match analytic
             if isinstance(self.obj_model, ObjectPixelated):
                 obj_grad_scale = self.dset.upsample_factor**2 / 2  # factor of 2 from l2 grad
-                self.obj_model._obj.grad.mul_(obj_grad_scale)  # type:ignore
+                if self.obj_model._obj.grad is not None:
+                    self.obj_model._obj.grad.mul_(obj_grad_scale)
 
             if isinstance(self.probe_model, ProbeParametric):
                 probe_grad_scale = np.sqrt(self.probe_model._mean_diffraction_intensity)
                 for par in self.probe_model.params:
-                    par.grad.mul_(probe_grad_scale)  # type:ignore
+                    if par.grad is not None:
+                        par.grad.mul_(probe_grad_scale)
 
         else:
             gradient = self.gradient_step(amplitudes, overlap)
@@ -713,7 +722,8 @@ def save(
         # Add other common skips for ptychography objects
         skips = skip
 
-        current_device = self.device
+        _dev = self.device
+        current_device: str = f"cuda:{_dev[0]}" if isinstance(_dev, list) else _dev
         self.to("cpu")
 
         if self.verbose and verbose:
@@ -730,8 +740,8 @@ def save(
         self.to(current_device)  # TODO figure out why this isn't working for DDIP sometimes?
 
         # Clean up temporary metadata
-        if not save_raw_data and hasattr(self, "_dataset_metadata"):
-            delattr(self, "_dataset_metadata")
+        if not save_raw_data and self._dataset_metadata is not None:
+            self._dataset_metadata = None
 
     @classmethod
     def from_file(
@@ -773,7 +783,7 @@ def from_file(
 
         # If no dataset was provided, try to reload it from saved metadata
         if dset is None and auto_reload_dataset and not hasattr(ptycho, "dset"):
-            if hasattr(ptycho, "_dataset_metadata") and ptycho._dataset_metadata:
+            if ptycho._dataset_metadata is not None:
                 metadata = ptycho._dataset_metadata
                 file_path = metadata.get("file_path")
 
@@ -816,13 +826,13 @@ def from_file(
         elif dset is not None:
             dset._set_initial_scan_positions_px(ptycho.obj_padding_px)
             dset._set_patch_indices(ptycho.obj_padding_px)
-            if hasattr(ptycho, "_dataset_metadata") and ptycho._dataset_metadata:
+            if ptycho._dataset_metadata is not None:
                 metadata = ptycho._dataset_metadata
                 # preserve learned scan positions and descan shifts
                 if "learned_scan_positions_px" in metadata:
-                    dset.scan_positions_px.data = metadata["learned_scan_positions_px"]
+                    dset.scan_positions_px.data = metadata["learned_scan_positions_px"]  # type: ignore[assignment]
                 if "learned_descan_shifts" in metadata:
-                    dset.descan_shifts.data = metadata["learned_descan_shifts"]
+                    dset.descan_shifts.data = metadata["learned_descan_shifts"]  # type: ignore[assignment]
 
         # check if dset was attached to ptycho object
         if dset is not None:
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 183e4cb1e..1f05b4bf4 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -95,6 +95,16 @@ def __init__(  # TODO prevent direct instantiation
             raise RuntimeError("the quantEM Ptychography module requires torch to be installed.")
 
         super().__init__()
+
+        # Pre-initialize private attributes so type checker sees them in __init__
+        self._verbose: int = 0
+        self._logger: LoggerPtychography | None = None
+        self._batch_size: int = 1
+        self._dset: DatasetModelType = dset
+        self._obj_model: ObjectModelType = obj_model
+        self._probe_model: ProbeModelType = probe_model
+        self._detector_model: DetectorModelType = detector_model
+
         self.verbose = verbose
         self.dset = dset
         self.device = device
@@ -109,7 +119,7 @@ def __init__(  # TODO prevent direct instantiation
         self._iter_losses: list[float] = []
         self._iter_val_losses: list[float] = []
         self._iter_recon_types: list[str] = []
-        self._iter_lrs: dict[str, list] = {}  # LRs/step_sizes across iterations
+        self._iter_lrs: dict[str, list[float]] = {}  # LRs/step_sizes across iterations
         self._snapshots: list[Snapshot] = []
         self._obj_padding_px = np.array([0, 0])
         self.obj_fov_mask = torch.ones(self.dset._obj_shape_full_2d(self.obj_padding_px).shape)
@@ -130,7 +140,7 @@ def __init__(  # TODO prevent direct instantiation
         self.detector_model = detector_model
         self.compute_propagator_arrays()
         self.logger = logger
-        self.to(self.device)
+        self.to(self._single_device)
 
     # region --- preprocessing ---
     ## hopefully will be able to remove some of thes preprocessing flags,
@@ -173,7 +183,7 @@ def preprocess(
                 self.roi_shape,
                 self.reciprocal_sampling,
                 self.dset.mean_diffraction_intensity,
-                device=self.device,
+                device=self._single_device,
             )
 
         # change obj_padding_px and whatever else needs to be changed
@@ -208,7 +218,7 @@ def _get_probe_overlap(self, max_batch_size: int | None = None) -> np.ndarray:
 
         batch_size = num_dps if max_batch_size is None else int(max_batch_size)
         probe_overlap = torch.zeros(
-            tuple(self.obj_shape_full[-2:]), dtype=self._dtype_real, device=self.device
+            tuple(self.obj_shape_full[-2:]), dtype=self._dtype_real, device=self._single_device
         )
         for start, end in generate_batches(num_dps, max_batch=batch_size):
             probe_overlap += sum_patches(
@@ -299,7 +309,7 @@ def slice_thicknesses(self) -> np.ndarray:
         return self._to_numpy(slice_thick)
 
     @slice_thicknesses.setter
-    def slice_thicknesses(self, val: float | Sequence | None) -> None:
+    def slice_thicknesses(self, val: float | Sequence[float] | None) -> None:
         self._obj_model.slice_thicknesses = val
         if hasattr(self, "_propagators"):  # propagators already set, update with new slices
             self.compute_propagator_arrays()
@@ -509,7 +519,7 @@ def obj_model(self, model: ObjectModelType | type):
             raise TypeError(f"obj_model must be a ObjectModelType, got {type(model)}")
 
         # Set object shape
-        model.to(self.device)
+        model.to(self._single_device)
         self._obj_model = cast(ObjectModelType, model)
 
     @property
@@ -530,12 +540,12 @@ def probe_model(self, model: ProbeModelType | type):
                 self.roi_shape,
                 self.reciprocal_sampling,
                 self.dset.mean_diffraction_intensity,
-                device=self.device,
+                device=self._single_device,
             )
         else:
             # will be set in ptycho.preprocess after dset is preprocessed
             pass
-        self._probe_model.to(self.device)
+        self._probe_model.to(self._single_device)
 
     @property
     def constraints(self) -> dict[str, Any]:
@@ -600,7 +610,7 @@ def logger(self, logger: LoggerPtychography | None):
     @property
     def device(self) -> str | list[int]:
         """Returns the active device: 'cuda:X'/'cpu' for single-GPU, or [gpu_ids] for multi-GPU."""
-        if getattr(self, "_multi_gpu_devices", None) is not None:
+        if self._multi_gpu_devices is not None:
             return self._multi_gpu_devices
         if hasattr(self, "_device"):
             return self._device
@@ -617,6 +627,11 @@ def device(self, device: str | int | None):
             except AttributeError:
                 pass
 
+    @property
+    def _single_device(self) -> str:
+        """Single-device string for internal tensor operations. Always str, never a list."""
+        return self._device if hasattr(self, "_device") else str(config.get("device"))
+
     @property
     def _obj_dtype(self) -> "torch.dtype":
         return self.obj_model.dtype
@@ -774,13 +789,13 @@ def _to_torch(
             raise TypeError(f"dtype should be string or torch.dtype, got {type(dtype)} {dtype}")
 
         if isinstance(array, np.ndarray):
-            t = torch.tensor(array.copy(), device=self.device, dtype=dt)
+            t = torch.tensor(array.copy(), device=self._single_device, dtype=dt)
         elif isinstance(array, torch.Tensor):
-            t = array.to(self.device)
+            t = array.to(self._single_device)
             if dt is not None:
                 t = t.type(dt)
         elif isinstance(array, (list, tuple)):
-            t = torch.tensor(array, device=self.device, dtype=dt)
+            t = torch.tensor(array, device=self._single_device, dtype=dt)
         else:
             raise TypeError(f"arr should be ndarray or Tensor, got {type(array)}")
         return t
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 9efec88f3..e399412c6 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -32,21 +32,21 @@ def from_dataset(
         *,
         # object settings
         num_slices: int = 1,
-        slice_thicknesses: float | Sequence | None = None,
+        slice_thicknesses: float | Sequence[float] | None = None,
         obj_type: Literal["complex", "pure_phase", "potential"] = "complex",
         # probe settings
         num_probes: int = 1,
         energy: float | None = None,
         defocus: float | None = None,
         semiangle_cutoff: float | None = None,
-        polar_parameters: dict | None = None,
+        polar_parameters: dict[str, Any] | None = None,
         middle_focus: bool = False,
         vacuum_probe_intensity: np.ndarray | Dataset4dstem | None = None,
         initial_probe_weights: list[float] | np.ndarray | None = None,
         # preprocessing
         obj_padding_px: tuple[int, int] = (0, 0),
         # logging/device
-        log_dir: os.PathLike | str | None = None,
+        log_dir: os.PathLike[str] | str | None = None,
         log_prefix: str = "",
         log_images_every: int = 10,
         log_probe_images: bool = False,
@@ -166,7 +166,7 @@ def from_dataset(
         )
         return ptycho
 
-    def reconstruct(  # type:ignore could do overloads but this is simpler...
+    def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         num_iters: int = 0,
         reset: bool = False,
@@ -178,7 +178,7 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
-        constraints: dict = {},  # TODO add constraints flags
+        constraints: dict[str, Any] = {},  # TODO add constraints flags
         store_iterations_every: int | None = None,
         device: "Literal['cpu', 'gpu'] | int | list[int] | None" = None,
         verbose: int | bool = True,
@@ -202,6 +202,8 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if not needs_dataset_optimizer and "dataset" in self.optimizers:
             self.remove_optimizer("dataset")
 
+        opt_params: dict[str, Any] | None
+        scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
             opt_params = {
                 "object": {
@@ -237,8 +239,6 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
             opt_params = None
             scheduler_params = None
 
-        constraints = constraints  # placeholder for constraints flags
-
         return super().reconstruct(
             num_iters=num_iters,
             reset=reset,
@@ -283,7 +283,7 @@ def from_file(
                 upgraded = cls._recursive_load_from_path(path)
                 return upgraded  # type: ignore[return-value]
 
-        return base  # type: ignore[return-value]
+        return base  # pyright: ignore[reportReturnType]
 
 
 class PtychoLiteDIP(Ptychography):
@@ -305,9 +305,9 @@ def from_ptycholite(
         normalize_object_plotting: bool = True,
         # model settings
         cnn_num_layers: int = 3,
-        final_activation: str | Callable = nn.Identity(),
+        final_activation: "str | Callable[..., Any]" = nn.Identity(),
         # logging/device
-        log_dir: os.PathLike | str | None = None,
+        log_dir: os.PathLike[str] | str | None = None,
         log_prefix: str = "",
         log_images_every: int = 10,
         log_probe_images: bool = False,
@@ -407,7 +407,7 @@ def from_ptycholite(
         )
         return ptycho
 
-    def reconstruct(  # type:ignore could do overloads but this is simpler...
+    def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         num_iters: int = 0,
         reset: bool = False,
@@ -419,9 +419,9 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
-        constraints: dict = {},  # TODO add constraints flags
+        constraints: dict[str, Any] = {},  # TODO add constraints flags
         store_iterations_every: int | None = None,
-        device: Literal["cpu", "gpu"] | None = None,
+        device: Literal["cpu", "gpu"] | int | list[int] | None = None,
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose
@@ -443,6 +443,8 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
         if not needs_dataset_optimizer and "dataset" in self.optimizers:
             self.remove_optimizer("dataset")
 
+        opt_params: dict[str, Any] | None
+        scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
             opt_params = {
                 "object": {
@@ -478,8 +480,6 @@ def reconstruct(  # type:ignore could do overloads but this is simpler...
             opt_params = None
             scheduler_params = None
 
-        constraints = constraints  # placeholder for constraints flags
-
         return super().reconstruct(
             num_iters=num_iters,
             reset=reset,
diff --git a/src/quantem/diffractive_imaging/ptychography_opt.py b/src/quantem/diffractive_imaging/ptychography_opt.py
index c18150e03..93eb1ff8a 100644
--- a/src/quantem/diffractive_imaging/ptychography_opt.py
+++ b/src/quantem/diffractive_imaging/ptychography_opt.py
@@ -1,5 +1,5 @@
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from quantem.core import config
 from quantem.core.ml.optimizer_mixin import (
@@ -54,7 +54,7 @@ def optimizer_params(self) -> dict[str, OptimizerType]:
         }
 
     @optimizer_params.setter
-    def optimizer_params(self, d: dict) -> None:
+    def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> None:
         """
         Takes a dictionary mapping optimizable keys to either an ``OptimizerType``
         dataclass or a plain dict (with optional ``"name"``/``"type"`` and ``"lr"``
@@ -67,10 +67,9 @@ def optimizer_params(self, d: dict) -> None:
         >>> ptycho.optimizer_params = {"object": {"name": "adam", "lr": 5e-3}}
         >>> ptycho.optimizer_params = ["object", "probe"]  # use all defaults
         """
-        if isinstance(d, (tuple, list)):
-            d = {k: {} for k in d}
+        _d: dict[str, Any] = {k: {} for k in d} if isinstance(d, (list, tuple)) else d
 
-        for k, v in d.items():
+        for k, v in _d.items():
             if isinstance(v, OptimizerType):
                 pass  # already a dataclass, pass through
             elif isinstance(v, dict):
@@ -85,11 +84,11 @@ def optimizer_params(self, d: dict) -> None:
                 raise TypeError(f"Expected OptimizerType or dict for key '{k}', got {type(v)}")
 
             if k == "object":
-                self.obj_model.optimizer_params = v
+                self.obj_model.optimizer_params = v  # type: ignore[assignment]
             elif k == "probe":
-                self.probe_model.optimizer_params = v
+                self.probe_model.optimizer_params = v  # type: ignore[assignment]
             elif k == "dataset":
-                self.dset.optimizer_params = v
+                self.dset.optimizer_params = v  # type: ignore[assignment]
             else:
                 raise ValueError(
                     f"key to be optimized, {k}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
@@ -140,7 +139,7 @@ def scheduler_params(self) -> dict[str, SchedulerType]:
         }
 
     @scheduler_params.setter
-    def scheduler_params(self, d: dict) -> None:
+    def scheduler_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> None:
         """
         Takes a dictionary mapping optimizable keys to either a ``SchedulerType``
         dataclass or a plain dict.  Keys not present in ``d`` are set to
@@ -151,10 +150,11 @@ def scheduler_params(self, d: dict) -> None:
         >>> ptycho.scheduler_params = {"object": SchedulerParams.Plateau(factor=0.5)}
         >>> ptycho.scheduler_params = {"object": {"name": "plateau", "factor": 0.5}}
         """
+        _d: dict[str, Any] = {k: {} for k in d} if isinstance(d, (list, tuple)) else dict(d)
         for key in self.OPTIMIZABLE_VALS:
-            if key not in d:
-                d[key] = SchedulerParams.NoneScheduler()
-        for k, v in d.items():
+            if key not in _d:
+                _d[key] = SchedulerParams.NoneScheduler()
+        for k, v in _d.items():
             if k == "object":
                 self.obj_model.scheduler_params = v
             elif k == "probe":
@@ -167,7 +167,7 @@ def scheduler_params(self, d: dict) -> None:
                 )
 
     @property
-    def schedulers(self) -> dict[str, "torch.optim.lr_scheduler._LRScheduler"]:
+    def schedulers(self) -> dict[str, "torch.optim.lr_scheduler.LRScheduler"]:
         """Get schedulers from all models."""
         schedulers = {}
         if self.obj_model.scheduler is not None:

From 578ace191f1d7759357f671358344836c3677aac Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Fri, 22 May 2026 09:55:02 -0700
Subject: [PATCH 13/59] fix for optimize hyperparameters

---
 .../optimize_hyperparameters.py               | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
index f7be892dc..88e212890 100644
--- a/src/quantem/diffractive_imaging/optimize_hyperparameters.py
+++ b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import gc
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Mapping, Optional
@@ -11,6 +12,7 @@
 from tqdm.auto import tqdm
 
 from quantem.core.visualization import show_2d
+from quantem.diffractive_imaging.dataset_models import PtychographyDatasetBase
 
 
 @dataclass
@@ -134,6 +136,7 @@ def _build_ptychography_instance(constructors, resolved_kwargs):
 
     init_kwargs = resolved_kwargs.get("init", {}).copy()
     init_kwargs["verbose"] = False
+    _isolate_trial_dataset(init_kwargs)
 
     return constructors["ptychography_class"](
         obj_model=obj_model,
@@ -147,10 +150,31 @@ def _build_ptycholite_instance(constructors, resolved_kwargs):
     """Build PtychoLite instance."""
     init_kwargs = resolved_kwargs.get("init", {}).copy()
     init_kwargs["verbose"] = False
+    _isolate_trial_dataset(init_kwargs)
 
     return constructors["ptychography_class"](**init_kwargs)
 
 
+def _isolate_trial_dataset(init_kwargs: dict[str, Any]) -> None:
+    """Give each optimization trial a private mutable dataset model."""
+    dset = init_kwargs.get("dset")
+    if not isinstance(dset, PtychographyDatasetBase):
+        return
+
+    try:
+        dset = copy.deepcopy(dset)
+    except Exception as exc:
+        raise RuntimeError(
+            "Could not copy the ptychography dataset for an optimization trial. "
+            "Pass a dataset_constructor instead so each trial can build a fresh dataset."
+        ) from exc
+
+    dset.reset()
+    dset.zero_grad(set_to_none=True)
+    dset.reset_optimizer()
+    init_kwargs["dset"] = dset
+
+
 def _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type):
     """Run the reconstruction pipeline for either class."""
     # Preprocess step

From 71840e08be2356a08ccc4df063d98949e01814a6 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 22 May 2026 18:20:12 -0700
Subject: [PATCH 14/59] converting ptycho to dataloader while maintaining
 jupyter notebook multi GPU compatibility

---
 src/quantem/core/ml/ddp.py                    |   5 +-
 src/quantem/core/ml/dist_utils.py             |  25 +++
 src/quantem/core/ml/optimizer_mixin.py        |   4 +
 .../diffractive_imaging/dataset_models.py     | 110 +++++------
 .../diffractive_imaging/object_models.py      |   2 +-
 .../diffractive_imaging/ptycho_utils.py       |  47 ++++-
 .../diffractive_imaging/ptychography.py       | 179 ++++++++++++++----
 .../diffractive_imaging/ptychography_base.py  |  29 ++-
 8 files changed, 286 insertions(+), 115 deletions(-)

diff --git a/src/quantem/core/ml/ddp.py b/src/quantem/core/ml/ddp.py
index dedd90dd2..9cc6e0f29 100644
--- a/src/quantem/core/ml/ddp.py
+++ b/src/quantem/core/ml/ddp.py
@@ -5,11 +5,10 @@
 import torch.nn as nn
 from torch.utils.data import DataLoader, Dataset, DistributedSampler, random_split
 
+from quantem.core.ml.dist_utils import worker_init_fn
 from quantem.tomography.dataset_models import DatasetModelType
 
-
-def worker_init_fn(worker_id):
-    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+__all__ = ["DDPMixin", "worker_init_fn"]
 
 
 class DDPMixin:
diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
index 0aeee4d0a..733e5be2e 100644
--- a/src/quantem/core/ml/dist_utils.py
+++ b/src/quantem/core/ml/dist_utils.py
@@ -63,3 +63,28 @@ def broadcast_params(*params: torch.Tensor, src: int = 0) -> None:
     """Broadcast .data of each parameter from rank src to all other ranks."""
     for p in params:
         _ = dist.broadcast(p.data, src=src)
+
+
+def worker_init_fn(worker_id: int) -> None:
+    """Hide CUDA from DataLoader workers so they only touch CPU-resident tensors."""
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+
+def spawn_distributed_workers(
+    worker_fn, devices: list[int], *worker_args, start_method: str = "forkserver"
+) -> None:
+    """Launch one worker per device via torch.multiprocessing.start_processes.
+
+    worker_fn must be a module-level callable with signature
+    (rank, world_size, *worker_args) — matches the mp.start_processes contract,
+    which passes rank as the first arg automatically.
+    """
+    import torch.multiprocessing as mp
+
+    mp.start_processes( # type: ignore 
+        worker_fn,
+        args=(len(devices), *worker_args),
+        nprocs=len(devices),
+        join=True,
+        start_method=start_method,
+    )
diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py
index 4ea263c07..9ade4b50a 100644
--- a/src/quantem/core/ml/optimizer_mixin.py
+++ b/src/quantem/core/ml/optimizer_mixin.py
@@ -612,6 +612,10 @@ def set_optimizer(self, opt_params: OptimizerType | dict | None = None) -> None:
         elif isinstance(params, Generator):
             params = list(params)
 
+        if not params:
+            self.remove_optimizer()
+            return
+
         # Ensure parameters require gradients
         for p in params:
             p.requires_grad_(True)
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 5f7799bfa..b000cf13e 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1,5 +1,4 @@
 from abc import abstractmethod
-from math import ceil
 from pathlib import Path
 from typing import Any, Literal, Self
 
@@ -7,6 +6,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.utils.data
 
 from quantem.core import config
 from quantem.core.datastructures.dataset3d import Dataset3d
@@ -29,7 +29,9 @@
 """
 
 
-class PtychographyDatasetBase(AutoSerialize, OptimizerMixin, torch.nn.Module):
+class PtychographyDatasetBase(
+    AutoSerialize, OptimizerMixin, torch.nn.Module, torch.utils.data.Dataset
+):
     _token = object()
     _patch_indices: torch.Tensor
 
@@ -63,6 +65,11 @@ def __init__(
 
         self.dset = dset
         self.verbose = verbose
+        # target_residency controls where loss targets live:
+        #   "device" (default) — targets are kept resident on the compute device (current behavior)
+        #   "cpu"              — targets live in CPU RAM and are streamed to the device per-batch,
+        #                        enabling datasets larger than a single GPU's VRAM
+        self.target_residency: Literal["device", "cpu"] = "device"
         self._preprocessed = False
         self._preprocessing_params = {}  # for serialization and reloading
         self._com_rotation_rad = 0  # default
@@ -86,7 +93,9 @@ def __init__(
         self._initial_scan_positions_px = torch.zeros_like(self._scan_positions_px)
         self._initial_descan_shifts = torch.zeros_like(self._descan_shifts)
 
-        self.register_buffer("_targets", torch.zeros(self.num_gpts, *self.roi_shape))
+        # _targets is a plain attribute (NOT a registered buffer) so that its device can be
+        # managed explicitly per target_residency; AutoSerialize does not serialize it either way.
+        self._targets = torch.zeros(self.num_gpts, *self.roi_shape)
         self.register_buffer(
             "_patch_indices", torch.zeros(self.num_gpts, *self.roi_shape, dtype=torch.int32)
         )
@@ -97,22 +106,31 @@ def __init__(
         self._probe_energy = None
 
     def get_optimization_parameters(self):
-        """Get the combined descan and scan position parameters for optimization."""
+        """Get the combined descan and scan position parameters for optimization.
+
+        Returns an empty list when neither learn flag is set; OptimizerMixin.set_optimizer
+        handles the empty-params case by removing the optimizer.
+        """
         params = []
         if self.learn_descan:
             params.append(self._descan_shifts)
         if self.learn_scan_positions:
             params.append(self._scan_positions_px)
-        if len(params) == 0:
-            raise RuntimeError(
-                "No parameters to optimize for dataset: learn_descan and learn_scan_positions are both False"
-            )
         return params
 
     def to(self, *args, **kwargs):
         """Move all relevant tensors to a different device."""
         # Call parent's to() method to handle PyTorch's internal device management
         super().to(*args, **kwargs)
+        # _targets is a plain attribute, so nn.Module.to() does not move it; do so explicitly
+        # unless residency is "cpu" (in which case targets intentionally stay on CPU and are
+        # streamed to the device per-batch).
+        if (
+            getattr(self, "target_residency", "device") != "cpu"
+            and getattr(self, "_targets", None) is not None
+        ):
+            # After super().to(), self.device reflects the new device (it reads off a Parameter).
+            self._targets = self._targets.to(self.device)
         # Reconnect optimizer to parameters on the new device
         self.reconnect_optimizer_to_parameters()
         return self
@@ -229,16 +247,19 @@ def _set_targets(
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ],
     ):
+        # When residency is "cpu", build targets on CPU so they can be streamed per-batch
+        # (and read by DataLoader workers); otherwise keep them resident on the compute device.
+        target_device = "cpu" if self.target_residency == "cpu" else self.device
         if "amplitude" in loss_type:
             if self.learn_descan and self.has_optimizer():
-                self._targets = self.amplitudes.clone().to(self.device)
+                self._targets = self.amplitudes.clone().to(target_device)
             else:
-                self._targets = self.centered_amplitudes.clone().to(self.device)
+                self._targets = self.centered_amplitudes.clone().to(target_device)
         elif "intensity" in loss_type or loss_type == "poisson":
             if self.learn_descan and self.has_optimizer():
-                self._targets = self.intensities.clone().to(self.device)
+                self._targets = self.intensities.clone().to(target_device)
             else:
-                self._targets = self.centered_intensities.clone().to(self.device)
+                self._targets = self.centered_intensities.clone().to(target_device)
         else:
             raise ValueError(f"Unknown loss type {loss_type}")
 
@@ -248,6 +269,21 @@ def patch_indices(self) -> torch.Tensor:
 
     # endregion --- buffers ---
 
+    # region --- torch.utils.data.Dataset interface ---
+    def __len__(self) -> int:
+        return self.num_gpts
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Return one sample for the DataLoader.
+
+        The target is returned as-is on whatever device target_residency dictates (CPU when
+        residency is "cpu" so DataLoader workers can read it). The integer index is collated
+        into a LongTensor by the default collate_fn.
+        """
+        return {"index": idx, "target": self._targets[idx]}
+
+    # endregion --- torch.utils.data.Dataset interface ---
+
     # region --- explicit properties (have setters) ---
     @property
     def dset(self) -> Dataset3d:
@@ -404,8 +440,6 @@ def roi_shape(self) -> np.ndarray:
 
     @property
     def num_gpts(self) -> int:
-        if hasattr(self, "_local_num_gpts") and self._local_num_gpts is not None:
-            return self._local_num_gpts
         return int(self.dset.shape[0])
 
     @property
@@ -553,54 +587,6 @@ def reset(self) -> None:
         self.descan_shifts = self.initial_descan_shifts.clone().to(self.device)
         self.scan_positions_px = self.initial_scan_positions_px.clone().to(self.device)
 
-    def shard(self, rank: int, world_size: int) -> None:
-        """Partition diffraction data across DDP ranks (call after preprocess, before to(device)).
-
-        Each rank retains a contiguous slice [start:end] of the N scan positions. The object
-        and probe are not touched — they remain full-size and are replicated on every GPU.
-        After sharding, num_gpts returns the local shard size.
-        """
-        if not self._preprocessed:
-            raise RuntimeError("shard() must be called after preprocess()")
-        n_total = int(self.dset.shape[0])
-        shard_size = ceil(n_total / world_size)
-        start = rank * shard_size
-        end = min(start + shard_size, n_total)
-
-        self._shard_start: int = start
-        self._shard_end: int = end
-        self._local_num_gpts: int = end - start
-
-        sl = slice(start, end)
-
-        # Slice preprocessed amplitude/intensity data (plain attributes, not buffers)
-        for attr in (
-            "_amplitudes",
-            "_centered_amplitudes",
-            "_centered_intensities",
-            "_intensities",
-        ):
-            if hasattr(self, attr):
-                setattr(self, attr, getattr(self, attr)[sl].clone())
-
-        # Re-register buffers with rank-local slices
-        self.register_buffer("_patch_indices", self._patch_indices[sl].clone())
-        self.register_buffer("_last_patch_positions_px", self._last_patch_positions_px[sl].clone())
-        # _targets will be rebuilt from the sliced amplitudes in _set_targets(); reset it here
-        self.register_buffer("_targets", self._targets[sl].clone())
-
-        # Replace learnable parameters with rank-local slices
-        self._scan_positions_px = nn.Parameter(
-            self._scan_positions_px.data[sl].clone(),
-            requires_grad=self.learn_scan_positions,
-        )
-        self._descan_shifts = nn.Parameter(
-            self._descan_shifts.data[sl].clone(),
-            requires_grad=self.learn_descan,
-        )
-        self._initial_scan_positions_px = self._initial_scan_positions_px[sl].clone()
-        self._initial_descan_shifts = self._initial_descan_shifts[sl].clone()
-
     # endregion --- class methods ---
 
 
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 786e953f1..ee4c672e8 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1243,7 +1243,7 @@ def visualize_pretrain(
                 ],
                 cmap="magma",
                 cbar=True,
-                norm=[norm_angle, norm_angle, norm_abs, norm_abs], # type:ignore 
+                norm=[norm_angle, norm_angle, norm_abs, norm_abs],  # type:ignore
             )
         else:
             norm = None
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index 023963ffc..f3c097261 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -117,6 +117,51 @@ def val_len(self) -> int:
         return int(ceil(len(self.val_indices) / self.batch_size)) if self.has_validation else 0
 
 
+def compute_train_val_split(
+    num: int,
+    val_ratio: float,
+    val_mode: Literal["grid", "random"],
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the train/validation index split.
+
+    Returns ``(train_indices, val_indices)`` as int numpy arrays. ``val_mode="grid"``
+    selects every k-th index (with ``k = round(1/val_ratio)``, inverted when
+    ``val_ratio > 0.5``); ``"random"`` selects a seeded ``rng.permutation`` slice.
+    """
+    indices = np.arange(num)
+    if val_ratio < 0 or val_ratio >= 1:
+        val_ratio = 0.0
+    n_val = int(round(len(indices) * val_ratio))
+    if n_val <= 0:
+        return indices, np.asarray([], dtype=int)
+
+    if val_mode == "random":
+        # Random unique selection for validation
+        perm = rng.permutation(indices)
+        val_indices = perm[:n_val]
+        train_indices = np.setdiff1d(indices, val_indices, assume_unique=False)
+    else:  # grid/regular selection: every k-th index
+        if val_ratio <= 0.5:
+            k = max(1, int(round(1.0 / val_ratio)))
+            invert = False
+        else:
+            k = max(1, int(round(1.0 / (1.0 - val_ratio))))
+            invert = True
+
+        grid_sel = indices[::k]
+        if len(grid_sel) > n_val:
+            grid_sel = grid_sel[:n_val]
+        if invert:
+            train_indices = grid_sel
+            val_indices = np.setdiff1d(indices, grid_sel, assume_unique=False)
+        else:
+            val_indices = grid_sel
+            train_indices = np.setdiff1d(indices, val_indices, assume_unique=False)
+
+    return np.asarray(train_indices, dtype=int), np.asarray(val_indices, dtype=int)
+
+
 @overload
 def fourier_shift_expand(
     array: np.ndarray, positions: np.ndarray, expand_dim: bool = True
@@ -136,7 +181,7 @@ def fourier_shift_expand(
     if af.is_complex(array):
         return shifted_array
     else:
-        return shifted_array.real # type:ignore ## will be numeric so this should be safe 
+        return shifted_array.real  # type:ignore ## will be numeric so this should be safe
 
 
 @overload
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 8b1e3e543..651d8b2a1 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -4,38 +4,32 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Self, Sequence, cast
+from typing import Any, Literal, Self, Sequence, cast
 from warnings import warn
 
 import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, DistributedSampler
 from tqdm.auto import tqdm
 
-from quantem.core import config
 from quantem.core.io.serialize import load as autoserialize_load
 from quantem.core.ml.dist_utils import (
     init_process_group,
     is_distributed_launch,
+    spawn_distributed_workers,
+    worker_init_fn,
 )
 from quantem.diffractive_imaging.dataset_models import DatasetModelType
 from quantem.diffractive_imaging.detector_models import DetectorModelType
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
 from quantem.diffractive_imaging.object_models import ObjectModelType, ObjectPixelated
 from quantem.diffractive_imaging.probe_models import ProbeModelType, ProbeParametric
-from quantem.diffractive_imaging.ptycho_utils import SimpleBatcher
+from quantem.diffractive_imaging.ptycho_utils import compute_train_val_split
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase
 from quantem.diffractive_imaging.ptychography_opt import PtychographyOpt
 from quantem.diffractive_imaging.ptychography_visualizations import PtychographyVisualizations
 
-if TYPE_CHECKING:
-    import torch
-    import torch.distributed as dist
-    import torch.multiprocessing as mp
-else:
-    if config.get("has_torch"):
-        import torch
-        import torch.distributed as dist
-        import torch.multiprocessing as mp
-
 
 def _ddp_ptycho_worker(
     rank: int,
@@ -54,8 +48,9 @@ def _ddp_ptycho_worker(
     device_id = devices[rank]
     init_process_group(rank, world_size, backend="nccl" if torch.cuda.is_available() else "gloo")
 
-    ptycho = torch.load(ptycho_path, map_location="cpu", weights_only=False)
-    ptycho.dset.shard(rank, world_size)
+    # mmap=True so all workers share one memory-mapped RAM copy of the (potentially large,
+    # CPU-resident) state instead of each duplicating it.
+    ptycho = torch.load(ptycho_path, map_location="cpu", weights_only=False, mmap=True)
     ptycho.to(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
 
     if dist.is_available() and dist.is_initialized():
@@ -218,6 +213,7 @@ def reconstruct(
         loss_type: Literal[
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ] = "l2_amplitude",
+        num_workers: int = 0,
     ) -> Self:
         """Run iterative ptychography reconstruction.
 
@@ -254,6 +250,7 @@ def reconstruct(
                 store_snapshots_every=store_snapshots_every,
                 autograd=autograd,
                 loss_type=loss_type,
+                num_workers=num_workers,
             )
 
         # Handle torchrun distributed launch (RANK env var present)
@@ -268,7 +265,6 @@ def reconstruct(
             local_rank = int(os.environ.get("LOCAL_RANK", rank))
             dev = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
             self.to(dev)
-            self.dset.shard(rank, world_size)
             self._broadcast_parameters(src=0)
         else:
             rank, world_size = 0, 1
@@ -286,10 +282,91 @@ def reconstruct(
             store_snapshots_every=store_snapshots_every,
             autograd=autograd,
             loss_type=loss_type,
+            num_workers=num_workers,
             _dist_rank=rank,
             _dist_world_size=world_size,
         )
 
+    def _build_dataloaders(
+        self,
+        train_indices: np.ndarray,
+        val_indices: np.ndarray,
+        world_size: int,
+        rank: int,
+        num_workers: int,
+    ) -> "tuple[DataLoader, DistributedSampler | None, DataLoader | None]":
+        """Build train + (optional) val DataLoaders for both single- and multi-GPU paths.
+
+        Mirrors the shape of ``DDPMixin.setup_dataloader`` but adapted to ptycho's device
+        contract (``str | list[int]``) and ptycho's precomputed ``val_mode`` index split.
+        ``world_size > 1`` uses ``DistributedSampler`` over a ``Subset``; ``world_size == 1``
+        uses ``shuffle=True`` with a seeded ``torch.Generator`` for run-to-run determinism.
+        ``__getitem__`` returns ``{"index": idx, ...}`` for the original dataset index, and
+        ``Subset[i]`` calls ``dataset[indices[i]]``, so ``batch["index"]`` is the original
+        dataset index under either branch.
+        """
+        pin_memory = self.dset.target_residency == "cpu" and str(self._single_device).startswith(
+            "cuda"
+        )
+        loader_kwargs: dict[str, Any] = {
+            "batch_size": self.batch_size,
+            "num_workers": num_workers,
+            "pin_memory": pin_memory,
+            "drop_last": False,
+        }
+        if num_workers > 0:
+            loader_kwargs.update(
+                multiprocessing_context="spawn",
+                persistent_workers=True,
+                worker_init_fn=worker_init_fn,
+            )
+
+        train_subset = torch.utils.data.Subset(self.dset, train_indices.tolist())
+        val_subset = (
+            torch.utils.data.Subset(self.dset, val_indices.tolist())
+            if len(val_indices) > 0
+            else None
+        )
+
+        if world_size > 1:
+            train_sampler = DistributedSampler(
+                train_subset,
+                num_replicas=world_size,
+                rank=rank,
+                shuffle=True,
+                seed=int(self.rng.integers(0, 2**31 - 1)),
+                drop_last=False,
+            )
+            train_loader = torch.utils.data.DataLoader(
+                train_subset, sampler=train_sampler, **loader_kwargs
+            )
+            if val_subset is not None:
+                val_sampler = DistributedSampler(
+                    val_subset,
+                    num_replicas=world_size,
+                    rank=rank,
+                    shuffle=False,
+                    drop_last=False,
+                )
+                val_loader = torch.utils.data.DataLoader(
+                    val_subset, sampler=val_sampler, **loader_kwargs
+                )
+            else:
+                val_loader = None
+        else:
+            train_sampler = None
+            shuffle_gen = torch.Generator().manual_seed(int(self.rng.integers(0, 2**31 - 1)))
+            train_loader = torch.utils.data.DataLoader(
+                train_subset, shuffle=True, generator=shuffle_gen, **loader_kwargs
+            )
+            val_loader = (
+                torch.utils.data.DataLoader(val_subset, shuffle=False, **loader_kwargs)
+                if val_subset is not None
+                else None
+            )
+
+        return train_loader, train_sampler, val_loader
+
     def _reconstruct_inner(
         self,
         num_iters: int = 0,
@@ -304,6 +381,7 @@ def _reconstruct_inner(
         loss_type: Literal[
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ] = "l2_amplitude",
+        num_workers: int = 0,
         _dist_rank: int = 0,
         _dist_world_size: int = 1,
     ) -> Self:
@@ -336,34 +414,36 @@ def _reconstruct_inner(
         self.dset._set_targets(loss_type)
         self.compute_propagator_arrays()  # required to avoid issue if stopped learning probe tilt
 
-        # Compute the global scan count once — needed to keep loss scale consistent across world sizes.
-        # In single-GPU mode num_gpts is already global; in distributed mode each rank holds a shard,
-        # so we sum across ranks to get the exact total.
-        if _dist_world_size > 1 and dist.is_available() and dist.is_initialized():
-            n_local = torch.tensor(
-                self.dset.num_gpts, device=self._single_device, dtype=torch.long
-            )
-            _ = dist.all_reduce(n_local, op=dist.ReduceOp.SUM)  # type: ignore[call-overload]
-            global_n = int(n_local.item())
-        else:
-            global_n = self.dset.num_gpts
+        # Compute the global scan count once — needed to keep loss scale consistent across world
+        global_n = self.dset.num_gpts
 
-        batcher = SimpleBatcher(
+        train_indices, val_indices = compute_train_val_split(
             self.dset.num_gpts,
-            self.batch_size,
-            rng=self.rng,
-            val_ratio=self.val_ratio,
-            val_mode=self.val_mode,
+            self.val_ratio,
+            self.val_mode,
+            self.rng,
         )
+        train_loader, train_sampler, val_loader = self._build_dataloaders(
+            train_indices,
+            val_indices,
+            world_size=_dist_world_size,
+            rank=_dist_rank,
+            num_workers=num_workers,
+        )
+
         pbar = tqdm(range(num_iters), disable=not self.verbose or _dist_rank != 0)
 
         for a0 in pbar:
+            if _dist_world_size > 1 and train_sampler is not None:
+                train_sampler.set_epoch(a0)
             consistency_loss = 0.0
             total_loss = 0.0
             self._reset_iter_constraints()
 
-            for batch_indices in batcher:
+            for batch in train_loader:
                 self.zero_grad_all()
+                batch_indices = batch["index"].to(self._single_device)
+                targets = batch["target"].to(self._single_device, non_blocking=True)
                 patch_indices, _positions_px, positions_px_fractional, descan_shifts = (
                     self.dset.forward(batch_indices, self.obj_padding_px)
                 )
@@ -377,6 +457,7 @@ def _reconstruct_inner(
                 batch_consistency_loss, targets = self.error_estimate(
                     pred_intensities,
                     batch_indices,
+                    targets=targets,
                     loss_type=loss_type,
                     global_n=global_n,
                 )
@@ -399,7 +480,7 @@ def _reconstruct_inner(
                 consistency_loss += batch_consistency_loss.item()
                 total_loss += batch_loss.item()
 
-            num_batches = len(batcher)
+            num_batches = len(train_loader)
             total_loss = total_loss / num_batches
             consistency_loss = consistency_loss / num_batches
 
@@ -413,11 +494,13 @@ def _reconstruct_inner(
 
             # Validation pass (no gradient, no optimizer steps)
             val_loss = None
-            if batcher.has_validation:
+            if val_loader is not None:
                 val_consistency_loss = 0.0
                 val_batches = 0
                 with torch.no_grad():
-                    for batch_indices in batcher.iter_val():
+                    for batch in val_loader:
+                        batch_indices = batch["index"].to(self._single_device)
+                        targets = batch["target"].to(self._single_device, non_blocking=True)
                         patch_indices, _positions_px, positions_px_fractional, descan_shifts = (
                             self.dset.forward(batch_indices, self.obj_padding_px)
                         )
@@ -428,12 +511,23 @@ def _reconstruct_inner(
                         )
                         pred_intensities = self.detector_model.forward(overlap)
                         batch_val_loss, _ = self.error_estimate(
-                            pred_intensities, batch_indices, loss_type=loss_type, global_n=global_n
+                            pred_intensities,
+                            batch_indices,
+                            targets=targets,
+                            loss_type=loss_type,
+                            global_n=global_n,
                         )
                         val_consistency_loss += batch_val_loss.item()
                         val_batches += 1
                 if val_batches > 0:
                     val_loss = val_consistency_loss / val_batches
+                    # Average the val loss across ranks so rank-0 records the global mean
+                    if _dist_world_size > 1:
+                        val_t = torch.tensor(
+                            val_loss, device=self._single_device, dtype=torch.float64
+                        )
+                        dist.all_reduce(val_t, op=dist.ReduceOp.AVG)
+                        val_loss = val_t.item()
                     if _dist_rank == 0:
                         self._iter_val_losses.append(val_loss)
 
@@ -501,12 +595,13 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
             # forkserver: workers fork from a clean pre-started server (no inherited
             # CUDA, no Jupyter FDs).  Only plain Python scalars/strings cross the
             # process boundary, so tensor pickling is never triggered.
-            mp.start_processes( # type: ignore 
+            spawn_distributed_workers(
                 _ddp_ptycho_worker,
-                args=(len(devices), ptycho_path, devices, recon_kwargs, result_path),
-                nprocs=len(devices),
-                join=True,
-                start_method="forkserver",
+                devices,
+                ptycho_path,
+                devices,
+                recon_kwargs,
+                result_path,
             )
             result = torch.load(result_path, map_location="cpu", weights_only=False)
 
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 1f05b4bf4..cd6ea9435 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -896,26 +896,43 @@ def get_probe_intensities(
             return intensities.sum(axis=(-2, -1)) / intensities.sum()
 
     def _broadcast_parameters(self, src: int = 0) -> None:
-        """Broadcast obj and probe parameters from rank src to all other ranks.
+        """Broadcast obj, probe, and dataset parameters from rank src to all other ranks.
 
-        Uses .parameters() so it works for both pixelated and DIP/INR models.
+        Uses .parameters() so it works for both pixelated and DIP/INR models. The dataset's
+        learnable params (scan positions / descan shifts) must also be broadcast: with the
+        DistributedSampler partitioning scan positions, the full position params are replicated
+        on every rank, so they must start identical and stay synchronized.
         """
         for p in self.obj_model.parameters():
             dist.broadcast(p.data, src=src)
         for p in self.probe_model.parameters():
             dist.broadcast(p.data, src=src)
+        for p in self.dset.get_optimization_parameters():
+            buf = p.data.contiguous()
+            dist.broadcast(buf, src=src)
+            p.data.copy_(buf)
 
     def _all_reduce_gradients(self) -> None:
-        """Average obj and probe gradients across all ranks (call after backward, before step).
+        """Average obj, probe, and dataset gradients across all ranks (call after backward,
+        before step).
 
-        Uses .parameters() so it works for both pixelated and DIP/INR models.
+        Uses .parameters() so it works for both pixelated and DIP/INR models. The dataset's
+        learnable params are included because each scan position's gradient is nonzero on
+        exactly one rank, so they must be reduced (AVG) to stay consistent across ranks.
         """
         params = [
             p
-            for p in list(self.obj_model.parameters()) + list(self.probe_model.parameters())
+            for p in (
+                list(self.obj_model.parameters())
+                + list(self.probe_model.parameters())
+                + list(self.dset.get_optimization_parameters())
+            )
             if p.grad is not None
         ]
         if params:
+            for p in params:
+                if p.grad is not None and not p.grad.is_contiguous():
+                    p.grad = p.grad.contiguous()
             all_reduce_params(*params)
 
     def to(self, device: str | int | torch.device):
@@ -956,12 +973,12 @@ def error_estimate(
         self,
         pred_intensities: torch.Tensor,
         batch_indices: np.ndarray,
+        targets: torch.Tensor,
         loss_type: Literal[
             "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
         ] = "l2_amplitude",
         global_n: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        targets = self.dset.targets[batch_indices]
         if "amplitude" in loss_type:
             preds = torch.sqrt(pred_intensities + 1e-9)  # add eps to avoid diverging gradients
         else:

From caa5c39d796c77910c0f6a37a6b345f720c059be Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 22 May 2026 18:26:43 -0700
Subject: [PATCH 15/59] move _build_dataloaders to ptychography_base

---
 .../diffractive_imaging/ptychography.py       | 82 ------------------
 .../diffractive_imaging/ptychography_base.py  | 83 ++++++++++++++++++-
 2 files changed, 82 insertions(+), 83 deletions(-)

diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 651d8b2a1..c6f8ffd73 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -10,7 +10,6 @@
 import numpy as np
 import torch
 import torch.distributed as dist
-from torch.utils.data import DataLoader, DistributedSampler
 from tqdm.auto import tqdm
 
 from quantem.core.io.serialize import load as autoserialize_load
@@ -18,7 +17,6 @@
     init_process_group,
     is_distributed_launch,
     spawn_distributed_workers,
-    worker_init_fn,
 )
 from quantem.diffractive_imaging.dataset_models import DatasetModelType
 from quantem.diffractive_imaging.detector_models import DetectorModelType
@@ -287,86 +285,6 @@ def reconstruct(
             _dist_world_size=world_size,
         )
 
-    def _build_dataloaders(
-        self,
-        train_indices: np.ndarray,
-        val_indices: np.ndarray,
-        world_size: int,
-        rank: int,
-        num_workers: int,
-    ) -> "tuple[DataLoader, DistributedSampler | None, DataLoader | None]":
-        """Build train + (optional) val DataLoaders for both single- and multi-GPU paths.
-
-        Mirrors the shape of ``DDPMixin.setup_dataloader`` but adapted to ptycho's device
-        contract (``str | list[int]``) and ptycho's precomputed ``val_mode`` index split.
-        ``world_size > 1`` uses ``DistributedSampler`` over a ``Subset``; ``world_size == 1``
-        uses ``shuffle=True`` with a seeded ``torch.Generator`` for run-to-run determinism.
-        ``__getitem__`` returns ``{"index": idx, ...}`` for the original dataset index, and
-        ``Subset[i]`` calls ``dataset[indices[i]]``, so ``batch["index"]`` is the original
-        dataset index under either branch.
-        """
-        pin_memory = self.dset.target_residency == "cpu" and str(self._single_device).startswith(
-            "cuda"
-        )
-        loader_kwargs: dict[str, Any] = {
-            "batch_size": self.batch_size,
-            "num_workers": num_workers,
-            "pin_memory": pin_memory,
-            "drop_last": False,
-        }
-        if num_workers > 0:
-            loader_kwargs.update(
-                multiprocessing_context="spawn",
-                persistent_workers=True,
-                worker_init_fn=worker_init_fn,
-            )
-
-        train_subset = torch.utils.data.Subset(self.dset, train_indices.tolist())
-        val_subset = (
-            torch.utils.data.Subset(self.dset, val_indices.tolist())
-            if len(val_indices) > 0
-            else None
-        )
-
-        if world_size > 1:
-            train_sampler = DistributedSampler(
-                train_subset,
-                num_replicas=world_size,
-                rank=rank,
-                shuffle=True,
-                seed=int(self.rng.integers(0, 2**31 - 1)),
-                drop_last=False,
-            )
-            train_loader = torch.utils.data.DataLoader(
-                train_subset, sampler=train_sampler, **loader_kwargs
-            )
-            if val_subset is not None:
-                val_sampler = DistributedSampler(
-                    val_subset,
-                    num_replicas=world_size,
-                    rank=rank,
-                    shuffle=False,
-                    drop_last=False,
-                )
-                val_loader = torch.utils.data.DataLoader(
-                    val_subset, sampler=val_sampler, **loader_kwargs
-                )
-            else:
-                val_loader = None
-        else:
-            train_sampler = None
-            shuffle_gen = torch.Generator().manual_seed(int(self.rng.integers(0, 2**31 - 1)))
-            train_loader = torch.utils.data.DataLoader(
-                train_subset, shuffle=True, generator=shuffle_gen, **loader_kwargs
-            )
-            val_loader = (
-                torch.utils.data.DataLoader(val_subset, shuffle=False, **loader_kwargs)
-                if val_subset is not None
-                else None
-            )
-
-        return train_loader, train_sampler, val_loader
-
     def _reconstruct_inner(
         self,
         num_iters: int = 0,
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index cd6ea9435..835487e2d 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -5,10 +5,11 @@
 import scipy.ndimage as ndi
 import torch
 import torch.distributed as dist
+from torch.utils.data import DataLoader, DistributedSampler
 
 from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
-from quantem.core.ml.dist_utils import all_reduce_params
+from quantem.core.ml.dist_utils import all_reduce_params, worker_init_fn
 from quantem.core.utils.rng import RNGMixin
 from quantem.core.utils.utils import (
     electron_wavelength_angstrom,
@@ -949,6 +950,86 @@ def to(self, device: str | int | torch.device):
         self._propagators = self._to_torch(self._propagators)
         self._rng_to_device(dev)
 
+    def _build_dataloaders(
+        self,
+        train_indices: np.ndarray,
+        val_indices: np.ndarray,
+        world_size: int,
+        rank: int,
+        num_workers: int,
+    ) -> "tuple[DataLoader, DistributedSampler | None, DataLoader | None]":
+        """Build train + (optional) val DataLoaders for both single- and multi-GPU paths.
+
+        Mirrors the shape of ``DDPMixin.setup_dataloader`` but adapted to ptycho's device
+        contract (``str | list[int]``) and ptycho's precomputed ``val_mode`` index split.
+        ``world_size > 1`` uses ``DistributedSampler`` over a ``Subset``; ``world_size == 1``
+        uses ``shuffle=True`` with a seeded ``torch.Generator`` for run-to-run determinism.
+        ``__getitem__`` returns ``{"index": idx, ...}`` for the original dataset index, and
+        ``Subset[i]`` calls ``dataset[indices[i]]``, so ``batch["index"]`` is the original
+        dataset index under either branch.
+        """
+        pin_memory = self.dset.target_residency == "cpu" and str(self._single_device).startswith(
+            "cuda"
+        )
+        loader_kwargs: dict[str, Any] = {
+            "batch_size": self.batch_size,
+            "num_workers": num_workers,
+            "pin_memory": pin_memory,
+            "drop_last": False,
+        }
+        if num_workers > 0:
+            loader_kwargs.update(
+                multiprocessing_context="spawn",
+                persistent_workers=True,
+                worker_init_fn=worker_init_fn,
+            )
+
+        train_subset = torch.utils.data.Subset(self.dset, train_indices.tolist())
+        val_subset = (
+            torch.utils.data.Subset(self.dset, val_indices.tolist())
+            if len(val_indices) > 0
+            else None
+        )
+
+        if world_size > 1:
+            train_sampler = DistributedSampler(
+                train_subset,
+                num_replicas=world_size,
+                rank=rank,
+                shuffle=True,
+                seed=int(self.rng.integers(0, 2**31 - 1)),
+                drop_last=False,
+            )
+            train_loader = DataLoader(
+                train_subset, sampler=train_sampler, **loader_kwargs
+            )
+            if val_subset is not None:
+                val_sampler = DistributedSampler(
+                    val_subset,
+                    num_replicas=world_size,
+                    rank=rank,
+                    shuffle=False,
+                    drop_last=False,
+                )
+                val_loader = DataLoader(
+                    val_subset, sampler=val_sampler, **loader_kwargs
+                )
+            else:
+                val_loader = None
+        else:
+            train_sampler = None
+            shuffle_gen = torch.Generator().manual_seed(int(self.rng.integers(0, 2**31 - 1)))
+            train_loader = DataLoader(
+                train_subset, shuffle=True, generator=shuffle_gen, **loader_kwargs
+            )
+            val_loader = (
+                DataLoader(val_subset, shuffle=False, **loader_kwargs)
+                if val_subset is not None
+                else None
+            )
+
+        return train_loader, train_sampler, val_loader
+
     # endregion
 
     # region --- ptychography foRcard model ---

From 5e36706d94bc9607e80c5eaa1e45a8f2ec75f3c7 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Sat, 23 May 2026 06:27:51 -0700
Subject: [PATCH 16/59] bug fix

---
 .../optimize_hyperparameters.py               | 75 ++++++++++++++++---
 1 file changed, 66 insertions(+), 9 deletions(-)

diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
index 88e212890..eb73c257e 100644
--- a/src/quantem/diffractive_imaging/optimize_hyperparameters.py
+++ b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
@@ -12,7 +12,10 @@
 from tqdm.auto import tqdm
 
 from quantem.core.visualization import show_2d
-from quantem.diffractive_imaging.dataset_models import PtychographyDatasetBase
+from quantem.diffractive_imaging.dataset_models import (
+    PtychographyDatasetBase,
+    PtychographyDatasetRaster,
+)
 
 
 @dataclass
@@ -161,20 +164,74 @@ def _isolate_trial_dataset(init_kwargs: dict[str, Any]) -> None:
     if not isinstance(dset, PtychographyDatasetBase):
         return
 
-    try:
-        dset = copy.deepcopy(dset)
-    except Exception as exc:
-        raise RuntimeError(
-            "Could not copy the ptychography dataset for an optimization trial. "
-            "Pass a dataset_constructor instead so each trial can build a fresh dataset."
-        ) from exc
-
+    dset = _clone_ptychography_dataset(dset)
     dset.reset()
     dset.zero_grad(set_to_none=True)
     dset.reset_optimizer()
     init_kwargs["dset"] = dset
 
 
+def _clone_ptychography_dataset(dset: PtychographyDatasetBase) -> PtychographyDatasetBase:
+    """Clone a ptychography dataset without using torch Module deepcopy."""
+    if not isinstance(dset, PtychographyDatasetRaster):
+        raise RuntimeError(
+            "Could not copy the ptychography dataset for an optimization trial. "
+            "Pass a dataset_constructor instead so each trial can build a fresh dataset."
+        )
+
+    detector_mask = dset.detector_mask.detach().cpu().clone()
+    origin = np.array([0, 0, *dset.dset.origin[-2:]])
+    sampling = np.array([*dset.scan_sampling, *dset.detector_sampling])
+    units = [*dset.scan_units, *dset.detector_units]
+    cloned = PtychographyDatasetRaster.from_array(
+        array=dset.intensities_4d.copy(),
+        name=dset.dset.name,
+        origin=origin,
+        sampling=sampling,
+        units=units,
+        signal_units=dset.dset.signal_units,
+        detector_mask=detector_mask,
+        verbose=dset.verbose,
+        learn_descan=dset.learn_descan,
+        learn_scan_positions=dset.learn_scan_positions,
+    )
+    cloned.constraints = copy.deepcopy(dset.constraints)
+    cloned._preprocessing_params = copy.deepcopy(dset._preprocessing_params)
+    cloned.com_rotation_rad = dset.com_rotation_rad
+    if hasattr(dset, "_transpose"):
+        cloned.com_transpose = dset.com_transpose
+    if dset.probe_energy is not None:
+        cloned.probe_energy = dset.probe_energy
+
+    if not dset.preprocessed:
+        return cloned
+
+    cloned.diffraction_padding = dset.diffraction_padding.copy()
+    cloned.com_measured = dset.com_measured.copy()
+    cloned.com_fit = dset.com_fit.copy()
+    cloned.centered_amplitudes = dset.centered_amplitudes.detach().cpu().clone()
+    cloned.amplitudes = dset.amplitudes.detach().cpu().clone()
+    cloned.centered_intensities = dset.centered_intensities.detach().cpu().clone()
+    cloned.intensities = dset.intensities.detach().cpu().clone()
+    cloned.detector_mask = detector_mask
+    cloned.mean_diffraction_intensity = dset.mean_diffraction_intensity
+    if hasattr(dset, "mean_diffraction_amplitude"):
+        cloned.mean_diffraction_amplitude = dset.mean_diffraction_amplitude
+    cloned._pattern_crop_mask = copy.deepcopy(getattr(dset, "_pattern_crop_mask", None))
+    cloned._pattern_crop_mask_shape = copy.deepcopy(
+        getattr(dset, "_pattern_crop_mask_shape", dset.roi_shape)
+    )
+    cloned.initial_descan_shifts = dset.initial_descan_shifts.detach().cpu().clone()
+    cloned.initial_scan_positions_px = dset.initial_scan_positions_px.detach().cpu().clone()
+    cloned.descan_shifts = cloned.initial_descan_shifts.clone()
+    cloned.scan_positions_px = cloned.initial_scan_positions_px.clone()
+    cloned._patch_indices = dset.patch_indices.detach().cpu().clone()
+    cloned._last_patch_positions_px = cloned.scan_positions_px.detach().clone()
+    cloned._targets = dset.targets.detach().cpu().clone()
+    cloned._preprocessed = True
+    return cloned
+
+
 def _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type):
     """Run the reconstruction pipeline for either class."""
     # Preprocess step

From 2ef8f9abab238d0a1c0d0c732905ad79125e7671 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Tue, 26 May 2026 14:06:49 -0700
Subject: [PATCH 17/59] cleaning up ptycho_opt

---
 .../diffractive_imaging/dataset_models.py     |  14 +-
 .../diffractive_imaging/ptychography_opt.py   | 209 +++++++-----------
 2 files changed, 99 insertions(+), 124 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index b000cf13e..835720908 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -69,7 +69,7 @@ def __init__(
         #   "device" (default) — targets are kept resident on the compute device (current behavior)
         #   "cpu"              — targets live in CPU RAM and are streamed to the device per-batch,
         #                        enabling datasets larger than a single GPU's VRAM
-        self.target_residency: Literal["device", "cpu"] = "device"
+        self._target_residency: Literal["device", "cpu"] = "device"
         self._preprocessed = False
         self._preprocessing_params = {}  # for serialization and reloading
         self._com_rotation_rad = 0  # default
@@ -241,6 +241,18 @@ def targets(self) -> torch.Tensor:
             raise ValueError("dset must be preprocessed before targets can be accessed")
         return self._targets
 
+    @property
+    def target_residency(self) -> Literal["device", "cpu"]:
+        """Where the loss targets live: ``"device"`` (resident, fastest) or
+        ``"cpu"`` (streamed per-batch, enables datasets larger than VRAM)."""
+        return self._target_residency
+
+    @target_residency.setter
+    def target_residency(self, value: str) -> None:
+        if value not in ("device", "cpu"):
+            raise ValueError(f"target_residency must be 'device' or 'cpu', got {value!r}")
+        self._target_residency = value
+
     def _set_targets(
         self,
         loss_type: Literal[
diff --git a/src/quantem/diffractive_imaging/ptychography_opt.py b/src/quantem/diffractive_imaging/ptychography_opt.py
index 93eb1ff8a..08bff4269 100644
--- a/src/quantem/diffractive_imaging/ptychography_opt.py
+++ b/src/quantem/diffractive_imaging/ptychography_opt.py
@@ -3,6 +3,7 @@
 
 from quantem.core import config
 from quantem.core.ml.optimizer_mixin import (
+    OptimizerMixin,
     OptimizerParams,
     OptimizerType,
     SchedulerParams,
@@ -19,7 +20,11 @@
 
 class PtychographyOpt(PtychographyBase):
     """
-    A class for performing phase retrieval using the Ptychography algorithm.
+    Optimizer/scheduler dispatch layer for `Ptychography`.
+
+    Each optimizable component (`object`, `probe`, `dataset`) lives on its own
+    `OptimizerMixin`-equipped model. The methods here are thin façades that fan a
+    single dict (`{key: params}`) out to the three models via the `_models` dict.
     """
 
     OPTIMIZABLE_VALS = ["object", "probe", "dataset"]
@@ -28,6 +33,25 @@ class PtychographyOpt(PtychographyBase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    @property
+    def _models(self) -> dict[str, OptimizerMixin]:
+        """Maps each optimization key to the model that owns its parameters.
+
+        Not cached: `obj_model`, `probe_model`, and `dset` can be reassigned
+        (e.g. by `from_ptychography`), and we must always see the current binding.
+        """
+        return {
+            "object": self.obj_model,
+            "probe": self.probe_model,
+            "dataset": self.dset,
+        }
+
+    def _check_key(self, key: str) -> None:
+        if key not in self.OPTIMIZABLE_VALS:
+            raise ValueError(
+                f"key to be optimized, {key}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
+            )
+
     def _get_default_lr(self, key: str) -> float:
         """Get default learning rate for a given optimization key."""
         if key == "object":
@@ -39,18 +63,14 @@ def _get_default_lr(self, key: str) -> float:
         else:
             raise ValueError(f"Unknown optimization key: {key}")
 
-    # region --- explicit properties and setters ---
+    # region --- optimizer params ---
 
     @property
     def optimizer_params(self) -> dict[str, OptimizerType]:
         return {
-            key: params
-            for key, params in [
-                ("object", self.obj_model.optimizer_params),
-                ("probe", self.probe_model.optimizer_params),
-                ("dataset", self.dset.optimizer_params),
-            ]
-            if not isinstance(params, OptimizerParams.NoneOptimizer)
+            key: model.optimizer_params
+            for key, model in self._models.items()
+            if not isinstance(model.optimizer_params, OptimizerParams.NoneOptimizer)
         }
 
     @optimizer_params.setter
@@ -58,7 +78,7 @@ def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
         """
         Takes a dictionary mapping optimizable keys to either an ``OptimizerType``
         dataclass or a plain dict (with optional ``"name"``/``"type"`` and ``"lr"``
-        keys).  Missing ``"name"`` / ``"lr"`` are filled from ``DEFAULT_OPTIMIZER_TYPE``
+        keys). Missing ``"name"`` / ``"lr"`` are filled from ``DEFAULT_OPTIMIZER_TYPE``
         and ``_get_default_lr`` respectively.
 
         Examples
@@ -70,6 +90,7 @@ def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
         _d: dict[str, Any] = {k: {} for k in d} if isinstance(d, (list, tuple)) else d
 
         for k, v in _d.items():
+            self._check_key(k)
             if isinstance(v, OptimizerType):
                 pass  # already a dataclass, pass through
             elif isinstance(v, dict):
@@ -83,66 +104,54 @@ def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
             else:
                 raise TypeError(f"Expected OptimizerType or dict for key '{k}', got {type(v)}")
 
-            if k == "object":
-                self.obj_model.optimizer_params = v  # type: ignore[assignment]
-            elif k == "probe":
-                self.probe_model.optimizer_params = v  # type: ignore[assignment]
-            elif k == "dataset":
-                self.dset.optimizer_params = v  # type: ignore[assignment]
-            else:
-                raise ValueError(
-                    f"key to be optimized, {k}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
-                )
+            self._models[k].optimizer_params = v  # type: ignore[assignment]
+
+    # endregion --- optimizer params ---
+
+    # region --- optimizers ---
 
     @property
     def optimizers(self) -> dict[str, "torch.optim.Optimizer"]:
-        """Get optimizers from all models."""
-        optimizers = {}
-        if self.obj_model.has_optimizer():
-            optimizers["object"] = self.obj_model.optimizer
-        if self.probe_model.has_optimizer():
-            optimizers["probe"] = self.probe_model.optimizer
-        if self.dset.has_optimizer():
-            optimizers["dataset"] = self.dset.optimizer
-        return optimizers
-
-    def set_optimizers(self):
-        """Set optimizers for each model."""
+        """Active optimizers, keyed by optimization key."""
+        return {
+            key: model.optimizer  # type: ignore[reportIncompatibleMethodOverride]
+            for key, model in self._models.items()
+            if model.has_optimizer()
+        }
+
+    def set_optimizers(self) -> None:
+        """(Re)create an optimizer on each model whose params are not `NoneOptimizer`."""
         for key, params in self.optimizer_params.items():
-            if key == "object":
-                self.obj_model.set_optimizer(params)
-            elif key == "probe":
-                self.probe_model.set_optimizer(params)
-            elif key == "dataset":
-                self.dset.set_optimizer(params)
-            else:
-                raise ValueError(
-                    f"key to be optimized, {key}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
-                )
+            self._models[key].set_optimizer(params)
 
     def remove_optimizer(self, key: str) -> None:
-        """Remove optimizer from a specific model."""
-        if key == "object":
-            self.obj_model.remove_optimizer()
-        elif key == "probe":
-            self.probe_model.remove_optimizer()
-        elif key == "dataset":
-            self.dset.remove_optimizer()
+        """Tear down the optimizer on the model for `key`."""
+        self._check_key(key)
+        self._models[key].remove_optimizer()
+
+    def step_optimizers(self) -> None:
+        for model in self._models.values():
+            if model.has_optimizer():
+                model.step_optimizer()
+
+    def zero_grad_all(self) -> None:
+        for model in self._models.values():
+            if model.has_optimizer():
+                model.zero_optimizer_grad()
+
+    # endregion --- optimizers ---
+
+    # region --- schedulers ---
 
     @property
     def scheduler_params(self) -> dict[str, SchedulerType]:
-        """Returns the parameters used to set the schedulers."""
-        return {
-            "object": self.obj_model.scheduler_params,
-            "probe": self.probe_model.scheduler_params,
-            "dataset": self.dset.scheduler_params,
-        }
+        return {key: model.scheduler_params for key, model in self._models.items()}
 
     @scheduler_params.setter
     def scheduler_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> None:
         """
         Takes a dictionary mapping optimizable keys to either a ``SchedulerType``
-        dataclass or a plain dict.  Keys not present in ``d`` are set to
+        dataclass or a plain dict. Keys not present in ``d`` are set to
         ``SchedulerParams.NoneScheduler()`` (disables scheduling for that model).
 
         Examples
@@ -152,75 +161,29 @@ def scheduler_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
         """
         _d: dict[str, Any] = {k: {} for k in d} if isinstance(d, (list, tuple)) else dict(d)
         for key in self.OPTIMIZABLE_VALS:
-            if key not in _d:
-                _d[key] = SchedulerParams.NoneScheduler()
+            _d.setdefault(key, SchedulerParams.NoneScheduler())
         for k, v in _d.items():
-            if k == "object":
-                self.obj_model.scheduler_params = v
-            elif k == "probe":
-                self.probe_model.scheduler_params = v
-            elif k == "dataset":
-                self.dset.scheduler_params = v
-            else:
-                raise ValueError(
-                    f"key to be optimized, {k}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
-                )
+            self._check_key(k)
+            self._models[k].scheduler_params = v
 
     @property
     def schedulers(self) -> dict[str, "torch.optim.lr_scheduler.LRScheduler"]:
-        """Get schedulers from all models."""
-        schedulers = {}
-        if self.obj_model.scheduler is not None:
-            schedulers["object"] = self.obj_model.scheduler
-        if self.probe_model.scheduler is not None:
-            schedulers["probe"] = self.probe_model.scheduler
-        if self.dset.scheduler is not None:
-            schedulers["dataset"] = self.dset.scheduler
-        return schedulers
-
-    def set_schedulers(self, params: dict[str, SchedulerType], num_iter: int | None = None):
-        """Set schedulers for each model."""
+        return {
+            key: model.scheduler
+            for key, model in self._models.items()
+            if model.scheduler is not None
+        }
+
+    def set_schedulers(
+        self, params: dict[str, SchedulerType], num_iter: int | None = None
+    ) -> None:
         for key, scheduler_params in params.items():
-            if key not in self.OPTIMIZABLE_VALS:
-                raise ValueError(
-                    f"key to be optimized, {key}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
-                )
-
-            if key == "object":
-                self.obj_model.set_scheduler(scheduler_params, num_iter)
-            elif key == "probe":
-                self.probe_model.set_scheduler(scheduler_params, num_iter)
-            elif key == "dataset":
-                self.dset.set_scheduler(scheduler_params, num_iter)
-
-    def step_optimizers(self):
-        """Step all active optimizers."""
-        for key in self.optimizer_params.keys():
-            if key == "object" and self.obj_model.has_optimizer():
-                self.obj_model.step_optimizer()
-            elif key == "probe" and self.probe_model.has_optimizer():
-                self.probe_model.step_optimizer()
-            elif key == "dataset" and self.dset.has_optimizer():
-                self.dset.step_optimizer()
-
-    def zero_grad_all(self):
-        """Zero gradients for all active optimizers."""
-        for key in self.optimizer_params.keys():
-            if key == "object" and self.obj_model.has_optimizer():
-                self.obj_model.zero_optimizer_grad()
-            elif key == "probe" and self.probe_model.has_optimizer():
-                self.probe_model.zero_optimizer_grad()
-            elif key == "dataset" and self.dset.has_optimizer():
-                self.dset.zero_optimizer_grad()
-
-    def step_schedulers(self, loss: float | None = None):
-        """Step all active schedulers."""
-        for key in self.scheduler_params.keys():
-            if key == "object" and self.obj_model.scheduler is not None:
-                self.obj_model.step_scheduler(loss)
-            elif key == "probe" and self.probe_model.scheduler is not None:
-                self.probe_model.step_scheduler(loss)
-            elif key == "dataset" and self.dset.scheduler is not None:
-                self.dset.step_scheduler(loss)
-
-    # endregion --- explicit properties and setters ---
+            self._check_key(key)
+            self._models[key].set_scheduler(scheduler_params, num_iter)
+
+    def step_schedulers(self, loss: float | None = None) -> None:
+        for model in self._models.values():
+            if model.scheduler is not None:
+                model.step_scheduler(loss)
+
+    # endregion --- schedulers ---

From 107ccd3b52687b1972bfa90ab6a37f74bc982596 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Tue, 26 May 2026 14:49:42 -0700
Subject: [PATCH 18/59] adding tests

---
 .../diffractive_imaging/probe_models.py       |   6 +-
 tests/diffractive_imaging/test_multi_gpu.py   | 241 ++++++++++++++++++
 .../diffractive_imaging/test_ptychography.py  | 122 ++++++++-
 3 files changed, 361 insertions(+), 8 deletions(-)
 create mode 100644 tests/diffractive_imaging/test_multi_gpu.py

diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index b110a3350..56bb89f6c 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -594,6 +594,8 @@ def from_array(
     ):
         if isinstance(probe_array, np.ndarray):
             probe_array = torch.tensor(probe_array, dtype=dtype, device=device)
+        else:
+            probe_array = probe_array.to(dtype=dtype, device=device)
         if probe_array.ndim == 3:
             if num_probes is None:
                 num_probes = probe_array.shape[0]
@@ -603,9 +605,7 @@ def from_array(
                 )
         else:
             num_probes = 1 if num_probes is None else num_probes
-            probe_array = torch.tensor(probe_array, dtype=dtype, device=device)
-            # probe_array = torch.tile(probe_array, (num_probes, 1, 1))
-            probe_array = torch.cat([probe_array] * num_probes, dim=0)
+            probe_array = torch.stack([probe_array] * num_probes, dim=0)
 
         probe_model = cls(
             num_probes=num_probes,
diff --git a/tests/diffractive_imaging/test_multi_gpu.py b/tests/diffractive_imaging/test_multi_gpu.py
new file mode 100644
index 000000000..1fb9b4dc5
--- /dev/null
+++ b/tests/diffractive_imaging/test_multi_gpu.py
@@ -0,0 +1,241 @@
+"""Multi-GPU state-management tests for iterative ptychography.
+
+Ported from the standalone verify_multi_gpu.py script. Tests are grouped by
+scenario (one test per related cluster of assertions) so that the spawn
+overhead is paid once per scenario rather than once per assertion.
+
+All tests are marked ``slow`` and skipped when fewer than 2 CUDA devices are
+available. Run with: ``uv run pytest tests/diffractive_imaging/test_multi_gpu.py --runslow``.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+# Helpers must live at module scope so forkserver-spawned DataLoader / DDP workers
+# can pickle and re-import them.
+
+N_SCAN = 8
+N_DET = 32
+PROBE_ENERGY = 80e3
+PROBE_SEMIANGLE = 20
+PROBE_DEFOCUS = 100
+N_ITERS = 4
+GPU_IDS = [0, 1]
+DEVICE_0 = "cuda:0"
+
+
+def _make_dataset():
+    from quantem.core.datastructures import Dataset4dstem
+
+    rng = np.random.default_rng(42)
+    array = rng.random((N_SCAN, N_SCAN, N_DET, N_DET)).astype(np.float32)
+    return Dataset4dstem.from_array(
+        array,
+        name="test",
+        sampling=[1.0, 1.0, 0.05, 0.05],
+        units=["A", "A", "A^-1", "A^-1"],
+    )
+
+
+def _make_ptycho():
+    from quantem.diffractive_imaging import (
+        DetectorPixelated,
+        ObjectPixelated,
+        ProbePixelated,
+        Ptychography,
+        PtychographyDatasetRaster,
+    )
+
+    pdset = PtychographyDatasetRaster.from_dataset4dstem(_make_dataset())
+    pdset.preprocess(com_fit_function="constant", plot_rotation=False, plot_com=False)
+    obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1)
+    probe = ProbePixelated.from_params(
+        probe_params={
+            "energy": PROBE_ENERGY,
+            "defocus": PROBE_DEFOCUS,
+            "semiangle_cutoff": PROBE_SEMIANGLE,
+        }
+    )
+    ptycho = Ptychography.from_models(
+        dset=pdset,
+        obj_model=obj,
+        probe_model=probe,
+        detector_model=DetectorPixelated(),
+        verbose=False,
+        rng=42,
+    )
+    ptycho.preprocess(obj_padding_px=(4, 4))
+    return ptycho
+
+
+def _make_dip_ptycho():
+    from quantem.core.ml import OptimizerParams
+    from quantem.diffractive_imaging import PtychoLite, PtychoLiteDIP
+
+    base = _make_ptycho()
+    base.reconstruct(
+        num_iters=5,
+        reset=True,
+        optimizer_params={
+            "object": OptimizerParams.Adam(lr=1e-2),
+            "probe": OptimizerParams.Adam(lr=1e-2),
+        },
+        batch_size=16,
+        device=0,
+    )
+    lite = PtychoLite.from_models(
+        dset=base.dset,
+        obj_model=base.obj_model,
+        probe_model=base.probe_model,
+        detector_model=base.detector_model,
+        verbose=False,
+        rng=42,
+    )
+    lite.preprocess(obj_padding_px=(4, 4))
+    return PtychoLiteDIP.from_ptycholite(lite, device="cpu", pretrain_iters=None)
+
+
+def _opt():
+    from quantem.core.ml import OptimizerParams
+
+    return {
+        "object": OptimizerParams.Adam(lr=1e-2),
+        "probe": OptimizerParams.Adam(lr=1e-2),
+    }
+
+
+# Module-level marks: all tests in this file are slow and require >= 2 GPUs.
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.device_count() < 2,
+        reason="requires >= 2 CUDA devices",
+    ),
+]
+
+
+class TestSingleGPUDevicePersistence:
+    """device= argument is sticky across follow-up reconstruct() calls."""
+
+    def test_int_device_persists(self):
+        p = _make_ptycho()
+        p.reconstruct(
+            num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=0
+        )
+        assert p.device == DEVICE_0
+        assert p.obj_model._obj.device.type == "cuda"
+
+        # follow-up call without device= keeps the previous device
+        p.reconstruct(num_iters=N_ITERS, batch_size=16)
+        assert p.device == DEVICE_0
+        assert p.obj_model._obj.device.type == "cuda"
+
+        # reset=True must not reset device tracking
+        p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16)
+        assert p.device == DEVICE_0
+
+
+class TestMultiGPUDeviceRestoration:
+    """device=[…] is stored and restored across spawn boundaries."""
+
+    def test_gpu_list_persists(self):
+        p = _make_ptycho()
+        p.reconstruct(
+            num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS
+        )
+        assert p.device == GPU_IDS
+        assert p.obj_model._obj.device.type == "cuda"
+
+        p.reconstruct(num_iters=N_ITERS, batch_size=16)
+        assert p.device == GPU_IDS
+
+
+class TestMultiGPULossTracking:
+    """_iter_losses extends correctly across spawn(reset=False)/reset=True."""
+
+    def test_losses_length_lifecycle(self):
+        p = _make_ptycho()
+        p.reconstruct(
+            num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS
+        )
+        assert len(p._iter_losses) == N_ITERS
+
+        p.reconstruct(num_iters=N_ITERS, reset=False, batch_size=16, device=GPU_IDS)
+        assert len(p._iter_losses) == 2 * N_ITERS, "continuation must not double-count"
+
+        p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16, device=GPU_IDS)
+        assert len(p._iter_losses) == N_ITERS
+
+
+class TestMultiGPULRTracking:
+    """iter_lrs extends correctly across spawn boundaries."""
+
+    def test_iter_lrs_lifecycle(self):
+        p = _make_ptycho()
+        p.reconstruct(
+            num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS
+        )
+        assert "object" in p.iter_lrs
+        assert len(p.iter_lrs["object"]) == N_ITERS
+
+        p.reconstruct(num_iters=N_ITERS, reset=False, batch_size=16, device=GPU_IDS)
+        assert len(p.iter_lrs["object"]) == 2 * N_ITERS
+
+        p.reconstruct(num_iters=N_ITERS, reset=True, batch_size=16, device=GPU_IDS)
+        assert len(p.iter_lrs["object"]) == N_ITERS
+
+
+class TestMultiGPUOptimizerState:
+    """Adam state survives the save/restore around the spawn worker."""
+
+    def test_adam_state_restored_on_device(self):
+        p = _make_ptycho()
+        p.reconstruct(
+            num_iters=N_ITERS, reset=True, optimizer_params=_opt(), batch_size=16, device=GPU_IDS
+        )
+        obj_opt = p.optimizers.get("object")
+        assert obj_opt is not None
+        assert len(obj_opt.state) > 0
+        first = next(iter(obj_opt.state.values()))
+        assert "exp_avg" in first, "Adam moments missing"
+        assert first["exp_avg"].device.type == "cuda"
+
+
+class TestDIPMultiGPU:
+    """DIP path mirrors the pixelated multi-GPU contract."""
+
+    def test_dip_device_and_loss_lifecycle(self):
+        d = _make_dip_ptycho()
+        d.reconstruct(
+            num_iters=N_ITERS,
+            reset=True,
+            lr_obj=1e-3,
+            lr_probe=1e-3,
+            batch_size=16,
+            device=GPU_IDS,
+        )
+        assert d.device == GPU_IDS
+        assert len(d._iter_losses) == N_ITERS
+        assert "object" in d.iter_lrs
+        assert len(d.iter_lrs["object"]) == N_ITERS
+
+        d.reconstruct(
+            num_iters=N_ITERS,
+            reset=False,
+            lr_obj=1e-3,
+            lr_probe=1e-3,
+            batch_size=16,
+            device=GPU_IDS,
+        )
+        assert len(d._iter_losses) == 2 * N_ITERS, "DIP continuation must not double-count"
+
+        d.reconstruct(
+            num_iters=N_ITERS,
+            reset=True,
+            lr_obj=1e-3,
+            lr_probe=1e-3,
+            batch_size=16,
+            device=GPU_IDS,
+        )
+        assert len(d._iter_losses) == N_ITERS
diff --git a/tests/diffractive_imaging/test_ptychography.py b/tests/diffractive_imaging/test_ptychography.py
index 1434a8ee8..6d9c16ce6 100644
--- a/tests/diffractive_imaging/test_ptychography.py
+++ b/tests/diffractive_imaging/test_ptychography.py
@@ -1,5 +1,6 @@
 """
-Tests for ptychography gradient equivalence between autograd and analytical methods
+Tests for ptychography gradient equivalence between autograd and analytical methods,
+plus property-style tests for state management and serialization.
 """
 
 import numpy as np
@@ -8,6 +9,8 @@
 
 from quantem.core import config
 from quantem.core.datastructures.dataset4dstem import Dataset4dstem
+from quantem.core.io.serialize import load as autoserialize_load
+from quantem.core.ml import OptimizerParams
 from quantem.core.utils.utils import electron_wavelength_angstrom
 from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster
 from quantem.diffractive_imaging.detector_models import DetectorPixelated
@@ -237,7 +240,7 @@ def test_single_probe_gradients(self, single_probe_ptycho_model):
         }
 
         ptycho.reconstruct(
-            num_iter=1,
+            num_iters=1,
             reset=True,
             autograd=True,
             constraints=constraints,
@@ -249,7 +252,7 @@ def test_single_probe_gradients(self, single_probe_ptycho_model):
         grads_probe_ad = ptycho.probe_model._probe.grad.clone().detach().cpu().numpy()
 
         ptycho.reconstruct(
-            num_iter=1,
+            num_iters=1,
             reset=True,
             autograd=False,
             constraints=constraints,
@@ -311,7 +314,7 @@ def test_mixed_probe_gradients(self, mixed_probe_ptycho_model):
         }
 
         ptycho.reconstruct(
-            num_iter=1,
+            num_iters=1,
             reset=True,
             autograd=True,
             constraints=constraints,
@@ -323,7 +326,7 @@ def test_mixed_probe_gradients(self, mixed_probe_ptycho_model):
         grads_probe_ad = ptycho.probe_model._probe.grad.clone().detach().cpu().numpy()
 
         ptycho.reconstruct(
-            num_iter=1,
+            num_iters=1,
             reset=True,
             autograd=False,
             constraints=constraints,
@@ -360,3 +363,112 @@ def test_mixed_probe_gradients(self, mixed_probe_ptycho_model):
 
         assert ssim_obj_abs > 0.99  # type: ignore
         assert ssim_probe_angle > 0.7  # type: ignore
+
+
+class TestTargetResidency:
+    """Property + serialization behavior for the streaming-target knob."""
+
+    def test_default_is_device(self, ptycho_dataset):
+        assert ptycho_dataset.target_residency == "device"
+
+    def test_setter_accepts_valid(self, ptycho_dataset):
+        ptycho_dataset.target_residency = "cpu"
+        assert ptycho_dataset.target_residency == "cpu"
+        ptycho_dataset.target_residency = "device"
+        assert ptycho_dataset.target_residency == "device"
+
+    @pytest.mark.parametrize("bad", ["gpu", "GPU", "CPU", "", "cuda", "Device"])
+    def test_setter_rejects_invalid(self, ptycho_dataset, bad):
+        with pytest.raises(ValueError, match="target_residency"):
+            ptycho_dataset.target_residency = bad
+        # value should be unchanged after a rejected set
+        assert ptycho_dataset.target_residency == "device"
+
+    def test_save_load_roundtrip(self, ptycho_dataset, tmp_path):
+        ptycho_dataset.target_residency = "cpu"
+        path = tmp_path / "pdset.zip"
+        ptycho_dataset.save(str(path))
+        reloaded = autoserialize_load(str(path))
+        assert reloaded.target_residency == "cpu"
+
+
+@pytest.mark.slow
+class TestPtychographySaveLoadRoundtrip:
+    """Reconstruct → save → load → continue training preserves training state.
+
+    The 0.3 threshold reflects that, on this synthetic ducky-style dataset with the
+    analytical probe already in hand, a well-formed reconstruction should drive the
+    loss down by at least 70% in 20 iterations on the right configuration. The bar
+    is deliberately strict — if you tune optimizer settings and this fires, the
+    config probably regressed.
+    """
+
+    NUM_ITERS = 50  # enough headroom for the strict 0.3 threshold at lr=5e-3
+
+    @pytest.fixture
+    def trained_ptycho(self, single_probe_ptycho_model):
+        ptycho = single_probe_ptycho_model
+        ptycho.reconstruct(
+            num_iters=self.NUM_ITERS,
+            reset=True,
+            optimizer_params={
+                "object": OptimizerParams.Adam(lr=5e-3),
+                "probe": OptimizerParams.Adam(lr=5e-3),
+            },
+            batch_size=N**2,
+            device=config.get_device(),
+        )
+        return ptycho
+
+    def test_iter_losses_preserved(self, trained_ptycho, tmp_path):
+        path = tmp_path / "ptycho.zip"
+        trained_ptycho.save(str(path), save_raw_data=True)
+        reloaded = autoserialize_load(str(path))
+        np.testing.assert_array_equal(reloaded._iter_losses, trained_ptycho._iter_losses)
+
+    def test_scan_positions_preserved(self, trained_ptycho, tmp_path):
+        path = tmp_path / "ptycho.zip"
+        trained_ptycho.save(str(path), save_raw_data=True)
+        reloaded = autoserialize_load(str(path))
+        original = trained_ptycho.dset.scan_positions_px.detach().cpu().numpy()
+        new = reloaded.dset.scan_positions_px.detach().cpu().numpy()
+        np.testing.assert_allclose(new, original, rtol=0, atol=0)
+
+    def test_object_preserved(self, trained_ptycho, tmp_path):
+        path = tmp_path / "ptycho.zip"
+        trained_ptycho.save(str(path), save_raw_data=True)
+        reloaded = autoserialize_load(str(path))
+        original = trained_ptycho.obj_model._obj.detach().cpu().numpy()
+        new = reloaded.obj_model._obj.detach().cpu().numpy()
+        np.testing.assert_allclose(new, original, rtol=0, atol=0)
+
+    def test_loss_decreases_below_strict_threshold(self, trained_ptycho):
+        losses = trained_ptycho._iter_losses
+        assert losses[-1] < 0.3 * losses[0], (
+            f"loss should drop below 30% of initial in {self.NUM_ITERS} iters: "
+            f"initial={losses[0]:.3e}, final={losses[-1]:.3e}, "
+            f"ratio={losses[-1] / losses[0]:.2f}"
+        )
+
+    def test_continue_training_after_reload(self, trained_ptycho, tmp_path):
+        """Reload a trained ptycho, continue training, and verify the loss keeps
+        decreasing — confirms optimizer state and parameter bindings survive the
+        save/load roundtrip end-to-end."""
+        path = tmp_path / "ptycho.zip"
+        trained_ptycho.save(str(path), save_raw_data=True)
+        reloaded = autoserialize_load(str(path))
+        loss_after_reload = reloaded._iter_losses[-1]
+
+        n_continue = 10
+        reloaded.reconstruct(
+            num_iters=n_continue,
+            reset=False,
+            batch_size=N**2,
+            device=config.get_device(),
+        )
+        assert len(reloaded._iter_losses) == self.NUM_ITERS + n_continue, (
+            "continuation must not reset history"
+        )
+        assert reloaded._iter_losses[-1] <= loss_after_reload, (
+            "loss must not regress after reload — optimizer state likely lost"
+        )

From 07d739b22669399f54a7f29ce704537acb4a3290 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 27 May 2026 11:28:44 -0700
Subject: [PATCH 19/59] adding iterative ptycho constraint params

---
 src/quantem/core/ml/constraints.py            |  98 +++++++++-
 src/quantem/diffractive_imaging/__init__.py   |   8 +-
 .../diffractive_imaging/constraints.py        |  99 ----------
 .../diffractive_imaging/dataset_models.py     |  89 ++++++---
 .../diffractive_imaging/object_models.py      | 146 +++++++++++----
 .../diffractive_imaging/probe_models.py       |  88 +++++++--
 .../diffractive_imaging/ptychography.py       |  74 +++++++-
 .../diffractive_imaging/ptychography_base.py  |  29 ++-
 .../diffractive_imaging/ptychography_lite.py  |  33 +++-
 tests/diffractive_imaging/test_constraints.py | 175 ++++++++++++++++++
 10 files changed, 641 insertions(+), 198 deletions(-)
 delete mode 100644 src/quantem/diffractive_imaging/constraints.py
 create mode 100644 tests/diffractive_imaging/test_constraints.py

diff --git a/src/quantem/core/ml/constraints.py b/src/quantem/core/ml/constraints.py
index 553b06115..0adf6b57a 100644
--- a/src/quantem/core/ml/constraints.py
+++ b/src/quantem/core/ml/constraints.py
@@ -1,12 +1,14 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Self
+from typing import Any, Generic, Self, TypeVar
 
 import numpy as np
 import torch
 from numpy.typing import NDArray
 
+from quantem.core import config
+
 
 @dataclass(slots=False)
 class Constraints(ABC):
@@ -47,17 +49,27 @@ def __str__(self) -> str:
         )
 
 
-class BaseConstraints(ABC):
+C = TypeVar("C", bound=Constraints)
+
+
+class BaseConstraints(ABC, Generic[C]):
     """
     Base class for constraints.
+
+    Generic over a concrete ``Constraints`` subclass so that subclasses (and the
+    type checker) can see the specific fields available on ``self.constraints``.
+    Subclasses parameterize like ``BaseConstraints[MyConstraintsType]`` and set
+    ``DEFAULT_CONSTRAINTS`` to an instance of that type.
     """
 
-    # Default constraints are the dataclasses themselves.
-    DEFAULT_CONSTRAINTS = Constraints()
+    DEFAULT_CONSTRAINTS: C
+    _constraints: C
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._soft_constraint_losses = []
+        self._soft_constraint_loss: dict[str, torch.Tensor | float] = {}
+        self._iter_constraint_losses: dict[str, float] = {}
         self.constraints = self.DEFAULT_CONSTRAINTS.copy()
 
     @property
@@ -65,30 +77,100 @@ def soft_constraint_losses(self) -> NDArray[np.float32]:
         return np.array(self._soft_constraint_losses, dtype=np.float32)
 
     @property
-    def constraints(self) -> Constraints:
+    def soft_constraint_loss(self) -> dict[str, torch.Tensor | float]:
+        return self._soft_constraint_loss
+
+    @property
+    def constraints(self) -> C:
         """
         Constraints for the model.
         """
         return self._constraints
 
     @constraints.setter
-    def constraints(self, constraints: Constraints | dict[str, Any]):
+    def constraints(self, constraints: C | dict[str, Any]):
         """
         Setter for constraints class, can be a Constraints instance or a dictionary.
+        Dict keys are validated against the active Constraints dataclass's allowed_keys.
         """
         if isinstance(constraints, Constraints):
             self._constraints = constraints
         elif isinstance(constraints, dict):
+            allowed = self._constraints.allowed_keys
             for key, value in constraints.items():
+                if key not in allowed:
+                    raise KeyError(
+                        f"Invalid constraint key '{key}' for {type(self._constraints).__name__}, "
+                        f"allowed keys are {allowed}"
+                    )
                 setattr(self._constraints, key, value)
         else:
             raise ValueError(f"Invalid constraints type: {type(constraints)}")
 
-    # --- Required methods tha tneeds to implemented in subclasses ---
+    def add_constraint(self, key: str, value: Any) -> None:
+        """
+        Set a single constraint field by name, with validation against allowed_keys.
+        """
+        allowed = self._constraints.allowed_keys
+        if key not in allowed:
+            raise KeyError(
+                f"Invalid constraint key '{key}' for {type(self._constraints).__name__}, "
+                f"allowed keys are {allowed}"
+            )
+        setattr(self._constraints, key, value)
+
+    # --- helpers for consistent loss logging ---
+    def _get_zero_loss_tensor(self) -> torch.Tensor:
+        """Helper method to create a zero loss tensor with proper device and dtype."""
+        device = getattr(self, "device", "cpu")
+        return torch.tensor(0, device=device, dtype=getattr(torch, config.get("dtype_real")))
+
+    def reset_soft_constraint_losses(self) -> None:
+        self._soft_constraint_loss = {}
+
+    def add_soft_constraint_loss(self, name: str, value: torch.Tensor | float) -> None:
+        """Record a single soft-constraint loss for logging without holding the graph."""
+        if isinstance(value, torch.Tensor):
+            val = value.detach()
+            if val.ndim != 0:
+                val = val.mean()
+            self._soft_constraint_loss[name] = val
+        else:
+            self._soft_constraint_loss[name] = float(value)
+
+    def accumulate_constraint_losses(
+        self, batch_constraint_losses: dict[str, torch.Tensor | float] | None = None
+    ) -> None:
+        """Accumulate constraint losses across batches."""
+        if batch_constraint_losses is None:
+            batch_constraint_losses = self.soft_constraint_loss
+
+        for loss_name, loss_value in batch_constraint_losses.items():
+            if isinstance(loss_value, torch.Tensor):
+                try:
+                    v = loss_value.item()
+                except Exception:
+                    v = loss_value.detach().mean().item()
+            else:
+                v = float(loss_value)
+            self._iter_constraint_losses[loss_name] = (
+                self._iter_constraint_losses.get(loss_name, 0.0) + v
+            )
+
+    def get_iter_constraint_losses(self) -> dict[str, float]:
+        return self._iter_constraint_losses
+
+    def reset_iter_constraint_losses(self) -> None:
+        self._iter_constraint_losses = {}
+
+    # --- Required methods that need to be implemented in subclasses ---
     @abstractmethod
-    def apply_hard_constraints(self, *args, **kwargs) -> torch.Tensor:
+    def apply_hard_constraints(self, *args, **kwargs) -> torch.Tensor | None:
         """
         Apply hard constraints to the model.
+
+        May return a projected tensor (most models) or ``None`` when the
+        implementation mutates state in place (e.g. ``DatasetConstraints``).
         """
         raise NotImplementedError
 
diff --git a/src/quantem/diffractive_imaging/__init__.py b/src/quantem/diffractive_imaging/__init__.py
index 2c26de60c..9db66f112 100644
--- a/src/quantem/diffractive_imaging/__init__.py
+++ b/src/quantem/diffractive_imaging/__init__.py
@@ -1,15 +1,21 @@
 from quantem.diffractive_imaging.dataset_models import (
+    PtychoDatasetConstraintParams as PtychoDatasetConstraintParams,
+    PtychoDatasetConstraintsType as PtychoDatasetConstraintsType,
     PtychographyDatasetRaster as PtychographyDatasetRaster,
 )
 from quantem.diffractive_imaging.detector_models import DetectorPixelated as DetectorPixelated
 from quantem.diffractive_imaging.object_models import (
     ObjectDIP as ObjectDIP,
     ObjectPixelated as ObjectPixelated,
+    PtychoObjConstraintParams as PtychoObjConstraintParams,
+    PtychoObjConstraintsType as PtychoObjConstraintsType,
 )
 from quantem.diffractive_imaging.probe_models import (
     ProbeDIP as ProbeDIP,
-    ProbePixelated as ProbePixelated,
     ProbeParametric as ProbeParametric,
+    ProbePixelated as ProbePixelated,
+    PtychoProbeConstraintParams as PtychoProbeConstraintParams,
+    PtychoProbeConstraintsType as PtychoProbeConstraintsType,
 )
 from quantem.diffractive_imaging.ptychography import Ptychography as Ptychography
 from quantem.diffractive_imaging.ptychography_lite import (
diff --git a/src/quantem/diffractive_imaging/constraints.py b/src/quantem/diffractive_imaging/constraints.py
deleted file mode 100644
index e907690e3..000000000
--- a/src/quantem/diffractive_imaging/constraints.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any
-
-import torch
-
-from quantem.core import config
-
-
-class BaseConstraints(ABC):
-    """Base class for constraint management with common functionality."""
-
-    # Subclasses should define their own DEFAULT_CONSTRAINTS
-    DEFAULT_CONSTRAINTS: dict[str, Any] = {}
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._soft_constraint_loss = {}
-        self._constraints = self.DEFAULT_CONSTRAINTS.copy()
-        self._iter_constraint_losses = {}
-
-    @property
-    def constraints(self) -> dict[str, Any]:
-        return self._constraints
-
-    @constraints.setter
-    def constraints(self, c: dict[str, Any]):
-        allowed_keys = self.DEFAULT_CONSTRAINTS.keys()
-        constraint_type = self.__class__.__name__.lower().replace("constraints", "")
-
-        for key, value in c.items():
-            if key not in allowed_keys:
-                raise KeyError(
-                    f"Invalid {constraint_type} constraint key '{key}', allowed keys are {list(allowed_keys)}"
-                )
-            self._constraints[key] = value
-
-    @property
-    def soft_constraint_loss(self) -> dict[str, torch.Tensor | float]:
-        return self._soft_constraint_loss
-
-    def add_constraint(self, key: str, value: Any):
-        allowed_keys = self.DEFAULT_CONSTRAINTS.keys()
-        constraint_type = self.__class__.__name__.lower().replace("constraints", "")
-
-        if key not in allowed_keys:
-            raise KeyError(
-                f"Invalid {constraint_type} constraint key '{key}', allowed keys are {list(allowed_keys)}"
-            )
-        self._constraints[key] = value
-
-    @abstractmethod
-    def apply_soft_constraints(self, *args, **kwargs) -> torch.Tensor:
-        """Apply soft constraints and return total constraint loss."""
-        pass
-
-    def _get_zero_loss_tensor(self) -> torch.Tensor:
-        """Helper method to create a zero loss tensor with proper device and dtype."""
-        device = getattr(self, "device", "cpu")
-        return torch.tensor(0, device=device, dtype=getattr(torch, config.get("dtype_real")))
-
-    # --- helpers for consistent loss logging ---
-    def reset_soft_constraint_losses(self) -> None:
-        self._soft_constraint_loss = {}
-
-    def add_soft_constraint_loss(self, name: str, value: torch.Tensor | float) -> None:
-        """Record a single soft-constraint loss for logging without holding the graph."""
-        if isinstance(value, torch.Tensor):
-            val = value.detach()
-            if val.ndim != 0:
-                val = val.mean()
-            self._soft_constraint_loss[name] = val
-        else:
-            self._soft_constraint_loss[name] = float(value)
-
-    def accumulate_constraint_losses(
-        self, batch_constraint_losses: dict[str, torch.Tensor | float] | None = None
-    ) -> None:
-        """Accumulate constraint losses across batches."""
-        if batch_constraint_losses is None:
-            batch_constraint_losses = self.soft_constraint_loss
-
-        for loss_name, loss_value in batch_constraint_losses.items():
-            if isinstance(loss_value, torch.Tensor):
-                try:
-                    v = loss_value.item()
-                except Exception:
-                    print("loss value not singular: ", loss_value)  # TODO remove
-                    v = loss_value.detach().mean().item()
-            else:
-                v = float(loss_value)
-            self._iter_constraint_losses[loss_name] = (
-                self._iter_constraint_losses.get(loss_name, 0.0) + v
-            )
-
-    def get_iter_constraint_losses(self) -> dict[str, float]:
-        return getattr(self, "_iter_constraint_losses", {})  # TODO clean this up
-
-    def reset_iter_constraint_losses(self) -> None:
-        self._iter_constraint_losses = {}
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 835720908..5b5b49b4f 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Literal, Self
 
@@ -12,6 +13,7 @@
 from quantem.core.datastructures.dataset3d import Dataset3d
 from quantem.core.datastructures.dataset4dstem import Dataset4dstem
 from quantem.core.io.serialize import AutoSerialize
+from quantem.core.ml.constraints import BaseConstraints, Constraints
 from quantem.core.ml.optimizer_mixin import OptimizerMixin
 from quantem.core.utils.utils import electron_wavelength_angstrom, tqdmnd
 from quantem.core.utils.validators import (
@@ -21,7 +23,6 @@
     validate_tensor,
 )
 from quantem.core.visualization import show_2d
-from quantem.diffractive_imaging.constraints import BaseConstraints
 from quantem.diffractive_imaging.ptycho_utils import AffineTransform, fit_origin, shift_array
 
 """
@@ -29,6 +30,61 @@
 """
 
 
+class PtychoDatasetConstraintParams:
+    """
+    Namespace class for ptychography dataset constraint dataclasses.
+
+    Tab-complete on ``PtychoDatasetConstraintParams`` in a notebook to discover the
+    available variants. Tab-complete inside a variant's constructor to see every
+    constraint field with its default value.
+
+    Variants
+    --------
+    Raster
+        Constraints for ``PtychographyDatasetRaster`` (descan TV penalty, descan
+        zero-out, scan position clipping and centering).
+    """
+
+    @dataclass
+    class Raster(Constraints):
+        """Constraints for raster-scan ptychography datasets."""
+
+        # hard constraints
+        descan_shifts_constant: bool = False
+        center_scan_positions: bool = False
+        clip_scan_positions: bool = True
+        # soft constraints
+        descan_tv_weight: float = 0.0
+        _name: str = "raster"
+
+        soft_constraint_keys = ["descan_tv_weight"]
+        hard_constraint_keys = [
+            "descan_shifts_constant",
+            "center_scan_positions",
+            "clip_scan_positions",
+        ]
+
+    @classmethod
+    def parse_dict(cls, d: dict) -> "PtychoDatasetConstraintsType":
+        d = dict(d)
+        name = d.pop("name", None) or d.pop("type", None)
+        if name is None:
+            raise ValueError("Must provide either 'name' or 'type' key")
+        if isinstance(name, type):
+            name = name.__name__.lower()
+        elif isinstance(name, str):
+            name = name.lower()
+        else:
+            raise ValueError(f"Unknown dataset constraint type: {name}")
+        if name == "raster":
+            return cls.Raster(**d)
+        else:
+            raise ValueError(f"Unknown dataset constraint type: {name}")
+
+
+PtychoDatasetConstraintsType = PtychoDatasetConstraintParams.Raster
+
+
 class PtychographyDatasetBase(
     AutoSerialize, OptimizerMixin, torch.nn.Module, torch.utils.data.Dataset
 ):
@@ -602,24 +658,17 @@ def reset(self) -> None:
     # endregion --- class methods ---
 
 
-class DatasetConstraints(BaseConstraints, PtychographyDatasetBase):
-    DEFAULT_CONSTRAINTS = {
-        "descan_tv_weight": 0.0,
-        "descan_shifts_constant": False,
-        "center_scan_positions": False,
-        "clip_scan_positions": True,
-    }
+class DatasetConstraints(
+    BaseConstraints[PtychoDatasetConstraintParams.Raster], PtychographyDatasetBase
+):
+    DEFAULT_CONSTRAINTS: PtychoDatasetConstraintParams.Raster = PtychoDatasetConstraintParams.Raster()
 
     def apply_soft_constraints(self, descan_shifts: torch.Tensor) -> torch.Tensor:
         self.reset_soft_constraint_losses()
         loss = self._get_zero_loss_tensor()
 
-        if (
-            self.constraints.get("descan_tv_weight", 0) > 0
-            and self.learn_descan
-            and self.has_optimizer()
-        ):
-            tv_loss = self.get_descan_tv_loss(descan_shifts, self.constraints["descan_tv_weight"])
+        if self.constraints.descan_tv_weight > 0 and self.learn_descan and self.has_optimizer():
+            tv_loss = self.get_descan_tv_loss(descan_shifts, self.constraints.descan_tv_weight)
             loss = loss + tv_loss
             self.add_soft_constraint_loss("descan_tv_weight", tv_loss)
 
@@ -639,23 +688,17 @@ def apply_descan_constraints(
         self,
         descan: torch.Tensor,
     ) -> torch.Tensor:
-        if self.constraints["descan_shifts_constant"]:
+        if self.constraints.descan_shifts_constant:
             descan = torch.zeros_like(descan)
         return descan
 
     def apply_hard_constraints(self, obj_padding_px: np.ndarray | tuple) -> None:
-        # could clip positions here if needed
         positions = self.scan_positions_px
         obj_shape = torch.tensor(self._obj_shape_full_2d(obj_padding_px), device=positions.device)
-        if self.constraints.get(
-            "clip_scan_positions", self.DEFAULT_CONSTRAINTS["clip_scan_positions"]
-        ):
+        if self.constraints.clip_scan_positions:
             positions = torch.clamp(positions, min=torch.zeros_like(obj_shape), max=obj_shape - 1)
 
-        if self.constraints.get(
-            "center_scan_positions", self.DEFAULT_CONSTRAINTS["center_scan_positions"]
-        ):
-            # shift all positions uniformly so that the mean position is at the center of the object
+        if self.constraints.center_scan_positions:
             positions = positions - positions.mean(dim=0, keepdim=True)
             positions = positions + obj_shape / 2
 
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index ee4c672e8..0fd00c69c 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1,6 +1,7 @@
 import math
 from abc import abstractmethod
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Callable, Literal, Self, Sequence, cast
 from warnings import warn
 
@@ -13,6 +14,7 @@
 from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.blocks import reset_weights
+from quantem.core.ml.constraints import BaseConstraints, Constraints
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType
 from quantem.core.utils.rng import RNGMixin
@@ -23,11 +25,104 @@
 )
 from quantem.core.visualization import show_2d
 from quantem.core.visualization.custom_normalizations import CustomNormalization
-from quantem.diffractive_imaging.constraints import BaseConstraints
 from quantem.diffractive_imaging.ptycho_utils import sum_patches
 
 object_type = Literal["potential", "pure_phase", "complex"]
 
+
+class PtychoObjConstraintParams:
+    """
+    Namespace class for ptychography object constraint dataclasses.
+
+    Tab-complete on ``PtychoObjConstraintParams`` in a notebook to discover the
+    available variants. Tab-complete inside a variant's constructor to see every
+    constraint field with its default value.
+
+    Variants
+    --------
+    Raster
+        Constraints for grid-based object representations (``ObjectPixelated`` and
+        ``ObjectDIP`` share this set today).
+    INR
+        Placeholder for the upcoming implicit-neural-representation object.
+
+    Examples
+    --------
+    >>> PtychoObjConstraintParams.Raster(tv_weight_z=5.0, identical_slices=True)
+    >>> PtychoObjConstraintParams.parse_dict({"name": "raster", "positivity": False})
+    """
+
+    @dataclass
+    class Raster(Constraints):
+        """Constraints for grid-based ptychography object models (Pixelated, DIP)."""
+
+        # hard constraints
+        positivity: bool = True
+        fix_potential_baseline: bool = False
+        fix_potential_baseline_factor: float = 1.0
+        identical_slices: bool = False
+        apply_fov_mask: bool = False
+        # filtering (treated as hard, applied post-update)
+        gaussian_sigma: float | None = None  # pixels
+        butterworth_order: int = 4
+        q_lowpass: float | None = None  # A^-1
+        q_highpass: float | None = None  # A^-1
+        # soft constraints
+        tv_weight_z: float = 0.0
+        tv_weight_xy: float = 0.0
+        surface_zero_weight: float = 0.0
+        _name: str = "raster"
+
+        soft_constraint_keys = ["tv_weight_z", "tv_weight_xy", "surface_zero_weight"]
+        hard_constraint_keys = [
+            "positivity",
+            "fix_potential_baseline",
+            "fix_potential_baseline_factor",
+            "identical_slices",
+            "apply_fov_mask",
+            "gaussian_sigma",
+            "butterworth_order",
+            "q_lowpass",
+            "q_highpass",
+        ]
+
+    @dataclass
+    class INR(Constraints):
+        """Placeholder for the upcoming ObjectINR variant. Not yet wired to a model."""
+
+        _name: str = "inr"
+
+        soft_constraint_keys = []
+        hard_constraint_keys = []
+
+    @classmethod
+    def parse_dict(cls, d: dict) -> "PtychoObjConstraintsType":
+        """Instantiate the appropriate Raster or INR dataclass from a config dict.
+
+        The dict must contain a ``'name'`` or ``'type'`` key (case-insensitive),
+        with value ``'raster'`` or ``'inr'``. All other keys are forwarded as
+        keyword arguments to the chosen dataclass.
+        """
+        d = dict(d)
+        name = d.pop("name", None) or d.pop("type", None)
+        if name is None:
+            raise ValueError("Must provide either 'name' or 'type' key")
+        if isinstance(name, type):
+            name = name.__name__.lower()
+        elif isinstance(name, str):
+            name = name.lower()
+        else:
+            raise ValueError(f"Unknown object constraint type: {name}")
+        if name == "raster":
+            return cls.Raster(**d)
+        elif name == "inr":
+            return cls.INR(**d)
+        else:
+            raise ValueError(f"Unknown object constraint type: {name}")
+
+
+PtychoObjConstraintsType = PtychoObjConstraintParams.Raster | PtychoObjConstraintParams.INR
+
 """
 Currently all object models.obj are complex valued for "complex" or "pure_phase" object types,
 and real valued for "potential" object types. This could be changed to be always complex valued, 
@@ -271,37 +366,25 @@ def backward(self, *args, **kwargs):
         )
 
 
-class ObjectConstraints(BaseConstraints, ObjectBase):
-    DEFAULT_CONSTRAINTS = {
-        "positivity": True,
-        "fix_potential_baseline": False,
-        "fix_potential_baseline_factor": 1.0,
-        "identical_slices": False,
-        "apply_fov_mask": False,
-        "tv_weight_z": 0,
-        "tv_weight_xy": 0,
-        "surface_zero_weight": 0,
-        "gaussian_sigma": None,  # pixels
-        "butterworth_order": 4,
-        "q_lowpass": None,  # A^-1
-        "q_highpass": None,  # A^-1
-    }
+class ObjectConstraints(BaseConstraints[PtychoObjConstraintParams.Raster], ObjectBase):
+    DEFAULT_CONSTRAINTS: PtychoObjConstraintParams.Raster = PtychoObjConstraintParams.Raster()
 
     def apply_hard_constraints(
         self, obj: torch.Tensor, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
+        c = self.constraints
         if self.obj_type in ["complex", "pure_phase"]:
             if self.obj_type == "complex":
                 amp = torch.clamp(torch.abs(obj), 0.0, 1.0)
             else:
                 amp = 1.0
             phase = obj.angle() - obj.angle().mean()
-            if mask is not None and self.constraints["apply_fov_mask"]:
+            if mask is not None and c.apply_fov_mask:
                 obj2 = amp * mask * torch.exp(1.0j * phase * mask)
             else:
                 obj2 = amp * torch.exp(1.0j * phase)
         else:  # potential
-            if self.constraints["fix_potential_baseline"]:
+            if c.fix_potential_baseline:
                 if mask is not None:
                     background = mask < 0.5 * mask.max()
                     if background.any():
@@ -311,29 +394,28 @@ def apply_hard_constraints(
                 else:
                     offset = obj.min()
                 offset = offset.detach()
-                offset *= self.constraints["fix_potential_baseline_factor"]
+                offset *= c.fix_potential_baseline_factor
             else:
                 offset = 0
 
-            if self.constraints.get("positivity", True):
+            if c.positivity:
                 obj2 = torch.clamp(obj - offset, min=0.0)
             else:
                 obj2 = obj - offset
 
-        if self.constraints["apply_fov_mask"] and mask is not None:
+        if c.apply_fov_mask and mask is not None:
             obj2 *= mask
 
-        # want backwards compatibility for gaussian_sigma and q_lowpass/q_highpass, so use get
-        if self.constraints.get("gaussian_sigma") is not None:
-            obj2 = self.gaussian_blur_2d(obj2, sigma=self.constraints["gaussian_sigma"])
+        if c.gaussian_sigma is not None:
+            obj2 = self.gaussian_blur_2d(obj2, sigma=c.gaussian_sigma)
 
-        if any([self.constraints["q_lowpass"], self.constraints["q_highpass"]]):
+        if any([c.q_lowpass, c.q_highpass]):
             obj2 = self.butterworth_constraint(
                 obj2,
                 sampling=self.sampling,
             )
         if self.num_slices > 1:
-            if self.constraints["identical_slices"]:
+            if c.identical_slices:
                 with torch.no_grad():
                     obj2[:] = torch.mean(obj2, dim=0, keepdim=True)
 
@@ -352,7 +434,7 @@ def apply_soft_constraints(
 
         surface_zero_loss = self.get_surface_zero_loss(
             obj,
-            weight=self.constraints["surface_zero_weight"],
+            weight=self.constraints.surface_zero_weight,
         )
         self.add_soft_constraint_loss("surface_zero_loss", surface_zero_loss)
         self.accumulate_constraint_losses()
@@ -364,8 +446,8 @@ def get_tv_loss(
         loss = self._get_zero_loss_tensor()
         if weights is None:
             w = (
-                self.constraints["tv_weight_z"],
-                self.constraints["tv_weight_xy"],
+                self.constraints.tv_weight_z,
+                self.constraints.tv_weight_xy,
             )
         elif isinstance(weights, (float, int)):
             if weights == 0:
@@ -489,9 +571,9 @@ def butterworth_constraint(
 
         """
 
-        q_lowpass = self.constraints["q_lowpass"]
-        q_highpass = self.constraints["q_highpass"]
-        butterworth_order = self.constraints["butterworth_order"]
+        q_lowpass = self.constraints.q_lowpass
+        q_highpass = self.constraints.q_highpass
+        butterworth_order = self.constraints.butterworth_order
 
         qx = torch.fft.fftfreq(tensor.shape[-2], sampling[0], device=tensor.device)
         qy = torch.fft.fftfreq(tensor.shape[-1], sampling[1], device=tensor.device)
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index 56bb89f6c..c8666fb86 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -1,5 +1,6 @@
 from abc import abstractmethod
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any, Callable, Self, Union
 from warnings import warn
 
@@ -14,6 +15,7 @@
 from quantem.core.datastructures import Dataset2d, Dataset4dstem
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.blocks import reset_weights
+from quantem.core.ml.constraints import BaseConstraints, Constraints
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType
 from quantem.core.utils.rng import RNGMixin
@@ -32,7 +34,6 @@
     POLAR_SYMBOLS,
     real_space_probe,
 )
-from quantem.diffractive_imaging.constraints import BaseConstraints
 from quantem.diffractive_imaging.ptycho_utils import (
     fourier_shift_expand,
     shift_array,
@@ -41,6 +42,72 @@
 DeviceType = Union[str, torch.device, int]
 
 
+class PtychoProbeConstraintParams:
+    """
+    Namespace class for ptychography probe constraint dataclasses.
+
+    Tab-complete on ``PtychoProbeConstraintParams`` in a notebook to discover the
+    available variants. Tab-complete inside a variant's constructor to see every
+    constraint field with its default value.
+
+    Variants
+    --------
+    Raster
+        Constraints for grid-based probe representations (``ProbePixelated`` and
+        ``ProbeDIP`` share this set today).
+    Parametric
+        Placeholder for parametric probe models, where Gram-Schmidt orthogonalization
+        and pixel-domain TV are moot.
+    """
+
+    @dataclass
+    class Raster(Constraints):
+        """Constraints for grid-based ptychography probe models (Pixelated, DIP)."""
+
+        # hard constraints
+        orthogonalize_probe: bool = True
+        center_probe: bool = False
+        # soft constraints
+        tv_weight: float = 0.0
+        _name: str = "raster"
+
+        soft_constraint_keys = ["tv_weight"]
+        hard_constraint_keys = ["orthogonalize_probe", "center_probe"]
+
+    @dataclass
+    class Parametric(Constraints):
+        """Placeholder for parametric probe constraints. Not yet wired to a model."""
+
+        _name: str = "parametric"
+
+        soft_constraint_keys = []
+        hard_constraint_keys = []
+
+    @classmethod
+    def parse_dict(cls, d: dict) -> "PtychoProbeConstraintsType":
+        d = dict(d)
+        name = d.pop("name", None) or d.pop("type", None)
+        if name is None:
+            raise ValueError("Must provide either 'name' or 'type' key")
+        if isinstance(name, type):
+            name = name.__name__.lower()
+        elif isinstance(name, str):
+            name = name.lower()
+        else:
+            raise ValueError(f"Unknown probe constraint type: {name}")
+        if name == "raster":
+            return cls.Raster(**d)
+        elif name == "parametric":
+            return cls.Parametric(**d)
+        else:
+            raise ValueError(f"Unknown probe constraint type: {name}")
+
+
+PtychoProbeConstraintsType = (
+    PtychoProbeConstraintParams.Raster | PtychoProbeConstraintParams.Parametric
+)
+
+
 class ProbeBase(nn.Module, RNGMixin, OptimizerMixin, AutoSerialize):
     DEFAULT_PROBE_PARAMS = {
         "energy": None,
@@ -427,13 +494,8 @@ def _compute_propagator_arrays(
         return propagators
 
 
-class ProbeConstraints(BaseConstraints, ProbeBase):
-    DEFAULT_CONSTRAINTS = {
-        # "fix_probe": False,
-        "orthogonalize_probe": True,
-        "center_probe": False,
-        "tv_weight": 0.0,
-    }
+class ProbeConstraints(BaseConstraints[PtychoProbeConstraintParams.Raster], ProbeBase):
+    DEFAULT_CONSTRAINTS: PtychoProbeConstraintParams.Raster = PtychoProbeConstraintParams.Raster()
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -441,8 +503,8 @@ def __init__(self, *args, **kwargs):
     def apply_soft_constraints(self, probe: torch.Tensor) -> torch.Tensor:
         self.reset_soft_constraint_losses()
         loss = self._get_zero_loss_tensor()
-        if self.constraints["tv_weight"]:
-            loss_tv = self._probe_tv_constraint(probe, self.constraints["tv_weight"])
+        if self.constraints.tv_weight:
+            loss_tv = self._probe_tv_constraint(probe, self.constraints.tv_weight)
             self.add_soft_constraint_loss("tv_loss", loss_tv)
             loss = loss + loss_tv
 
@@ -450,11 +512,9 @@ def apply_soft_constraints(self, probe: torch.Tensor) -> torch.Tensor:
         return loss
 
     def apply_hard_constraints(self, probe: torch.Tensor) -> torch.Tensor:
-        # if self.constraints["fix_probe"]:
-        #     return self.initial_probe
-        if self.constraints["orthogonalize_probe"]:
+        if self.constraints.orthogonalize_probe:
             probe = self._probe_orthogonalization_constraint(probe)
-        if self.constraints["center_probe"]:
+        if self.constraints.center_probe:
             probe = self._probe_center_of_mass_constraint(probe)
         return probe
 
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index c6f8ffd73..c6e0673bf 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -18,17 +18,61 @@
     is_distributed_launch,
     spawn_distributed_workers,
 )
-from quantem.diffractive_imaging.dataset_models import DatasetModelType
+from quantem.diffractive_imaging.dataset_models import (
+    DatasetModelType,
+    PtychoDatasetConstraintParams,
+    PtychoDatasetConstraintsType,
+)
 from quantem.diffractive_imaging.detector_models import DetectorModelType
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
-from quantem.diffractive_imaging.object_models import ObjectModelType, ObjectPixelated
-from quantem.diffractive_imaging.probe_models import ProbeModelType, ProbeParametric
+from quantem.diffractive_imaging.object_models import (
+    ObjectModelType,
+    ObjectPixelated,
+    PtychoObjConstraintParams,
+    PtychoObjConstraintsType,
+)
+from quantem.diffractive_imaging.probe_models import (
+    ProbeModelType,
+    ProbeParametric,
+    PtychoProbeConstraintParams,
+    PtychoProbeConstraintsType,
+)
 from quantem.diffractive_imaging.ptycho_utils import compute_train_val_split
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase
 from quantem.diffractive_imaging.ptychography_opt import PtychographyOpt
 from quantem.diffractive_imaging.ptychography_visualizations import PtychographyVisualizations
 
 
+def _merge_constraints(
+    constraints: dict[str, Any] | None,
+    obj_constraints: dict | PtychoObjConstraintsType | None,
+    probe_constraints: dict | PtychoProbeConstraintsType | None,
+    dset_constraints: dict | PtychoDatasetConstraintsType | None,
+) -> dict[str, Any]:
+    """Merge the legacy ``constraints`` dict with the new per-model kwargs.
+
+    Each new kwarg may be a plain dict (routed through the relevant ``parse_dict``)
+    or a Constraints dataclass instance. Passing the same model via both
+    ``constraints`` and a per-model kwarg raises ``ValueError``.
+    """
+    merged: dict[str, Any] = dict(constraints) if constraints else {}
+
+    def _set(slot: str, value, parser):
+        if value is None:
+            return
+        if slot in merged:
+            raise ValueError(
+                f"Constraints for '{slot}' provided via both `constraints=` and "
+                f"`{slot[:3]}_constraints=`; pass only one."
+            )
+        merged[slot] = parser(value) if isinstance(value, dict) else value
+
+    _set("object", obj_constraints, PtychoObjConstraintParams.parse_dict)
+    _set("probe", probe_constraints, PtychoProbeConstraintParams.parse_dict)
+    _set("dataset", dset_constraints, PtychoDatasetConstraintParams.parse_dict)
+    return merged
+
+
 def _ddp_ptycho_worker(
     rank: int,
     world_size: int,
@@ -202,7 +246,10 @@ def reconstruct(
         reset: bool = False,
         optimizer_params: dict[str, Any] | None = None,
         scheduler_params: dict[str, Any] | None = None,
-        constraints: dict[str, Any] = {},
+        constraints: dict[str, Any] | None = None,
+        obj_constraints: dict | PtychoObjConstraintsType | None = None,
+        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
+        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
@@ -221,12 +268,26 @@ def reconstruct(
           - ``int``              — specific GPU index, e.g. ``device=2`` → cuda:2
           - ``list[int]``        — multi-GPU, e.g. ``device=[0,1,2,3]``
 
+        Constraints can be set in two equivalent ways:
+
+        - ``constraints={"object": {...}, "probe": {...}, "dataset": {...}}`` — legacy form;
+          each leaf may be a dict or a ``Constraints`` dataclass instance.
+        - ``obj_constraints=``, ``probe_constraints=``, ``dset_constraints=`` — new form;
+          each accepts a dict (parsed via the relevant ``parse_dict``) or a dataclass
+          instance like ``PtychoObjConstraintParams.Raster(...)``.
+
+        Passing the same model via both forms raises ``ValueError``.
+
         Multi-GPU (``device`` is a list) launches worker processes via ``mp.spawn`` when called
         from a notebook, or uses the existing distributed process group when launched with
         ``torchrun``. Only autograd mode is supported for multi-GPU in this release.
         """
         self._check_preprocessed()
 
+        constraints = _merge_constraints(
+            constraints, obj_constraints, probe_constraints, dset_constraints
+        )
+
         # Determine effective device list: explicit arg takes priority, else fall back to stored.
         devices_to_use = (
             device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
@@ -291,7 +352,7 @@ def _reconstruct_inner(
         reset: bool = False,
         optimizer_params: dict[str, Any] | None = None,
         scheduler_params: dict[str, Any] | None = None,
-        constraints: dict[str, Any] = {},
+        constraints: dict[str, Any] | None = None,
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
@@ -314,7 +375,8 @@ def _reconstruct_inner(
 
         if reset:
             self.reset_recon()
-        self.constraints = constraints
+        if constraints:
+            self.constraints = constraints
 
         new_scheduler = reset
         if optimizer_params is not None:
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 835487e2d..617196f20 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -9,6 +9,7 @@
 
 from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
+from quantem.core.ml.constraints import Constraints
 from quantem.core.ml.dist_utils import all_reduce_params, worker_init_fn
 from quantem.core.utils.rng import RNGMixin
 from quantem.core.utils.utils import (
@@ -562,7 +563,12 @@ def constraints(self) -> dict[str, Any]:
 
     @constraints.setter
     def constraints(self, c: dict[str, Any]):
-        """Set constraints by forwarding to individual models."""
+        """Set constraints by forwarding to individual models.
+
+        Each leaf value may be either a plain ``dict`` (validated per-key against
+        the model's constraint dataclass) or a ``Constraints`` dataclass instance
+        (assigned wholesale to the model).
+        """
         constraint_handlers = {
             "object": self.obj_model,
             "probe": self.probe_model,
@@ -570,9 +576,16 @@ def constraints(self, c: dict[str, Any]):
         }
 
         for key, value in c.items():
-            if key in constraint_handlers and isinstance(value, dict):
-                for subkey, subvalue in value.items():
-                    constraint_handlers[key].add_constraint(subkey, subvalue)
+            if key in constraint_handlers:
+                if isinstance(value, Constraints):
+                    constraint_handlers[key].constraints = value
+                elif isinstance(value, dict):
+                    constraint_handlers[key].constraints = value
+                else:
+                    raise TypeError(
+                        f"Constraints for '{key}' must be a dict or Constraints dataclass, "
+                        f"got {type(value).__name__}"
+                    )
             elif key == "detector" and isinstance(value, dict):
                 warn("Detector constraints not implemented, skipping")
             else:
@@ -1000,9 +1013,7 @@ def _build_dataloaders(
                 seed=int(self.rng.integers(0, 2**31 - 1)),
                 drop_last=False,
             )
-            train_loader = DataLoader(
-                train_subset, sampler=train_sampler, **loader_kwargs
-            )
+            train_loader = DataLoader(train_subset, sampler=train_sampler, **loader_kwargs)
             if val_subset is not None:
                 val_sampler = DistributedSampler(
                     val_subset,
@@ -1011,9 +1022,7 @@ def _build_dataloaders(
                     shuffle=False,
                     drop_last=False,
                 )
-                val_loader = DataLoader(
-                    val_subset, sampler=val_sampler, **loader_kwargs
-                )
+                val_loader = DataLoader(val_subset, sampler=val_sampler, **loader_kwargs)
             else:
                 val_loader = None
         else:
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index e399412c6..4a980fa2e 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -9,11 +9,22 @@
 from quantem.core import config
 from quantem.core.datastructures import Dataset4dstem
 from quantem.core.ml.cnn import CNN2d
-from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster
+from quantem.diffractive_imaging.dataset_models import (
+    PtychoDatasetConstraintsType,
+    PtychographyDatasetRaster,
+)
 from quantem.diffractive_imaging.detector_models import DetectorPixelated
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
-from quantem.diffractive_imaging.object_models import ObjectDIP, ObjectPixelated
-from quantem.diffractive_imaging.probe_models import ProbeDIP, ProbePixelated
+from quantem.diffractive_imaging.object_models import (
+    ObjectDIP,
+    ObjectPixelated,
+    PtychoObjConstraintsType,
+)
+from quantem.diffractive_imaging.probe_models import (
+    ProbeDIP,
+    ProbePixelated,
+    PtychoProbeConstraintsType,
+)
 from quantem.diffractive_imaging.ptychography import Ptychography
 
 
@@ -178,7 +189,10 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
-        constraints: dict[str, Any] = {},  # TODO add constraints flags
+        constraints: dict[str, Any] | None = None,
+        obj_constraints: dict | PtychoObjConstraintsType | None = None,
+        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
+        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         store_iterations_every: int | None = None,
         device: "Literal['cpu', 'gpu'] | int | list[int] | None" = None,
         verbose: int | bool = True,
@@ -245,6 +259,9 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
             optimizer_params=opt_params,
             scheduler_params=scheduler_params,
             constraints=constraints,
+            obj_constraints=obj_constraints,
+            probe_constraints=probe_constraints,
+            dset_constraints=dset_constraints,
             batch_size=batch_size,
             store_snapshots_every=store_iterations_every,
             device=device,
@@ -419,7 +436,10 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         scheduler_type: Literal["exp", "cyclic", "plateau", "none"] = "none",
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
-        constraints: dict[str, Any] = {},  # TODO add constraints flags
+        constraints: dict[str, Any] | None = None,
+        obj_constraints: dict | PtychoObjConstraintsType | None = None,
+        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
+        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         store_iterations_every: int | None = None,
         device: Literal["cpu", "gpu"] | int | list[int] | None = None,
         verbose: int | bool = True,
@@ -486,6 +506,9 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
             optimizer_params=opt_params,
             scheduler_params=scheduler_params,
             constraints=constraints,
+            obj_constraints=obj_constraints,
+            probe_constraints=probe_constraints,
+            dset_constraints=dset_constraints,
             batch_size=batch_size,
             store_snapshots_every=store_iterations_every,
             device=device,
diff --git a/tests/diffractive_imaging/test_constraints.py b/tests/diffractive_imaging/test_constraints.py
new file mode 100644
index 000000000..775bdfb99
--- /dev/null
+++ b/tests/diffractive_imaging/test_constraints.py
@@ -0,0 +1,175 @@
+"""Tests for the ptychography constraint dataclass API."""
+
+import numpy as np
+import pytest
+
+from quantem.core.datastructures import Dataset4dstem
+from quantem.diffractive_imaging import (
+    DetectorPixelated,
+    ObjectPixelated,
+    ProbePixelated,
+    PtychoDatasetConstraintParams,
+    PtychoObjConstraintParams,
+    PtychoProbeConstraintParams,
+    Ptychography,
+    PtychographyDatasetRaster,
+)
+
+N_SCAN = 8
+N_DET = 16
+PROBE_ENERGY = 80e3
+PROBE_SEMIANGLE = 20
+PROBE_DEFOCUS = 100
+
+
+@pytest.fixture
+def ptycho():
+    rng = np.random.default_rng(42)
+    array = rng.random((N_SCAN, N_SCAN, N_DET, N_DET)).astype(np.float32)
+    dset = Dataset4dstem.from_array(
+        array,
+        name="test",
+        sampling=[1.0, 1.0, 0.05, 0.05],
+        units=["A", "A", "A^-1", "A^-1"],
+    )
+    pdset = PtychographyDatasetRaster.from_dataset4dstem(dset)
+    pdset.preprocess(com_fit_function="constant", plot_rotation=False, plot_com=False)
+    obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1)
+    probe = ProbePixelated.from_params(
+        probe_params={
+            "energy": PROBE_ENERGY,
+            "defocus": PROBE_DEFOCUS,
+            "semiangle_cutoff": PROBE_SEMIANGLE,
+        }
+    )
+    p = Ptychography.from_models(
+        dset=pdset,
+        obj_model=obj,
+        probe_model=probe,
+        detector_model=DetectorPixelated(),
+        verbose=False,
+        rng=42,
+    )
+    p.preprocess(obj_padding_px=(4, 4))
+    return p
+
+
+# --- parse_dict tests ---------------------------------------------------------
+
+
+class TestParseDict:
+    def test_object_raster_by_name(self):
+        c = PtychoObjConstraintParams.parse_dict({"name": "raster", "tv_weight_z": 5.0})
+        assert isinstance(c, PtychoObjConstraintParams.Raster)
+        assert c.tv_weight_z == 5.0
+        assert c.positivity is True  # default preserved
+
+    def test_object_inr_by_type(self):
+        c = PtychoObjConstraintParams.parse_dict({"type": "inr"})
+        assert isinstance(c, PtychoObjConstraintParams.INR)
+
+    def test_object_unknown_raises(self):
+        with pytest.raises(ValueError, match="Unknown object constraint type"):
+            PtychoObjConstraintParams.parse_dict({"name": "nope"})
+
+    def test_object_missing_name_raises(self):
+        with pytest.raises(ValueError, match="Must provide either 'name' or 'type'"):
+            PtychoObjConstraintParams.parse_dict({"tv_weight_z": 5.0})
+
+    def test_probe_raster_with_fields(self):
+        c = PtychoProbeConstraintParams.parse_dict(
+            {"name": "raster", "center_probe": True, "tv_weight": 0.1}
+        )
+        assert isinstance(c, PtychoProbeConstraintParams.Raster)
+        assert c.center_probe is True
+        assert c.tv_weight == 0.1
+
+    def test_dataset_raster_default(self):
+        c = PtychoDatasetConstraintParams.parse_dict({"name": "raster"})
+        assert isinstance(c, PtychoDatasetConstraintParams.Raster)
+        assert c.clip_scan_positions is True  # default preserved
+
+
+# --- Constraint typo catching -------------------------------------------------
+
+
+class TestTypoCatching:
+    def test_setting_unknown_field_via_dict_raises(self, ptycho):
+        with pytest.raises(KeyError, match="Invalid constraint key"):
+            ptycho.obj_model.constraints = {"not_a_real_field": True}
+
+    def test_add_constraint_unknown_key_raises(self, ptycho):
+        with pytest.raises(KeyError, match="Invalid constraint key"):
+            ptycho.obj_model.add_constraint("not_a_real_field", True)
+
+
+# --- Round-trip: pass dataclass via reconstruct(), read back through getter ---
+
+
+class TestRoundtrip:
+    def test_obj_constraints_dataclass(self, ptycho):
+        obj_c = PtychoObjConstraintParams.Raster(tv_weight_z=2.5, identical_slices=True)
+        ptycho.constraints = {"object": obj_c}
+        assert ptycho.obj_model.constraints is obj_c
+        assert ptycho.obj_model.constraints.tv_weight_z == 2.5
+        assert ptycho.obj_model.constraints.identical_slices is True
+
+    def test_probe_constraints_dataclass(self, ptycho):
+        probe_c = PtychoProbeConstraintParams.Raster(center_probe=True, tv_weight=0.05)
+        ptycho.constraints = {"probe": probe_c}
+        assert ptycho.probe_model.constraints is probe_c
+
+    def test_dataset_constraints_dataclass(self, ptycho):
+        dset_c = PtychoDatasetConstraintParams.Raster(descan_tv_weight=0.01)
+        ptycho.constraints = {"dataset": dset_c}
+        assert ptycho.dset.constraints is dset_c
+
+    def test_dict_form_still_works(self, ptycho):
+        """Backward compatibility: nested-dict form sets individual fields."""
+        ptycho.constraints = {
+            "object": {"tv_weight_z": 3.0, "positivity": False},
+            "probe": {"tv_weight": 0.02},
+        }
+        assert ptycho.obj_model.constraints.tv_weight_z == 3.0
+        assert ptycho.obj_model.constraints.positivity is False
+        assert ptycho.probe_model.constraints.tv_weight == 0.02
+
+
+# --- Reconstruct() kwargs and mutual exclusion --------------------------------
+
+
+class TestReconstructKwargs:
+    def test_obj_constraints_kwarg_applied(self, ptycho):
+        from quantem.core.ml import OptimizerParams
+
+        obj_c = PtychoObjConstraintParams.Raster(tv_weight_z=1.5)
+        ptycho.reconstruct(
+            num_iters=1,
+            reset=True,
+            optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)},
+            obj_constraints=obj_c,
+            batch_size=4,
+            device="cpu",
+        )
+        assert ptycho.obj_model.constraints.tv_weight_z == 1.5
+
+    def test_dict_kwarg_parsed(self, ptycho):
+        from quantem.core.ml import OptimizerParams
+
+        ptycho.reconstruct(
+            num_iters=1,
+            reset=True,
+            optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)},
+            obj_constraints={"name": "raster", "surface_zero_weight": 0.7},
+            batch_size=4,
+            device="cpu",
+        )
+        assert ptycho.obj_model.constraints.surface_zero_weight == 0.7
+
+    def test_mutual_exclusion_with_legacy_dict(self, ptycho):
+        with pytest.raises(ValueError, match="provided via both"):
+            ptycho.reconstruct(
+                num_iters=0,
+                constraints={"object": {"tv_weight_z": 1.0}},
+                obj_constraints=PtychoObjConstraintParams.Raster(tv_weight_z=2.0),
+            )

From e5207c8b9aab29ac470ae02e0f5e512a65c87ed0 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 27 May 2026 11:48:06 -0700
Subject: [PATCH 20/59] removing additional constraint flags from reconstruct

---
 src/quantem/core/ml/dist_utils.py             |  2 +-
 .../diffractive_imaging/dataset_models.py     |  4 +-
 .../diffractive_imaging/ptychography.py       | 71 +++----------------
 .../diffractive_imaging/ptychography_lite.py  | 29 +-------
 tests/diffractive_imaging/test_constraints.py | 39 ++++++----
 5 files changed, 41 insertions(+), 104 deletions(-)

diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
index 733e5be2e..856c84d25 100644
--- a/src/quantem/core/ml/dist_utils.py
+++ b/src/quantem/core/ml/dist_utils.py
@@ -81,7 +81,7 @@ def spawn_distributed_workers(
     """
     import torch.multiprocessing as mp
 
-    mp.start_processes( # type: ignore 
+    mp.start_processes(  # type: ignore
         worker_fn,
         args=(len(devices), *worker_args),
         nprocs=len(devices),
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 5b5b49b4f..d3c5935df 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -661,7 +661,9 @@ def reset(self) -> None:
 class DatasetConstraints(
     BaseConstraints[PtychoDatasetConstraintParams.Raster], PtychographyDatasetBase
 ):
-    DEFAULT_CONSTRAINTS: PtychoDatasetConstraintParams.Raster = PtychoDatasetConstraintParams.Raster()
+    DEFAULT_CONSTRAINTS: PtychoDatasetConstraintParams.Raster = (
+        PtychoDatasetConstraintParams.Raster()
+    )
 
     def apply_soft_constraints(self, descan_shifts: torch.Tensor) -> torch.Tensor:
         self.reset_soft_constraint_losses()
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index c6e0673bf..05e0b5af3 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -18,61 +18,17 @@
     is_distributed_launch,
     spawn_distributed_workers,
 )
-from quantem.diffractive_imaging.dataset_models import (
-    DatasetModelType,
-    PtychoDatasetConstraintParams,
-    PtychoDatasetConstraintsType,
-)
+from quantem.diffractive_imaging.dataset_models import DatasetModelType
 from quantem.diffractive_imaging.detector_models import DetectorModelType
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
-from quantem.diffractive_imaging.object_models import (
-    ObjectModelType,
-    ObjectPixelated,
-    PtychoObjConstraintParams,
-    PtychoObjConstraintsType,
-)
-from quantem.diffractive_imaging.probe_models import (
-    ProbeModelType,
-    ProbeParametric,
-    PtychoProbeConstraintParams,
-    PtychoProbeConstraintsType,
-)
+from quantem.diffractive_imaging.object_models import ObjectModelType, ObjectPixelated
+from quantem.diffractive_imaging.probe_models import ProbeModelType, ProbeParametric
 from quantem.diffractive_imaging.ptycho_utils import compute_train_val_split
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase
 from quantem.diffractive_imaging.ptychography_opt import PtychographyOpt
 from quantem.diffractive_imaging.ptychography_visualizations import PtychographyVisualizations
 
 
-def _merge_constraints(
-    constraints: dict[str, Any] | None,
-    obj_constraints: dict | PtychoObjConstraintsType | None,
-    probe_constraints: dict | PtychoProbeConstraintsType | None,
-    dset_constraints: dict | PtychoDatasetConstraintsType | None,
-) -> dict[str, Any]:
-    """Merge the legacy ``constraints`` dict with the new per-model kwargs.
-
-    Each new kwarg may be a plain dict (routed through the relevant ``parse_dict``)
-    or a Constraints dataclass instance. Passing the same model via both
-    ``constraints`` and a per-model kwarg raises ``ValueError``.
-    """
-    merged: dict[str, Any] = dict(constraints) if constraints else {}
-
-    def _set(slot: str, value, parser):
-        if value is None:
-            return
-        if slot in merged:
-            raise ValueError(
-                f"Constraints for '{slot}' provided via both `constraints=` and "
-                f"`{slot[:3]}_constraints=`; pass only one."
-            )
-        merged[slot] = parser(value) if isinstance(value, dict) else value
-
-    _set("object", obj_constraints, PtychoObjConstraintParams.parse_dict)
-    _set("probe", probe_constraints, PtychoProbeConstraintParams.parse_dict)
-    _set("dataset", dset_constraints, PtychoDatasetConstraintParams.parse_dict)
-    return merged
-
-
 def _ddp_ptycho_worker(
     rank: int,
     world_size: int,
@@ -247,9 +203,6 @@ def reconstruct(
         optimizer_params: dict[str, Any] | None = None,
         scheduler_params: dict[str, Any] | None = None,
         constraints: dict[str, Any] | None = None,
-        obj_constraints: dict | PtychoObjConstraintsType | None = None,
-        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
-        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
@@ -268,15 +221,13 @@ def reconstruct(
           - ``int``              — specific GPU index, e.g. ``device=2`` → cuda:2
           - ``list[int]``        — multi-GPU, e.g. ``device=[0,1,2,3]``
 
-        Constraints can be set in two equivalent ways:
+        ``constraints`` is a dict keyed by ``"object"``, ``"probe"``, ``"dataset"``
+        (any subset). Each leaf may be:
 
-        - ``constraints={"object": {...}, "probe": {...}, "dataset": {...}}`` — legacy form;
-          each leaf may be a dict or a ``Constraints`` dataclass instance.
-        - ``obj_constraints=``, ``probe_constraints=``, ``dset_constraints=`` — new form;
-          each accepts a dict (parsed via the relevant ``parse_dict``) or a dataclass
-          instance like ``PtychoObjConstraintParams.Raster(...)``.
-
-        Passing the same model via both forms raises ``ValueError``.
+        - a ``Constraints`` dataclass instance (e.g. ``PtychoObjConstraintParams.Raster(...)``),
+          which replaces that model's constraint state wholesale, or
+        - a plain ``dict`` of field-name -> value, which does a per-key partial update
+          on the existing constraint state.
 
         Multi-GPU (``device`` is a list) launches worker processes via ``mp.spawn`` when called
         from a notebook, or uses the existing distributed process group when launched with
@@ -284,10 +235,6 @@ def reconstruct(
         """
         self._check_preprocessed()
 
-        constraints = _merge_constraints(
-            constraints, obj_constraints, probe_constraints, dset_constraints
-        )
-
         # Determine effective device list: explicit arg takes priority, else fall back to stored.
         devices_to_use = (
             device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 4a980fa2e..0bca7a494 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -9,22 +9,11 @@
 from quantem.core import config
 from quantem.core.datastructures import Dataset4dstem
 from quantem.core.ml.cnn import CNN2d
-from quantem.diffractive_imaging.dataset_models import (
-    PtychoDatasetConstraintsType,
-    PtychographyDatasetRaster,
-)
+from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster
 from quantem.diffractive_imaging.detector_models import DetectorPixelated
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
-from quantem.diffractive_imaging.object_models import (
-    ObjectDIP,
-    ObjectPixelated,
-    PtychoObjConstraintsType,
-)
-from quantem.diffractive_imaging.probe_models import (
-    ProbeDIP,
-    ProbePixelated,
-    PtychoProbeConstraintsType,
-)
+from quantem.diffractive_imaging.object_models import ObjectDIP, ObjectPixelated
+from quantem.diffractive_imaging.probe_models import ProbeDIP, ProbePixelated
 from quantem.diffractive_imaging.ptychography import Ptychography
 
 
@@ -190,9 +179,6 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
         constraints: dict[str, Any] | None = None,
-        obj_constraints: dict | PtychoObjConstraintsType | None = None,
-        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
-        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         store_iterations_every: int | None = None,
         device: "Literal['cpu', 'gpu'] | int | list[int] | None" = None,
         verbose: int | bool = True,
@@ -259,9 +245,6 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
             optimizer_params=opt_params,
             scheduler_params=scheduler_params,
             constraints=constraints,
-            obj_constraints=obj_constraints,
-            probe_constraints=probe_constraints,
-            dset_constraints=dset_constraints,
             batch_size=batch_size,
             store_snapshots_every=store_iterations_every,
             device=device,
@@ -437,9 +420,6 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         scheduler_factor: float = 0.5,
         new_optimizers: bool = False,  # not sure what the default should be
         constraints: dict[str, Any] | None = None,
-        obj_constraints: dict | PtychoObjConstraintsType | None = None,
-        probe_constraints: dict | PtychoProbeConstraintsType | None = None,
-        dset_constraints: dict | PtychoDatasetConstraintsType | None = None,
         store_iterations_every: int | None = None,
         device: Literal["cpu", "gpu"] | int | list[int] | None = None,
         verbose: int | bool = True,
@@ -506,9 +486,6 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
             optimizer_params=opt_params,
             scheduler_params=scheduler_params,
             constraints=constraints,
-            obj_constraints=obj_constraints,
-            probe_constraints=probe_constraints,
-            dset_constraints=dset_constraints,
             batch_size=batch_size,
             store_snapshots_every=store_iterations_every,
             device=device,
diff --git a/tests/diffractive_imaging/test_constraints.py b/tests/diffractive_imaging/test_constraints.py
index 775bdfb99..9cce6639b 100644
--- a/tests/diffractive_imaging/test_constraints.py
+++ b/tests/diffractive_imaging/test_constraints.py
@@ -9,10 +9,10 @@
     ObjectPixelated,
     ProbePixelated,
     PtychoDatasetConstraintParams,
-    PtychoObjConstraintParams,
-    PtychoProbeConstraintParams,
     Ptychography,
     PtychographyDatasetRaster,
+    PtychoObjConstraintParams,
+    PtychoProbeConstraintParams,
 )
 
 N_SCAN = 8
@@ -135,11 +135,11 @@ def test_dict_form_still_works(self, ptycho):
         assert ptycho.probe_model.constraints.tv_weight == 0.02
 
 
-# --- Reconstruct() kwargs and mutual exclusion --------------------------------
+# --- Reconstruct() with constraints= ------------------------------------------
 
 
 class TestReconstructKwargs:
-    def test_obj_constraints_kwarg_applied(self, ptycho):
+    def test_dataclass_leaf_applied(self, ptycho):
         from quantem.core.ml import OptimizerParams
 
         obj_c = PtychoObjConstraintParams.Raster(tv_weight_z=1.5)
@@ -147,29 +147,40 @@ def test_obj_constraints_kwarg_applied(self, ptycho):
             num_iters=1,
             reset=True,
             optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)},
-            obj_constraints=obj_c,
+            constraints={"object": obj_c},
             batch_size=4,
             device="cpu",
         )
         assert ptycho.obj_model.constraints.tv_weight_z == 1.5
 
-    def test_dict_kwarg_parsed(self, ptycho):
+    def test_dict_leaf_partial_update(self, ptycho):
         from quantem.core.ml import OptimizerParams
 
         ptycho.reconstruct(
             num_iters=1,
             reset=True,
             optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)},
-            obj_constraints={"name": "raster", "surface_zero_weight": 0.7},
+            constraints={"object": {"surface_zero_weight": 0.7}},
             batch_size=4,
             device="cpu",
         )
         assert ptycho.obj_model.constraints.surface_zero_weight == 0.7
+        # other fields keep their defaults
+        assert ptycho.obj_model.constraints.positivity is True
 
-    def test_mutual_exclusion_with_legacy_dict(self, ptycho):
-        with pytest.raises(ValueError, match="provided via both"):
-            ptycho.reconstruct(
-                num_iters=0,
-                constraints={"object": {"tv_weight_z": 1.0}},
-                obj_constraints=PtychoObjConstraintParams.Raster(tv_weight_z=2.0),
-            )
+    def test_mixed_dataclass_and_dict_leaves(self, ptycho):
+        from quantem.core.ml import OptimizerParams
+
+        ptycho.reconstruct(
+            num_iters=1,
+            reset=True,
+            optimizer_params={"object": OptimizerParams.Adam(lr=1e-2)},
+            constraints={
+                "object": PtychoObjConstraintParams.Raster(tv_weight_xy=0.4),
+                "probe": {"center_probe": True},
+            },
+            batch_size=4,
+            device="cpu",
+        )
+        assert ptycho.obj_model.constraints.tv_weight_xy == 0.4
+        assert ptycho.probe_model.constraints.center_probe is True

From 10d353812e007375f08c3c776fdca91443d16fa3 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 27 May 2026 12:15:22 -0700
Subject: [PATCH 21/59] improving constraint params docstrings, parse_dict

---
 src/quantem/core/ml/constraints.py            | 41 ++++++++++
 .../diffractive_imaging/dataset_models.py     | 48 +++++++----
 .../diffractive_imaging/object_models.py      | 80 ++++++++++++++-----
 .../diffractive_imaging/probe_models.py       | 55 ++++++++-----
 4 files changed, 167 insertions(+), 57 deletions(-)

diff --git a/src/quantem/core/ml/constraints.py b/src/quantem/core/ml/constraints.py
index 0adf6b57a..59f9b14ef 100644
--- a/src/quantem/core/ml/constraints.py
+++ b/src/quantem/core/ml/constraints.py
@@ -49,6 +49,47 @@ def __str__(self) -> str:
         )
 
 
+def parse_constraint_dict(
+    namespace: type,
+    d: dict,
+    *,
+    kind: str = "constraint",
+) -> Constraints:
+    """Dispatch a config dict to one of ``namespace``'s nested ``Constraints`` variants.
+
+    ``namespace`` is a class with one or more nested ``@dataclass``\\ -decorated
+    ``Constraints`` subclasses. The dict must contain a ``"name"`` or ``"type"`` key
+    whose value (case-insensitive) matches one variant's ``_name`` field; the
+    remaining keys are forwarded as constructor kwargs to that variant.
+
+    ``kind`` is a short human-readable label ("object", "probe", "dataset", ...)
+    used only in error messages.
+    """
+    d = dict(d)
+    name = d.pop("name", None) or d.pop("type", None)
+    if name is None:
+        raise ValueError(f"Must provide either 'name' or 'type' key for {kind} constraints")
+    if isinstance(name, type):
+        name = name.__name__.lower()
+    elif isinstance(name, str):
+        name = name.lower()
+    else:
+        raise ValueError(f"Unknown {kind} constraint type: {name!r}")
+
+    variants: dict[str, type[Constraints]] = {}
+    for attr in vars(namespace).values():
+        if isinstance(attr, type) and issubclass(attr, Constraints) and attr is not Constraints:
+            variant_name = getattr(attr, "_name", None)
+            if isinstance(variant_name, str):
+                variants[variant_name.lower()] = attr
+
+    if name not in variants:
+        raise ValueError(
+            f"Unknown {kind} constraint type: {name!r}; expected one of {sorted(variants)}"
+        )
+    return variants[name](**d)
+
+
 C = TypeVar("C", bound=Constraints)
 
 
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index d3c5935df..7b538e488 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1,7 +1,7 @@
 from abc import abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal, Self
+from typing import Any, Literal, Self, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -13,7 +13,7 @@
 from quantem.core.datastructures.dataset3d import Dataset3d
 from quantem.core.datastructures.dataset4dstem import Dataset4dstem
 from quantem.core.io.serialize import AutoSerialize
-from quantem.core.ml.constraints import BaseConstraints, Constraints
+from quantem.core.ml.constraints import BaseConstraints, Constraints, parse_constraint_dict
 from quantem.core.ml.optimizer_mixin import OptimizerMixin
 from quantem.core.utils.utils import electron_wavelength_angstrom, tqdmnd
 from quantem.core.utils.validators import (
@@ -47,7 +47,28 @@ class PtychoDatasetConstraintParams:
 
     @dataclass
     class Raster(Constraints):
-        """Constraints for raster-scan ptychography datasets."""
+        """Constraints for raster-scan ptychography datasets (``PtychographyDatasetRaster``).
+
+        Attributes
+        ----------
+        descan_shifts_constant : bool, default ``False``
+            Forces all descan shifts to zero after each update. Useful when you
+            want to keep the descan optimizer in the parameter group but freeze
+            its effect.
+        center_scan_positions : bool, default ``False``
+            Shifts all scan positions uniformly so their mean sits at the object
+            center after each update. Prevents the reconstruction from
+            translating during long runs with ``lr_scan_positions > 0``.
+        clip_scan_positions : bool, default ``True``
+            Clamps scan positions to lie within ``[0, obj_shape - 1]`` after each
+            update. On by default to prevent positions from drifting off the
+            padded object during refinement.
+        descan_tv_weight : float, default ``0.0``
+            Soft penalty. Weight on the total-variation of the descan-shift
+            sequence (x and y averaged). Encourages smoothly varying descan, but
+            only contributes when ``learn_descan`` is on and the dataset has a
+            descan optimizer attached.
+        """
 
         # hard constraints
         descan_shifts_constant: bool = False
@@ -66,20 +87,13 @@ class Raster(Constraints):
 
     @classmethod
     def parse_dict(cls, d: dict) -> "PtychoDatasetConstraintsType":
-        d = dict(d)
-        name = d.pop("name", None) or d.pop("type", None)
-        if name is None:
-            raise ValueError("Must provide either 'name' or 'type' key")
-        if isinstance(name, type):
-            name = name.__name__.lower()
-        elif isinstance(name, str):
-            name = name.lower()
-        else:
-            raise ValueError(f"Unknown dataset constraint type: {name}")
-        if name == "raster":
-            return cls.Raster(**d)
-        else:
-            raise ValueError(f"Unknown dataset constraint type: {name}")
+        """Instantiate the appropriate variant from a config dict.
+
+        The dict must contain a ``'name'`` or ``'type'`` key (case-insensitive),
+        with value ``'raster'``. All other keys are forwarded as keyword
+        arguments to the chosen dataclass.
+        """
+        return cast(PtychoDatasetConstraintsType, parse_constraint_dict(cls, d, kind="dataset"))
 
 
 PtychoDatasetConstraintsType = PtychoDatasetConstraintParams.Raster
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 0fd00c69c..b1e1a65da 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -14,7 +14,7 @@
 from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.blocks import reset_weights
-from quantem.core.ml.constraints import BaseConstraints, Constraints
+from quantem.core.ml.constraints import BaseConstraints, Constraints, parse_constraint_dict
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType
 from quantem.core.utils.rng import RNGMixin
@@ -54,7 +54,56 @@ class PtychoObjConstraintParams:
 
     @dataclass
     class Raster(Constraints):
-        """Constraints for grid-based ptychography object models (Pixelated, DIP)."""
+        """Constraints for grid-based ptychography object models (``ObjectPixelated``,
+        ``ObjectDIP``).
+
+        Fields are applied each iteration in two flavors: **hard** constraints
+        project / filter the object after the optimizer step; **soft** constraints
+        add a penalty term to the training loss.
+
+        Attributes
+        ----------
+        positivity : bool, default ``True``
+            Clamps the object to be non-negative after each update.
+            Only consulted when ``obj_type="potential"``; for ``"complex"`` /
+            ``"pure_phase"`` the amplitude is clamped to ``[0, 1]`` (or fixed to 1)
+            regardless of this flag.
+        fix_potential_baseline : bool, default ``False``
+            ``obj_type="potential"`` only. Subtracts an offset from the object so
+            background regions sit at zero. If an FOV mask is set the offset is
+            the mean of the background (``mask < 0.5 * mask.max()``); otherwise
+            it's ``obj.min()``.
+        fix_potential_baseline_factor : float, default ``1.0``
+            Scales the baseline offset. Values ``<1`` relax the anchoring
+            (subtract less of the background); ``>1`` over-correct.
+        identical_slices : bool, default ``False``
+            Multislice (``num_slices > 1``) only. Replaces every slice with the
+            mean across slices, forcing an effectively 2D object.
+        apply_fov_mask : bool, default ``False``
+            Multiplies the object by the precomputed FOV mask after each update.
+            Useful when the scan does not cover the full padded object area.
+        gaussian_sigma : float | None, default ``None``
+            Standard deviation (in pixels) of a 2D Gaussian blur applied to each
+            slice after each update. Smoothing prior; ``None`` disables.
+        butterworth_order : int, default ``4``
+            Order of the Butterworth filter used by ``q_lowpass`` / ``q_highpass``.
+        q_lowpass : float | None, default ``None``
+            Lowpass cutoff in inverse Angstroms. Fourier components above this
+            spatial frequency are suppressed via a Butterworth filter.
+        q_highpass : float | None, default ``None``
+            Highpass cutoff in inverse Angstroms. Components below this frequency
+            are suppressed; typically used to remove a slowly varying background.
+        tv_weight_z : float, default ``0.0``
+            Soft penalty. Weight on the depth-axis total-variation term in the
+            loss. Multislice (``num_slices > 1``) only.
+        tv_weight_xy : float, default ``0.0``
+            Soft penalty. Weight on the in-plane total-variation term;
+            encourages piecewise-smooth regions while preserving edges.
+        surface_zero_weight : float, default ``0.0``
+            Soft penalty pulling the first and last slices toward zero. Useful
+            for thick samples embedded in vacuum. Multislice only and requires
+            ``num_slices >= 3``.
+        """
 
         # hard constraints
         positivity: bool = True
@@ -88,7 +137,13 @@ class Raster(Constraints):
 
     @dataclass
     class INR(Constraints):
-        """Placeholder for the upcoming ObjectINR variant. Not yet wired to a model."""
+        """Placeholder for the upcoming ``ObjectINR`` variant.
+
+        INR-specific constraints (e.g. sparsity / TV penalties evaluated at
+        sampled coordinates) will land here when the model is implemented.
+        Until then this exists so ``parse_dict`` accepts ``"inr"`` and downstream
+        code can pattern-match on the variant.
+        """
 
         _name: str = "inr"
 
@@ -97,28 +152,13 @@ class INR(Constraints):
 
     @classmethod
     def parse_dict(cls, d: dict) -> "PtychoObjConstraintsType":
-        """Instantiate the appropriate Raster or INR dataclass from a config dict.
+        """Instantiate the appropriate variant from a config dict.
 
         The dict must contain a ``'name'`` or ``'type'`` key (case-insensitive),
         with value ``'raster'`` or ``'inr'``. All other keys are forwarded as
         keyword arguments to the chosen dataclass.
         """
-        d = dict(d)
-        name = d.pop("name", None) or d.pop("type", None)
-        if name is None:
-            raise ValueError("Must provide either 'name' or 'type' key")
-        if isinstance(name, type):
-            name = name.__name__.lower()
-        elif isinstance(name, str):
-            name = name.lower()
-        else:
-            raise ValueError(f"Unknown object constraint type: {name}")
-        if name == "raster":
-            return cls.Raster(**d)
-        elif name == "inr":
-            return cls.INR(**d)
-        else:
-            raise ValueError(f"Unknown object constraint type: {name}")
+        return cast(PtychoObjConstraintsType, parse_constraint_dict(cls, d, kind="object"))
 
 
 PtychoObjConstraintsType = PtychoObjConstraintParams.Raster | PtychoObjConstraintParams.INR
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index c8666fb86..1ac76c540 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -1,7 +1,7 @@
 from abc import abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Callable, Self, Union
+from typing import Any, Callable, Self, Union, cast
 from warnings import warn
 
 import matplotlib.pyplot as plt
@@ -15,7 +15,7 @@
 from quantem.core.datastructures import Dataset2d, Dataset4dstem
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.blocks import reset_weights
-from quantem.core.ml.constraints import BaseConstraints, Constraints
+from quantem.core.ml.constraints import BaseConstraints, Constraints, parse_constraint_dict
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType
 from quantem.core.utils.rng import RNGMixin
@@ -62,7 +62,25 @@ class PtychoProbeConstraintParams:
 
     @dataclass
     class Raster(Constraints):
-        """Constraints for grid-based ptychography probe models (Pixelated, DIP)."""
+        """Constraints for grid-based ptychography probe models (``ProbePixelated``,
+        ``ProbeDIP``).
+
+        Attributes
+        ----------
+        orthogonalize_probe : bool, default ``True``
+            Mixed-state probe (``num_probes > 1``) only. After each update applies
+            Gram-Schmidt orthogonalization across the probe stack and then sorts
+            the resulting probes by total intensity (descending). For
+            ``num_probes == 1`` this is effectively a renormalization no-op.
+        center_probe : bool, default ``False``
+            Shifts the probe's intensity center-of-mass back to the array center
+            via a Fourier shift after each update. Useful when probe drift
+            competes with scan-position refinement; if both move freely the
+            reconstruction can wander while still fitting the diffraction data.
+        tv_weight : float, default ``0.0``
+            Soft penalty. Weight on the in-plane total-variation of the (complex)
+            probe; encourages smooth probe magnitude / phase.
+        """
 
         # hard constraints
         orthogonalize_probe: bool = True
@@ -76,7 +94,13 @@ class Raster(Constraints):
 
     @dataclass
     class Parametric(Constraints):
-        """Placeholder for parametric probe constraints. Not yet wired to a model."""
+        """Placeholder for parametric probe constraints (``ProbeParametric``).
+
+        Parametric probes are pure functions of aberration / aperture coefficients,
+        so pixel-domain projections like ``orthogonalize_probe`` and ``tv_weight``
+        don't apply. Parametric-specific fields (e.g. bounds on individual
+        aberration coefficients) will land here when needed.
+        """
 
         _name: str = "parametric"
 
@@ -85,22 +109,13 @@ class Parametric(Constraints):
 
     @classmethod
     def parse_dict(cls, d: dict) -> "PtychoProbeConstraintsType":
-        d = dict(d)
-        name = d.pop("name", None) or d.pop("type", None)
-        if name is None:
-            raise ValueError("Must provide either 'name' or 'type' key")
-        if isinstance(name, type):
-            name = name.__name__.lower()
-        elif isinstance(name, str):
-            name = name.lower()
-        else:
-            raise ValueError(f"Unknown probe constraint type: {name}")
-        if name == "raster":
-            return cls.Raster(**d)
-        elif name == "parametric":
-            return cls.Parametric(**d)
-        else:
-            raise ValueError(f"Unknown probe constraint type: {name}")
+        """Instantiate the appropriate variant from a config dict.
+
+        The dict must contain a ``'name'`` or ``'type'`` key (case-insensitive),
+        with value ``'raster'`` or ``'parametric'``. All other keys are forwarded
+        as keyword arguments to the chosen dataclass.
+        """
+        return cast(PtychoProbeConstraintsType, parse_constraint_dict(cls, d, kind="probe"))
 
 
 PtychoProbeConstraintsType = (

From e3d15ae562d04fce377b5082087cc03d4a51e438 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 27 May 2026 18:53:45 -0700
Subject: [PATCH 22/59] changing pure_phase to be unwrapped real values

---
 .../diffractive_imaging/object_models.py      | 225 ++++++++++--------
 .../optimize_hyperparameters.py               |   3 +-
 .../diffractive_imaging/ptychography_base.py  |  28 ++-
 .../ptychography_visualizations.py            |  18 +-
 tests/diffractive_imaging/test_constraints.py | 103 ++++++++
 5 files changed, 268 insertions(+), 109 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index b1e1a65da..7d85d8b95 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -164,11 +164,14 @@ def parse_dict(cls, d: dict) -> "PtychoObjConstraintsType":
 PtychoObjConstraintsType = PtychoObjConstraintParams.Raster | PtychoObjConstraintParams.INR
 
 """
-Currently all object models.obj are complex valued for "complex" or "pure_phase" object types,
-and real valued for "potential" object types. This could be changed to be always complex valued, 
-(after applying constraints) as currently the real-valued potential is made complex in get_obj_patches, 
-which will not be used for implicit NNs, which leads to an inconsistency. Leaving for now as I'm not 
-sure if this would lead to other issues, so a bit of testing will be needed.
+Object representation by obj_type:
+- "complex"    : _obj is complex (amplitude * exp(1j * phase))
+- "pure_phase" : _obj is a real, unwrapped phase array
+- "potential"  : _obj is a real potential array
+
+The forward boundary (`_get_obj_patches`) wraps real `_obj` to `exp(1j * _obj)` for
+both pure_phase and potential, so the rest of the forward model never has to
+branch on obj_type.
 """
 
 
@@ -220,10 +223,9 @@ def shape_2d(self) -> tuple[int, int]:
 
     @property
     def dtype(self) -> "torch.dtype":
-        if self.obj_type == "potential":
-            return getattr(torch, config.get("dtype_real"))
-        else:
+        if self.obj_type == "complex":
             return getattr(torch, config.get("dtype_complex"))
+        return getattr(torch, config.get("dtype_real"))
 
     @property
     def device(self) -> str:
@@ -413,53 +415,70 @@ def apply_hard_constraints(
         self, obj: torch.Tensor, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
         c = self.constraints
-        if self.obj_type in ["complex", "pure_phase"]:
-            if self.obj_type == "complex":
-                amp = torch.clamp(torch.abs(obj), 0.0, 1.0)
-            else:
-                amp = 1.0
-            phase = obj.angle() - obj.angle().mean()
-            if mask is not None and c.apply_fov_mask:
-                obj2 = amp * mask * torch.exp(1.0j * phase * mask)
-            else:
-                obj2 = amp * torch.exp(1.0j * phase)
+        if self.obj_type == "complex":
+            obj2 = self._apply_hard_complex(obj, c)
+        elif self.obj_type == "pure_phase":
+            obj2 = self._apply_hard_pure_phase(obj, c)
         else:  # potential
-            if c.fix_potential_baseline:
-                if mask is not None:
-                    background = mask < 0.5 * mask.max()
-                    if background.any():
-                        offset = obj[background].mean()
-                    else:
-                        offset = obj.min()
+            obj2 = self._apply_hard_potential(obj, c, mask)
+        return self._apply_shared_hard(obj2, c, mask)
+
+    def _apply_hard_complex(
+        self, obj: torch.Tensor, c: PtychoObjConstraintParams.Raster
+    ) -> torch.Tensor:
+        amp = torch.clamp(torch.abs(obj), 0.0, 1.0)
+        phase = obj.angle() - obj.angle().mean()
+        return amp * torch.exp(1.0j * phase)
+
+    def _apply_hard_pure_phase(
+        self, obj: torch.Tensor, c: PtychoObjConstraintParams.Raster
+    ) -> torch.Tensor:
+        # phase stored directly as a real tensor; recenter to zero mean
+        return obj - obj.mean()
+
+    def _apply_hard_potential(
+        self,
+        obj: torch.Tensor,
+        c: PtychoObjConstraintParams.Raster,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if c.fix_potential_baseline:
+            if mask is not None:
+                background = mask < 0.5 * mask.max()
+                if background.any():
+                    offset = obj[background].mean()
                 else:
                     offset = obj.min()
-                offset = offset.detach()
-                offset *= c.fix_potential_baseline_factor
             else:
-                offset = 0
+                offset = obj.min()
+            offset = offset.detach()
+            offset = offset * c.fix_potential_baseline_factor
+        else:
+            offset = 0
 
-            if c.positivity:
-                obj2 = torch.clamp(obj - offset, min=0.0)
-            else:
-                obj2 = obj - offset
+        if c.positivity:
+            return torch.clamp(obj - offset, min=0.0)
+        return obj - offset
 
+    def _apply_shared_hard(
+        self,
+        obj: torch.Tensor,
+        c: PtychoObjConstraintParams.Raster,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor:
         if c.apply_fov_mask and mask is not None:
-            obj2 *= mask
+            obj = obj * mask
 
         if c.gaussian_sigma is not None:
-            obj2 = self.gaussian_blur_2d(obj2, sigma=c.gaussian_sigma)
+            obj = self.gaussian_blur_2d(obj, sigma=c.gaussian_sigma)
 
         if any([c.q_lowpass, c.q_highpass]):
-            obj2 = self.butterworth_constraint(
-                obj2,
-                sampling=self.sampling,
-            )
-        if self.num_slices > 1:
-            if c.identical_slices:
-                with torch.no_grad():
-                    obj2[:] = torch.mean(obj2, dim=0, keepdim=True)
+            obj = self.butterworth_constraint(obj, sampling=self.sampling)
 
-        return obj2
+        if self.num_slices > 1 and c.identical_slices:
+            with torch.no_grad():
+                obj[:] = torch.mean(obj, dim=0, keepdim=True)
+        return obj
 
     def apply_soft_constraints(
         self, obj: torch.Tensor, mask: torch.Tensor | None = None
@@ -484,38 +503,46 @@ def get_tv_loss(
         self, array: torch.Tensor, weights: None | tuple[float, float] = None
     ) -> torch.Tensor:
         loss = self._get_zero_loss_tensor()
+        w = self._resolve_tv_weights(weights)
+        if not any(w):
+            return loss
+
+        if self.obj_type == "complex":
+            return self._tv_complex(array, w)
+        # pure_phase and potential are both real tensors; phase wrapping is gone.
+        return self._calc_tv_loss(array, w)
+
+    def _resolve_tv_weights(
+        self, weights: None | tuple[float, float] | float | int
+    ) -> tuple[float, float]:
         if weights is None:
-            w = (
+            w: tuple[float, float] = (
                 self.constraints.tv_weight_z,
                 self.constraints.tv_weight_xy,
             )
         elif isinstance(weights, (float, int)):
-            if weights == 0:
-                return loss
-            w = (weights, weights)
+            w = (float(weights), float(weights))
         else:
             if len(weights) != 2:
                 raise ValueError(f"weights must be a tuple of length 2, got {weights}")
-            w = weights
-
-        if not any(w):
-            return loss
-
+            w = (float(weights[0]), float(weights[1]))
         if self.num_slices == 1:
-            w = (0, w[1])
-
-        if array.is_complex():
-            ph = array.angle()
-            warn(
-                "calculating TV loss for phase, need to check phase wrapping. Easiest fix is scalar phase array."
-            )
-            loss = loss + self._calc_tv_loss(ph, w)
-            amp = array.abs()
-            if self.obj_type == "complex":
-                loss = loss + self._calc_tv_loss(amp, w)
-        else:
-            loss = loss + self._calc_tv_loss(array, w)
+            w = (0.0, w[1])
+        return w
 
+    def _tv_complex(self, array: torch.Tensor, w: tuple[float, float]) -> torch.Tensor:
+        # complex objects carry information in both amplitude and phase. We
+        # still extract phase via angle() here, so the wrap warning stays —
+        # but only for obj_type == "complex".
+        loss = self._get_zero_loss_tensor()
+        ph = array.angle()
+        warn(
+            "calculating TV loss for phase of complex object, "
+            "phase wrapping may distort the gradient. Consider obj_type='pure_phase'."
+        )
+        loss = loss + self._calc_tv_loss(ph, w)
+        amp = array.abs()
+        loss = loss + self._calc_tv_loss(amp, w)
         return loss
 
     def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> torch.Tensor:
@@ -537,24 +564,27 @@ def get_surface_zero_loss(
         self, array: torch.Tensor, weight: float | int = 0.0
     ) -> torch.Tensor:
         loss = self._get_zero_loss_tensor()
-        if weight == 0:
-            return loss
-        if array.shape[0] < 3:
+        if weight == 0 or array.shape[0] < 3:
             return loss
-        if array.is_complex():
-            ph = array.angle().abs()
-            if self.obj_type == "complex":
-                amp = array.abs()
-                loss = loss + weight * (torch.mean(1.0 - amp[0]) + torch.mean(1.0 - amp[-1]))
-            warn("calculating surface zero loss for phase, need to check phase wrapping.")
-            loss = loss + weight * (
-                torch.mean(torch.abs(ph[0] - ph[0].mean()))
-                + torch.mean(torch.abs(ph[-1] - ph[-1].mean()))
-            )
-        else:
-            loss = loss + weight * (
-                torch.mean(torch.abs(array[0])) + torch.mean(torch.abs(array[-1]))
-            )
+        if self.obj_type == "complex":
+            return self._surface_zero_complex(array, weight)
+        # pure_phase and potential: real array, penalize first/last slice magnitude
+        return loss + weight * (torch.mean(torch.abs(array[0])) + torch.mean(torch.abs(array[-1])))
+
+    def _surface_zero_complex(self, array: torch.Tensor, weight: float | int) -> torch.Tensor:
+        # complex: pull amp toward 1 (vacuum) at the surfaces, and phase toward its mean
+        loss = self._get_zero_loss_tensor()
+        amp = array.abs()
+        loss = loss + weight * (torch.mean(1.0 - amp[0]) + torch.mean(1.0 - amp[-1]))
+        ph = array.angle().abs()
+        warn(
+            "calculating surface zero loss for phase of complex object, "
+            "phase wrapping may distort the gradient. Consider obj_type='pure_phase'."
+        )
+        loss = loss + weight * (
+            torch.mean(torch.abs(ph[0] - ph[0].mean()))
+            + torch.mean(torch.abs(ph[-1] - ph[-1].mean()))
+        )
         return loss
 
     def gaussian_blur_2d(self, tensor, sigma=1.0):
@@ -637,8 +667,9 @@ def butterworth_constraint(
 
         tensor = tensor + tensor_mean
 
-        # Take real part for potential tensorects
-        if self.obj_type == "potential":
+        # FFT-based filter returns complex even for real inputs; cast back to real
+        # for any non-complex object type (pure_phase, potential).
+        if self.obj_type != "complex":
             tensor = tensor.real
 
         return tensor
@@ -774,20 +805,26 @@ def _initialize_obj(
             return
         init_shape = tuple(int(x) for x in shape)
         if self._initialize_mode == "uniform":
-            if self.obj_type in ["complex", "pure_phase"]:
+            if self.obj_type == "complex":
+                # amp=1, phase=0 -> complex ones
                 arr = torch.ones(init_shape) * torch.exp(1.0j * torch.zeros(init_shape))
             else:
+                # pure_phase (phase=0) and potential start as real zeros
                 arr = torch.zeros(init_shape)
         elif self._initialize_mode == "random":
             ph = (
                 torch.randn(init_shape, dtype=torch.float32, generator=self._rng_torch) - 0.5
             ) * 1e-6
-            if self.obj_type == "potential":
-                arr = ph
-            else:
+            if self.obj_type == "complex":
                 arr = torch.exp(1.0j * ph)
+            else:
+                # pure_phase stores phase directly; potential stores real values
+                arr = ph
         elif self._initialize_mode == "array":
             arr = self._initial_obj
+            if self.obj_type == "pure_phase" and arr.is_complex():
+                # Convert legacy complex initial_obj (amp*exp(1j*phase)) to bare phase
+                arr = arr.angle()
         else:
             raise ValueError(f"Invalid initialize mode: {self._initialize_mode}")
 
@@ -938,7 +975,7 @@ def from_pixelated(
             else:
                 model_dtype = "real"
 
-        if pixelated.obj_type == "pure_phase" and model_dtype == "real":
+        if pixelated.obj_type == "complex" and model_dtype == "real":
             obj = pixelated.obj.angle().clone().detach()
         else:
             obj = pixelated.obj.clone().detach()
@@ -957,9 +994,6 @@ def from_pixelated(
 
         return obj_model
 
-    # TODO add a from_params that sets the model input and target from params,
-    # will need to specify a shape as well, at least before pre-training (so just set here)
-
     @property
     def num_slices(self) -> int:
         return self._num_slices
@@ -1099,9 +1133,8 @@ def pretrain_lrs(self) -> np.ndarray:
     def obj(self):
         """get the full object"""
         obj = self.model(self._model_input)[0]
-        if self.obj_type == "pure_phase" and "complex" not in str(self.dtype):
-            # using a real-valued model for a pure-phase (complex) object
-            obj = torch.ones_like(obj) * torch.exp(1j * obj)
+        # pure_phase and potential stay real here; complex models output complex.
+        # _get_obj_patches wraps real -> exp(1j*obj) at the forward boundary.
         # TODO -- single channel 2D with identical slices, view as 3D num_slices
         return self.apply_hard_constraints(obj, mask=self.mask)
 
@@ -1258,8 +1291,6 @@ def _pretrain(
 
             if apply_constraints:
                 output = self.apply_hard_constraints(self.model(model_input)[0])
-                if self.obj_type == "pure_phase":
-                    output = output.angle()
             else:
                 output = self.model(model_input)[0]
             loss: torch.Tensor = loss_fn(output, self.pretrain_target)
diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
index eb73c257e..28155a073 100644
--- a/src/quantem/diffractive_imaging/optimize_hyperparameters.py
+++ b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
@@ -867,7 +867,8 @@ def _plot_grid_objects(self, results, param_names, figsize):
             if recon_obj.obj_type == "potential":
                 obj = np.abs(obj).sum(0)
             elif recon_obj.obj_type == "pure_phase":
-                obj = np.angle(obj).sum(0)
+                # pure_phase obj_cropped is a real phase array — plot directly
+                obj = obj.sum(0)
             else:
                 obj = np.angle(obj).sum(0)
 
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 617196f20..6a0f24daf 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -326,8 +326,14 @@ def verbose(self, v: bool | int | float) -> None:
 
     @property
     def obj(self) -> np.ndarray:
+        """Object array in its native representation per ``obj_type``:
+
+        - ``"complex"`` → complex ndarray (amp * exp(1j*phase)); phase recentered.
+        - ``"pure_phase"`` → real ndarray of phase values.
+        - ``"potential"`` → real ndarray of potential values.
+        """
         obj = self._to_numpy(self.obj_model.obj)
-        if self.obj_type in ["pure_phase", "complex"]:
+        if self.obj_type == "complex":
             ph = np.angle(obj)
             obj = np.abs(obj) * np.exp(1j * (ph - ph.mean()))
         return obj
@@ -497,11 +503,10 @@ def get_snapshot_by_iter(
         if cropped:
             snp2 = snp.copy()
             cropped_obj = self._crop_rotate_obj_fov(snp2["obj"])
-            # same logic as self.obj_cropped
-            if self.obj_type == "pure_phase":
-                ph = np.angle(cropped_obj)
-                cropped_obj = np.exp(1j * (ph - ph.mean()))
-            if self.obj_type in ["pure_phase", "complex"]:
+            # same logic as self.obj_cropped: only re-center for complex (which
+            # carries phase inside a complex tensor); pure_phase and potential
+            # are already real and recentered upstream.
+            if self.obj_type == "complex":
                 ph = np.angle(cropped_obj)
                 cropped_obj = np.abs(cropped_obj) * np.exp(1j * (ph - ph.mean()))
             snp2["obj"] = cropped_obj
@@ -661,8 +666,17 @@ def _dtype_complex(self) -> "torch.dtype":
 
     @property
     def obj_cropped(self) -> np.ndarray:
+        """Cropped + FOV-rotated object, in its native representation.
+
+        - ``obj_type="complex"`` → complex array (amp * exp(1j*phase)); phase is
+          recentered to zero mean here as a defensive duplicate of
+          ``ObjectConstraints._apply_hard_complex``.
+        - ``obj_type="pure_phase"`` → real array of phase values (already
+          recentered upstream by ``_apply_hard_pure_phase``).
+        - ``obj_type="potential"`` → real array of potential values.
+        """
         cropped = self._crop_rotate_obj_fov(self.obj, padding=self.obj_padding_px)
-        if self.obj_type in ["pure_phase", "complex"]:
+        if self.obj_type == "complex":
             ph = np.angle(cropped)
             cropped = np.abs(cropped) * np.exp(1j * (ph - ph.mean()))
         return cropped
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index 3470a5f36..f1e7010de 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -70,12 +70,16 @@ def show_obj(
         ims = []
         titles = []
         cmaps = []
+        # obj_np dtype depends on obj_type:
+        #   potential  -> real values to plot directly
+        #   pure_phase -> real phase to plot directly (no np.angle wrap)
+        #   complex    -> complex; extract amp & phase via np.abs / np.angle
         if self.obj_type == "potential":
             ims.append(np.abs(obj_np).sum(0))
             titles.append(t + "Potential")
             cmaps.append(ph_cmap)
         elif self.obj_type == "pure_phase":
-            ims.append(np.angle(obj_np).sum(0))
+            ims.append(obj_np.sum(0))
             titles.append(t + "Pure Phase")
             cmaps.append(ph_cmap)
         else:
@@ -145,8 +149,14 @@ def show_obj_fft(
             tukey(obj_np.shape[-2], tukey_alpha)[:, None]
             * tukey(obj_np.shape[-1], tukey_alpha)[None, :]
         )
+        # Build the complex transmission function and apply the spatial window:
+        #   potential  -> real values, plotted in real space
+        #   pure_phase -> real phase; transmission = exp(1j*phase)
+        #   complex    -> already complex transmission
         if self.obj_type == "potential":
             windowed_obj = obj_np.sum(0) * window_2d
+        elif self.obj_type == "pure_phase":
+            windowed_obj = np.exp(1j * obj_np.sum(0)) * window_2d
         else:
             windowed_obj = (
                 np.abs(obj_np).sum(0)
@@ -398,7 +408,7 @@ def show_obj_slices(
             objs_flat = [np.abs(obj[i]) for i in range(len(obj))]
             titles_flat = [f"Potential {t_parts[i]}" for i in range(len(obj))]
         elif self.obj_type == "pure_phase":
-            objs_flat = [np.angle(obj[i]) for i in range(len(obj))]
+            objs_flat = [obj[i] for i in range(len(obj))]
             titles_flat = [f"Pure Phase {t_parts[i]}" for i in range(len(obj))]
         else:
             objs_flat = [np.angle(obj[i]) for i in range(len(obj))]
@@ -690,7 +700,7 @@ def _show_object_iters_only(
                 all_titles.append(title_prefix + "Potential")
                 all_cmaps.append(ph_cmap)
             elif self.obj_type == "pure_phase":
-                all_images.append(np.angle(obj).sum(0))
+                all_images.append(obj.sum(0))
                 all_titles.append(title_prefix + "Phase")
                 all_cmaps.append(ph_cmap)
             else:  # complex
@@ -805,7 +815,7 @@ def _show_object_and_probe_iters(
                     row_titles.append(f"Iter {iteration} Potential")
                     row_cmaps.append(ph_cmap)
                 elif self.obj_type == "pure_phase":
-                    row_images.append(np.angle(obj).sum(0))
+                    row_images.append(obj.sum(0))
                     row_titles.append(f"Iter {iteration} Phase")
                     row_cmaps.append(ph_cmap)
                 else:  # complex
diff --git a/tests/diffractive_imaging/test_constraints.py b/tests/diffractive_imaging/test_constraints.py
index 9cce6639b..eeb7e1a0b 100644
--- a/tests/diffractive_imaging/test_constraints.py
+++ b/tests/diffractive_imaging/test_constraints.py
@@ -1,7 +1,10 @@
 """Tests for the ptychography constraint dataclass API."""
 
+import warnings
+
 import numpy as np
 import pytest
+import torch
 
 from quantem.core.datastructures import Dataset4dstem
 from quantem.diffractive_imaging import (
@@ -184,3 +187,103 @@ def test_mixed_dataclass_and_dict_leaves(self, ptycho):
         )
         assert ptycho.obj_model.constraints.tv_weight_xy == 0.4
         assert ptycho.probe_model.constraints.center_probe is True
+
+
+# --- Real-valued pure_phase representation -----------------------------------
+
+
+class TestPurePhaseRealValued:
+    def test_pure_phase_pixelated_obj_is_real(self):
+        obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        assert not obj._obj.is_complex(), f"pure_phase _obj should be real, got {obj._obj.dtype}"
+
+    def test_complex_pixelated_obj_is_complex(self):
+        obj = ObjectPixelated.from_uniform(obj_type="complex", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        assert obj._obj.is_complex()
+
+    def test_potential_pixelated_obj_is_real(self):
+        obj = ObjectPixelated.from_uniform(obj_type="potential", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        assert not obj._obj.is_complex()
+
+    def test_pure_phase_tv_emits_no_phase_warning(self):
+        obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        obj.constraints.tv_weight_xy = 0.1
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            obj.get_tv_loss(obj._obj)
+        phase_warnings = [w for w in caught if "phase wrapping" in str(w.message)]
+        assert not phase_warnings, (
+            f"pure_phase should not emit phase-wrap warning, got {phase_warnings}"
+        )
+
+    def test_complex_tv_still_emits_phase_warning(self):
+        obj = ObjectPixelated.from_uniform(obj_type="complex", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        obj.constraints.tv_weight_xy = 0.1
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            obj.get_tv_loss(obj._obj)
+        assert any("phase wrapping" in str(w.message) for w in caught), (
+            "complex obj_type should still emit phase-wrap warning"
+        )
+
+    def test_pure_phase_apply_hard_constraints_stays_real(self):
+        obj = ObjectPixelated.from_uniform(obj_type="pure_phase", num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        out = obj.apply_hard_constraints(obj._obj)
+        assert not out.is_complex()
+
+
+# --- FOV-mask single application ---------------------------------------------
+
+
+class TestFovMaskSingleApplication:
+    def _make_obj(self, obj_type) -> ObjectPixelated:
+        obj = ObjectPixelated.from_uniform(obj_type=obj_type, num_slices=1)
+        obj._initialize_obj((1, 16, 16), sampling=(0.1, 0.1))
+        obj.constraints.apply_fov_mask = True
+        # Force a non-trivial _obj so masking is observable
+        if obj_type == "complex":
+            obj._obj = torch.nn.Parameter(
+                torch.ones(1, 16, 16, dtype=torch.complex64) * (0.5 + 0.3j)
+            )
+        else:
+            obj._obj = torch.nn.Parameter(torch.full((1, 16, 16), 0.7))
+        return obj
+
+    @pytest.mark.parametrize("obj_type", ["pure_phase", "complex", "potential"])
+    def test_mask_applied_once(self, obj_type):
+        obj = self._make_obj(obj_type)
+        # Half-mask: ones on the left, zeros on the right; if mask is applied
+        # twice the masked region squares the multiplication (no observable
+        # difference for 0/1 masks), so use a non-binary mask.
+        mask = torch.full((1, 16, 16), 0.5)
+        obj._mask = mask
+        out = obj.apply_hard_constraints(obj._obj, mask=mask)
+        # Verify nothing crashed and shape is preserved.
+        assert out.shape == obj._obj.shape
+        # If mask had been applied twice, |out| would scale by 0.5**2 = 0.25
+        # of the unmasked value; once it scales by 0.5. We compare to the
+        # per-obj-type expected post-constraint value.
+        if obj_type == "pure_phase":
+            # phase recentered to zero mean, then *= 0.5 mask
+            expected_mag = 0.0  # phase=constant -> recenter to 0 -> *0.5 = 0
+        elif obj_type == "potential":
+            # positivity clamp keeps 0.7, * 0.5 -> 0.35 (one application)
+            expected_mag = 0.35
+        else:  # complex
+            # amp clamp keeps 0.5+0.3j, * 0.5 -> magnitude 0.5 * |0.5+0.3j|
+            expected_mag = 0.5 * abs(0.5 + 0.3j)
+        # Sample the magnitude in the masked region
+        if out.is_complex():
+            sampled = out.abs().mean().item()
+        else:
+            sampled = out.abs().mean().item()
+        assert abs(sampled - expected_mag) < 1e-4, (
+            f"{obj_type}: expected mag ~{expected_mag}, got {sampled} "
+            f"(would be {expected_mag * 0.5} if mask were applied twice)"
+        )

From 85dfa3ff01d0e198e55cea68b8b25d197747eaee Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 28 May 2026 14:23:34 -0700
Subject: [PATCH 23/59] bugfix of initializing to first devices even if not
 specified

---
 src/quantem/core/ml/dist_utils.py              | 18 +++++++++++++++---
 .../diffractive_imaging/ptychography.py        | 17 +++++++++++++++--
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
index 856c84d25..450c4fead 100644
--- a/src/quantem/core/ml/dist_utils.py
+++ b/src/quantem/core/ml/dist_utils.py
@@ -25,17 +25,29 @@ def init_process_group(
     backend: str = "nccl",
     master_addr: str = "127.0.0.1",
     master_port: str = "29500",
+    local_device: int | None = None,
 ) -> None:
-    """Initialize the distributed process group from within an mp.spawn worker."""
+    """Initialize the distributed process group from within an mp.spawn worker.
+
+    ``local_device`` is the physical CUDA device index this rank should bind to
+    (e.g. with ``GPU_IDS=[2, 3]``, rank 0 should get ``local_device=2``).
+    NCCL allocates communicator buffers on the *current* CUDA device at
+    ``init_process_group`` time, so the device must be set *before* that call
+    or the buffers will land on whichever device was current — typically
+    ``cuda:0``. Falling back to ``rank`` matches PyTorch's
+    ``LOCAL_RANK == device_index`` convention used by ``torchrun`` when each
+    process maps to a contiguous device starting at 0.
+    """
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
+    if backend == "nccl":
+        device_index = local_device if local_device is not None else rank
+        torch.cuda.set_device(device_index)
     dist.init_process_group(
         backend=backend,
         rank=rank,
         world_size=world_size,
     )
-    if backend == "nccl":
-        torch.cuda.set_device(rank)
 
 
 def get_rank() -> int:
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 05e0b5af3..4d9345e73 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -44,7 +44,16 @@ def _ddp_ptycho_worker(
     shared-memory tensor mechanism and fails in some Linux environments).
     """
     device_id = devices[rank]
-    init_process_group(rank, world_size, backend="nccl" if torch.cuda.is_available() else "gloo")
+    # Bind the CUDA device BEFORE init_process_group so NCCL allocates its
+    # communicator buffers on the correct GPU. Without this, NCCL grabs cuda:0
+    # at init time, stranding small per-rank buffers on GPUs the user didn't
+    # ask for.
+    init_process_group(
+        rank,
+        world_size,
+        backend="nccl" if torch.cuda.is_available() else "gloo",
+        local_device=device_id if torch.cuda.is_available() else None,
+    )
 
     # mmap=True so all workers share one memory-mapped RAM copy of the (potentially large,
     # CPU-resident) state instead of each duplicating it.
@@ -263,12 +272,16 @@ def reconstruct(
         if is_distributed_launch():
             rank = int(os.environ["RANK"])
             world_size = int(os.environ["WORLD_SIZE"])
+            local_rank = int(os.environ.get("LOCAL_RANK", rank))
             if not torch.distributed.is_initialized():
+                # Bind the device BEFORE init_process_group so NCCL allocates
+                # its communicator buffers on the correct GPU.
+                if torch.cuda.is_available():
+                    torch.cuda.set_device(local_rank)
                 torch.distributed.init_process_group(
                     backend="nccl" if torch.cuda.is_available() else "gloo",
                     init_method="env://",
                 )
-            local_rank = int(os.environ.get("LOCAL_RANK", rank))
             dev = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
             self.to(dev)
             self._broadcast_parameters(src=0)

From 74f52ea6a6434df12f4e3404cdd4baca78811d2a Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 28 May 2026 14:46:38 -0700
Subject: [PATCH 24/59] adding TODO for amp/phase tv weight splitting

---
 src/quantem/diffractive_imaging/object_models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 7d85d8b95..f66156a3d 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -540,6 +540,9 @@ def _tv_complex(self, array: torch.Tensor, w: tuple[float, float]) -> torch.Tens
             "calculating TV loss for phase of complex object, "
             "phase wrapping may distort the gradient. Consider obj_type='pure_phase'."
         )
+        # TODO: amp and phase share `w` here. Consider splitting `tv_weight_xy`
+        # into separate amp/phase weights on PtychoObjConstraintParams.Raster
+        # so users can tune them independently for obj_type="complex".
         loss = loss + self._calc_tv_loss(ph, w)
         amp = array.abs()
         loss = loss + self._calc_tv_loss(amp, w)

From b7d084bbb14006fcc0931838d2466978c5180585 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 28 May 2026 19:52:56 -0700
Subject: [PATCH 25/59] moving hard constraints outside of computational graph

---
 .../diffractive_imaging/object_models.py      | 118 ++++++++++++------
 1 file changed, 78 insertions(+), 40 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index f66156a3d..1c3d60ba1 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -388,6 +388,15 @@ def _propagate_array(
         return propagated
 
     def _get_obj_patches(self, obj_array, patch_indices):
+        """Forward boundary: wrap real obj to ``exp(1j * obj)`` and gather patches.
+
+        ``obj_array`` may be complex (``obj_type="complex"``) or real (``"pure_phase"``,
+        ``"potential"``). Real inputs are wrapped to the complex transmission
+        function ``exp(1j * obj_array)`` here, so the rest of the forward model
+        never has to branch on ``obj_type``. ``patch_indices`` is a
+        ``(num_gpts, Hroi, Wroi)`` int tensor of flattened-index lookups into the
+        2D padded object.
+        """
         if not obj_array.is_complex():  # potential or pure_phase DIP -> float
             obj_array2 = torch.exp(1.0j * obj_array)
         else:
@@ -412,16 +421,22 @@ class ObjectConstraints(BaseConstraints[PtychoObjConstraintParams.Raster], Objec
     DEFAULT_CONSTRAINTS: PtychoObjConstraintParams.Raster = PtychoObjConstraintParams.Raster()
 
     def apply_hard_constraints(
-        self, obj: torch.Tensor, mask: torch.Tensor | None = None
+        self, raw: torch.Tensor, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
+        """
+        Apply hard constraints: range clamping and filtering. All hard constaints are applied in
+        place with torch.no_grad().
+        """
         c = self.constraints
-        if self.obj_type == "complex":
-            obj2 = self._apply_hard_complex(obj, c)
-        elif self.obj_type == "pure_phase":
-            obj2 = self._apply_hard_pure_phase(obj, c)
-        else:  # potential
-            obj2 = self._apply_hard_potential(obj, c, mask)
-        return self._apply_shared_hard(obj2, c, mask)
+        with torch.no_grad():
+            if self.obj_type == "complex":
+                constrained = self._apply_hard_complex(raw, c)
+            elif self.obj_type == "pure_phase":
+                constrained = self._apply_hard_pure_phase(raw, c)
+            else:  # potential
+                constrained = self._apply_hard_potential(raw, c, mask)
+            constrained = self._apply_shared_hard(constrained, c, mask)
+        return raw + (constrained - raw).detach()
 
     def _apply_hard_complex(
         self, obj: torch.Tensor, c: PtychoObjConstraintParams.Raster
@@ -476,19 +491,24 @@ def _apply_shared_hard(
             obj = self.butterworth_constraint(obj, sampling=self.sampling)
 
         if self.num_slices > 1 and c.identical_slices:
-            with torch.no_grad():
-                obj[:] = torch.mean(obj, dim=0, keepdim=True)
+            # In-place mutation is safe because apply_hard_constraints is
+            # always called under outer torch.no_grad (see its docstring).
+            obj[:] = torch.mean(obj, dim=0, keepdim=True)
         return obj
 
     def apply_soft_constraints(
         self, obj: torch.Tensor, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
+        """Sum of the per-iteration soft penalties.
+
+        Returns a scalar tensor that is added to the data-fidelity loss before
+        ``backward()``. Individual contributions are also recorded via
+        ``add_soft_constraint_loss`` for logging.
+        """
         # reset recorded losses each call
         self.reset_soft_constraint_losses()
 
-        tv_loss = self.get_tv_loss(
-            obj,
-        )
+        tv_loss = self.get_tv_loss(obj)
         self.add_soft_constraint_loss("tv_loss", tv_loss)
 
         surface_zero_loss = self.get_surface_zero_loss(
@@ -502,6 +522,13 @@ def apply_soft_constraints(
     def get_tv_loss(
         self, array: torch.Tensor, weights: None | tuple[float, float] = None
     ) -> torch.Tensor:
+        """Total-variation soft penalty on the object.
+
+        ``weights`` is a ``(z_weight, xy_weight)`` tuple. When ``None``, defaults
+        to ``(self.constraints.tv_weight_z, self.constraints.tv_weight_xy)``. A single
+        scalar is broadcast to both axes. The z weight is zeroed for
+        ``num_slices == 1``.
+        """
         loss = self._get_zero_loss_tensor()
         w = self._resolve_tv_weights(weights)
         if not any(w):
@@ -549,6 +576,12 @@ def _tv_complex(self, array: torch.Tensor, w: tuple[float, float]) -> torch.Tens
         return loss
 
     def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> torch.Tensor:
+        """Mean-|diff| TV on a real array. ``weight = (w_z, w_xy)``.
+
+        For a 3D ``(slices, H, W)`` array, dim 0 uses ``w_z`` and dims 1+2 use
+        ``w_xy``. The result is averaged over the number of axes that actually
+        contributed (i.e. had a non-zero weight).
+        """
         loss = self._get_zero_loss_tensor()
         calc_dim = 0
         for dim in range(array.ndim):
@@ -566,6 +599,13 @@ def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> tor
     def get_surface_zero_loss(
         self, array: torch.Tensor, weight: float | int = 0.0
     ) -> torch.Tensor:
+        """Penalize the first and last slices to be near vacuum.
+
+        Real ``pure_phase`` / ``potential`` arrays: penalizes ``|array[0]|`` and
+        ``|array[-1]|`` directly. ``complex`` arrays pull amplitude toward 1
+        and phase toward its mean (see ``_surface_zero_complex``). A no-op for
+        single- or double-slice objects (``array.shape[0] < 3``).
+        """
         loss = self._get_zero_loss_tensor()
         if weight == 0 or array.shape[0] < 3:
             return loss
@@ -591,12 +631,16 @@ def _surface_zero_complex(self, array: torch.Tensor, weight: float | int) -> tor
         return loss
 
     def gaussian_blur_2d(self, tensor, sigma=1.0):
-        """
-        Apply Gaussian blur along dimensions 2 and 3 of a 3D tensor.
+        """Separable 2D Gaussian blur over the last two dimensions.
 
-        Args:
-            tensor: Can be real or complex
-            sigma: Standard deviation for Gaussian kernel
+        Parameters
+        ----------
+        tensor : torch.Tensor
+            Real or complex, shape ``(slices, H, W)``. Complex inputs are
+            filtered as independent real/imag channels.
+        sigma : float
+            Standard deviation of the Gaussian kernel, in pixels. The
+            kernel size is ``2 * ceil(3 * sigma) + 1``.
         """
         kernel_size = int(2 * math.ceil(3 * sigma) + 1)
         if kernel_size % 2 == 0:
@@ -639,9 +683,21 @@ def butterworth_constraint(
         tensor: torch.Tensor,
         sampling: tuple[float, float],
     ) -> torch.Tensor:
-        """
-        Butterworth filter used for low/high-pass filtering.
+        """Apply a Fourier-domain Butterworth low/high-pass to each 2D slice.
+
+        Reads ``q_lowpass``, ``q_highpass``, and ``butterworth_order`` off
+        ``self.constraints``. The DC component is subtracted before filtering
+        and added back so the mean is preserved.
 
+        Parameters
+        ----------
+        tensor : torch.Tensor
+            Shape ``(slices, H, W)``. May be real or complex. Real inputs are
+            re-cast to real after the FFT round-trip when ``obj_type != "complex"``.
+        sampling : tuple[float, float]
+            ``(dy, dx)`` real-space sampling in Ångström per pixel. Sets the
+            inverse-Å scale of the Butterworth response; ``q_lowpass`` and
+            ``q_highpass`` are in *inverse Ångström (cycles / Å).
         """
 
         q_lowpass = self.constraints.q_lowpass
@@ -1072,21 +1128,6 @@ def model_input(self, input_tensor: torch.Tensor | np.ndarray):
 
         self._model_input = input_tensor.type(self.dtype).to(self.device)
 
-    # def _generate_model_input(self, mode: Literal["random", "zeros", "ones"]) -> None:
-    #     input_shape = (1, *self.shape)
-    #     # could support for 3D CNN models, single channel 2D with identical slices
-    #     if mode == "random":
-    #         inp = torch.randn(
-    #             input_shape, device=self.device, dtype=self.dtype, generator=self._rng_torch
-    #         )
-    #     elif mode == "zeros":
-    #         inp = torch.zeros(input_shape, device=self.device, dtype=self.dtype)
-    #     elif mode == "ones":
-    #         inp = torch.ones(input_shape, device=self.device, dtype=self.dtype)
-    #     else:
-    #         raise ValueError(f"Invalid mode: {mode} | must be one of: 'random', 'zeros', 'ones'")
-    #     self._model_input = inp
-
     @property
     def pretrain_target(self) -> torch.Tensor:
         """get the pretrain target"""
@@ -1135,15 +1176,12 @@ def pretrain_lrs(self) -> np.ndarray:
     @property
     def obj(self):
         """get the full object"""
-        obj = self.model(self._model_input)[0]
-        # pure_phase and potential stay real here; complex models output complex.
-        # _get_obj_patches wraps real -> exp(1j*obj) at the forward boundary.
+        raw = self.model(self._model_input)[0]
         # TODO -- single channel 2D with identical slices, view as 3D num_slices
-        return self.apply_hard_constraints(obj, mask=self.mask)
+        return self.apply_hard_constraints(raw, mask=self.mask)
 
     @property
     def _obj(self):
-        # TODO -- single channel 2D with identical slices, view as 3D num_slices??
         return self.model(self._model_input)[0]
 
     def forward(self, patch_indices: torch.Tensor):

From 72efce95493d737dfa42f23d54fa855faa1afa8f Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 29 May 2026 13:45:33 -0700
Subject: [PATCH 26/59] workingish ptycho inr

---
 src/quantem/diffractive_imaging/__init__.py   |   1 +
 .../diffractive_imaging/dataset_models.py     |  76 ++-
 .../diffractive_imaging/object_models.py      | 617 +++++++++++++++++-
 .../diffractive_imaging/ptychography.py       |  21 +-
 .../diffractive_imaging/ptychography_base.py  |   4 +
 5 files changed, 679 insertions(+), 40 deletions(-)

diff --git a/src/quantem/diffractive_imaging/__init__.py b/src/quantem/diffractive_imaging/__init__.py
index 9db66f112..d9058a3fa 100644
--- a/src/quantem/diffractive_imaging/__init__.py
+++ b/src/quantem/diffractive_imaging/__init__.py
@@ -6,6 +6,7 @@
 from quantem.diffractive_imaging.detector_models import DetectorPixelated as DetectorPixelated
 from quantem.diffractive_imaging.object_models import (
     ObjectDIP as ObjectDIP,
+    ObjectINR as ObjectINR,
     ObjectPixelated as ObjectPixelated,
     PtychoObjConstraintParams as PtychoObjConstraintParams,
     PtychoObjConstraintsType as PtychoObjConstraintsType,
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 7b538e488..473eb8db5 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -174,6 +174,10 @@ def __init__(
         self.detector_mask = detector_mask
         self._constraints = {}
         self._probe_energy = None
+        # Set by the ptychography wiring from obj_model.is_implicit. When True, forward()
+        # emits continuous per-patch coordinates instead of integer patch_indices (and zeroed
+        # fractional positions), so the probe is not subpixel-shifted.
+        self._implicit_object = False
 
     def get_optimization_parameters(self):
         """Get the combined descan and scan position parameters for optimization.
@@ -512,6 +516,20 @@ def device(self) -> torch.device:
     def preprocessed(self) -> bool:
         return self._preprocessed
 
+    @property
+    def implicit_object(self) -> bool:
+        """Whether the paired object model is an implicit (coordinate-queried) representation.
+
+        Set by the ptychography wiring from ``obj_model.is_implicit``. When True, ``forward``
+        emits continuous per-patch coordinates (instead of integer ``patch_indices``) along with
+        zeroed fractional positions, so the probe is not subpixel-shifted.
+        """
+        return self._implicit_object
+
+    @implicit_object.setter
+    def implicit_object(self, val: bool) -> None:
+        self._implicit_object = bool(val)
+
     @property
     def shape(self) -> np.ndarray:
         return np.array(self.dset.shape)
@@ -665,6 +683,38 @@ def patch_indices_need_update(self) -> bool:
         new_pos = torch.round(self.scan_positions_px)
         return not torch.equal(old_pos, new_pos)
 
+    def _scan_coords(
+        self, batch_indices: np.ndarray | torch.Tensor, obj_padding_px: np.ndarray | tuple
+    ) -> torch.Tensor:
+        """Continuous normalized ``(row, col)`` patch coordinates for implicit objects.
+
+        For each scan position (NOT rounded) we add the same integer ROI offsets used by
+        ``_set_patch_indices``, then normalize to ``[-1, 1]`` over the padded object extent
+        via ``idx / (N - 1) * 2 - 1`` (matching ``torch.linspace(-1, 1, N)``). Unlike
+        ``_set_patch_indices`` there is no modulo wrap: samples that fall outside the object
+        map outside ``[-1, 1]``, and the implicit object treats them as vacuum.
+
+        Because the un-rounded (fractional) position is baked into the coordinates here, the
+        probe is queried at zero fractional shift in this case.
+
+        Returns
+        -------
+        torch.Tensor
+            ``(batch, Hroi, Wroi, 2)`` normalized ``(row, col)`` coordinates.
+        """
+        obj_shape = self._obj_shape_full_2d(obj_padding_px)
+        positions = self.scan_positions_px[batch_indices]  # (batch, 2), un-rounded
+        hroi, wroi = int(self.roi_shape[0]), int(self.roi_shape[1])
+        x_ind = torch.fft.fftfreq(hroi, d=1 / hroi).to(self.device)
+        y_ind = torch.fft.fftfreq(wroi, d=1 / wroi).to(self.device)
+        rows = positions[:, 0][:, None, None] + x_ind[None, :, None]  # (batch, Hroi, 1)
+        cols = positions[:, 1][:, None, None] + y_ind[None, None, :]  # (batch, 1, Wroi)
+        rows = rows.expand(-1, -1, wroi)
+        cols = cols.expand(-1, hroi, -1)
+        rows_n = rows / float(obj_shape[-2] - 1) * 2.0 - 1.0
+        cols_n = cols / float(obj_shape[-1] - 1) * 2.0 - 1.0
+        return torch.stack([rows_n, cols_n], dim=-1)  # (batch, Hroi, Wroi, 2)
+
     def reset(self) -> None:
         self.descan_shifts = self.initial_descan_shifts.clone().to(self.device)
         self.scan_positions_px = self.initial_scan_positions_px.clone().to(self.device)
@@ -1592,19 +1642,31 @@ def forward(
         batch_indices: np.ndarray | torch.Tensor,
         obj_padding_px: np.ndarray | tuple,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        """Forward pass to compute the diffraction intensities from the object and scan positions."""
+        """Forward pass to compute the diffraction intensities from the object and scan positions.
+
+        The first return value is the object-query payload: integer ``patch_indices`` for a
+        grid-based object, or continuous normalized coordinates for an implicit object
+        (``implicit_object=True``). In the implicit case the fractional position is baked into
+        the coordinates, so the returned fractional shift is zero (the probe is not shifted).
+        """
         self.apply_hard_constraints(obj_padding_px)
         positions_px = self.scan_positions_px[batch_indices]
-        positions_px_fractional = positions_px - torch.round(positions_px)
-        with torch.no_grad():
-            if self.patch_indices_need_update():
-                self._set_patch_indices(obj_padding_px)
-        patch_indices = self.patch_indices[batch_indices]
         if self.learn_descan and self.has_optimizer():
             descan_shifts = self.apply_descan_constraints(self.descan_shifts)[batch_indices]
         else:
             descan_shifts = None
-        return patch_indices, positions_px, positions_px_fractional, descan_shifts
+
+        if self._implicit_object:
+            patch_data = self._scan_coords(batch_indices, obj_padding_px)
+            positions_px_fractional = torch.zeros_like(positions_px)
+            return patch_data, positions_px, positions_px_fractional, descan_shifts
+
+        positions_px_fractional = positions_px - torch.round(positions_px)
+        with torch.no_grad():
+            if self.patch_indices_need_update():
+                self._set_patch_indices(obj_padding_px)
+        patch_data = self.patch_indices[batch_indices]
+        return patch_data, positions_px, positions_px_fractional, descan_shifts
 
 
 DatasetModelType = PtychographyDatasetRaster  # | PtychographyDatasetSpiral
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 1c3d60ba1..78d713bb0 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -15,6 +15,7 @@
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.blocks import reset_weights
 from quantem.core.ml.constraints import BaseConstraints, Constraints, parse_constraint_dict
+from quantem.core.ml.inr import HSiren
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerType, SchedulerType
 from quantem.core.utils.rng import RNGMixin
@@ -137,17 +138,30 @@ class Raster(Constraints):
 
     @dataclass
     class INR(Constraints):
-        """Placeholder for the upcoming ``ObjectINR`` variant.
+        """Constraints for the implicit (``ObjectINR``) object representation.
 
-        INR-specific constraints (e.g. sparsity / TV penalties evaluated at
-        sampled coordinates) will land here when the model is implemented.
-        Until then this exists so ``parse_dict`` accepts ``"inr"`` and downstream
-        code can pattern-match on the variant.
+        An INR has no grid to project, so the grid-based hard constraints of
+        ``Raster`` (positivity, filtering, FOV masking) do not apply. The
+        regularizers that do carry over are soft penalties evaluated at sampled
+        coordinates.
+
+        Attributes
+        ----------
+        tv_weight_z : float, default ``0.0``
+            Soft penalty. Weight on the depth-axis (``z``) total-variation term,
+            evaluated via finite differences at randomly sampled coordinates.
+            Multislice (``num_slices > 1``) only.
+        tv_weight_xy : float, default ``0.0``
+            Soft penalty. Weight on the in-plane (``y``, ``x``) total-variation
+            term, evaluated via finite differences at randomly sampled coordinates.
         """
 
+        # soft constraints (evaluated at sampled coordinates)
+        tv_weight_z: float = 0.0
+        tv_weight_xy: float = 0.0
         _name: str = "inr"
 
-        soft_constraint_keys = []
+        soft_constraint_keys = ["tv_weight_z", "tv_weight_xy"]
         hard_constraint_keys = []
 
     @classmethod
@@ -211,6 +225,17 @@ def __init__(
     def shape(self) -> tuple[int, int, int]:
         return self.obj.shape
 
+    @property
+    def is_implicit(self) -> bool:
+        """Whether this is an implicit (coordinate-queried) object representation.
+
+        Pixelated/DIP objects are grid-based and consume integer ``patch_indices``;
+        an implicit object (``ObjectINR``) instead consumes continuous coordinates,
+        which the paired dataset produces when this is True. Overridden to True by
+        implicit subclasses.
+        """
+        return False
+
     @property
     @abstractmethod
     def num_slices(self) -> int:
@@ -1472,28 +1497,570 @@ def visualize_pretrain(
         plt.show()
 
 
-# class ObjectImplicit(ObjectBase):
-#     """
-#     Object model for implicit objects. Importantly, the forward call from scan positions
-#     for this model will not require subpixel shifting of the object probe, as subpixel shifting
-#     will be done in the object model itself, so it is properly aligned around the probe positions
-#     """
+class ObjectINR(ObjectConstraints):
+    """Implicit (coordinate-queried) object model.
+
+    Wraps an implicit neural representation (INR; an ``HSiren`` by default) that maps
+    normalized 3D coordinates ``(z, y, x)`` in ``[-1, 1]`` to the object's real phase
+    (``obj_type="pure_phase"``). Rather than gathering grid-aligned patches at integer
+    scan positions like ``ObjectPixelated``, the paired dataset produces continuous
+    per-patch ``(y, x)`` coordinates at the *true* (fractional) scan positions; this
+    model augments them with each slice's ``z`` coordinate, queries the INR, and returns
+    the complex transmission patches ``exp(1j * phase)``. Because the object is sampled
+    directly at the true position, the probe no longer needs subpixel shifting.
+
+    Coordinate convention
+    ----------------------
+    Coordinates are normalized to ``[-1, 1]`` over the padded object extent (matching
+    ``torch.linspace(-1, 1, N)``). The ``z`` axis spans the slices: a single slice sits
+    at ``z = 0``; multislice z-positions come from the cumulative ``slice_thicknesses``,
+    mapped so the first slice is at ``-1`` and the last at ``+1``. Samples that fall
+    outside the object (``|y| > 1`` or ``|x| > 1``) are treated as vacuum (phase 0,
+    transmission 1) rather than wrapped toroidally as the pixelated path does.
+
+    Notes
+    -----
+    - Autograd-only: analytical gradients are not implemented (see ``backward``).
+    - Hard constraints have no grid to project; only the soft, coordinate-sampled TV
+      penalties of ``PtychoObjConstraintParams.INR`` apply.
+    """
 
-#     def __init__(self, *args, **kwargs):
-#         super().__init__(*args, **kwargs)
-#         self._obj = None
-#         self._obj_shape = None
-#         self._num_slices = None
+    DEFAULT_LRS = {
+        "object": 8e-6,
+        "tv_weight_z": 0,
+        "tv_weight_xy": 0,
+    }
+    DEFAULT_CONSTRAINTS: PtychoObjConstraintParams.INR = PtychoObjConstraintParams.INR()
+
+    def __init__(
+        self,
+        model: "torch.nn.Module",
+        num_slices: int = 1,
+        slice_thicknesses: float | Sequence | torch.Tensor | None = None,
+        obj_type: object_type = "pure_phase",
+        device: str = "cpu",
+        rng: np.random.Generator | int | None = None,
+        _token: object | None = None,
+    ):
+        super().__init__(
+            device=device,
+            obj_type=obj_type,
+            rng=rng,
+            _token=_token,
+        )
+        if self.obj_type != "pure_phase":
+            raise NotImplementedError(
+                f"ObjectINR currently only supports obj_type='pure_phase', got '{self.obj_type}'. "
+                "Complex/potential INR objects are planned."
+            )
+        if num_slices < 1:
+            raise ValueError(f"num_slices must be greater than 0, got {num_slices}")
+        self._num_slices = int(num_slices)
+        self._model = model.to(self._device)
+        self.slice_thicknesses = slice_thicknesses
+        self._set_pretrained_weights(self._model)
+
+        # Padded object extent [num_slices, H, W]; set in _initialize_obj. Defines the
+        # [-1, 1] coordinate domain and the grid on which .obj is materialized.
+        self._obj_shape: tuple[int, int, int] | None = None
+        # Lazily materialized full-grid object (detached); invalidated each forward().
+        self._obj_cache: torch.Tensor | None = None
+        # Pretraining state (used by pretrain() / from_pixelated()).
+        self.register_buffer("_pretrain_target", torch.tensor([]))
+        self._pretrain_losses: list[float] = []
+        self._pretrain_lrs: list[float] = []
+
+    @classmethod
+    def from_inr(
+        cls,
+        model: "torch.nn.Module",
+        num_slices: int = 1,
+        slice_thicknesses: float | Sequence | torch.Tensor | None = None,
+        obj_type: object_type = "pure_phase",
+        device: str = "cpu",
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectINR":
+        """Create an ObjectINR from a user-supplied INR ``nn.Module``.
+
+        The model must map coordinates of shape ``(N, 3)`` (``z, y, x``) to a single
+        real output ``(N, 1)``.
+        """
+        return cls(
+            model=model,
+            num_slices=num_slices,
+            slice_thicknesses=slice_thicknesses,
+            obj_type=obj_type,
+            device=device,
+            rng=rng,
+            _token=cls._token,
+        )
+
+    @classmethod
+    def from_uniform(
+        cls,
+        num_slices: int = 1,
+        slice_thicknesses: float | Sequence | torch.Tensor | None = None,
+        hidden_features: int = 128,
+        hidden_layers: int = 3,
+        first_omega_0: float = 10.0,
+        hidden_omega_0: float = 10.0,
+        obj_type: object_type = "pure_phase",
+        device: str = "cpu",
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectINR":
+        """Create an ObjectINR backed by a default ``HSiren``, initialized to vacuum.
+
+        The HSiren's final layer is zero-initialized so the object starts as a flat,
+        phase-0 (vacuum) transmission, matching ``ObjectPixelated.from_uniform``.
+
+        Note
+        ----
+        ``first_omega_0`` / ``hidden_omega_0`` set the SIREN's frequency content and are the
+        most important knobs to tune (INRs are sensitive to this, more so than DIPs). The
+        default of ``10`` suits smooth phase objects; the SIREN image-fitting default of ``30``
+        is typically too high here (optimization stalls near vacuum), while objects with fine
+        features may want a larger value. Pair omega_0 with the object learning rate.
+        """
+        model = HSiren(
+            in_features=3,
+            out_features=1,
+            hidden_layers=hidden_layers,
+            hidden_features=hidden_features,
+            first_omega_0=first_omega_0,
+            hidden_omega_0=hidden_omega_0,
+            dtype=getattr(torch, config.get("dtype_real")),
+        )
+        # Zero the final linear layer so the INR outputs phase 0 everywhere (vacuum start).
+        with torch.no_grad():
+            final_linear = model.net[-2]
+            final_linear.weight.zero_()
+            if final_linear.bias is not None:
+                final_linear.bias.zero_()
+        return cls.from_inr(
+            model=model,
+            num_slices=num_slices,
+            slice_thicknesses=slice_thicknesses,
+            obj_type=obj_type,
+            device=device,
+            rng=rng,
+        )
+
+    @classmethod
+    def from_pixelated(
+        cls,
+        pixelated: "ObjectModelType",
+        hidden_features: int = 256,
+        hidden_layers: int = 3,
+        first_omega_0: float = 10.0,
+        hidden_omega_0: float = 10.0,
+        device: str | None = None,
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectINR":
+        """Create an ObjectINR matching a pixelated object, with it as the pretrain target.
+
+        The INR is built to the pixelated object's geometry (``num_slices``,
+        ``slice_thicknesses``, ``obj_type``, padded shape) and the current pixelated object is
+        stored as the pretrain target, so ``pretrain()`` warm-starts the INR to reproduce the
+        pixelated reconstruction -- mirroring ``ObjectDIP.from_pixelated`` + ``pretrain``.
+        """
+        if not (
+            isinstance(pixelated, ObjectPixelated) or "ObjectPixelated" in str(type(pixelated))
+        ):
+            raise ValueError(f"pixelated must be an ObjectPixelated, got {type(pixelated)}")
+        dev = pixelated.device if device is None else device
+        inr = cls.from_uniform(
+            num_slices=pixelated.num_slices,
+            slice_thicknesses=pixelated.slice_thicknesses,
+            obj_type=pixelated.obj_type,
+            hidden_features=hidden_features,
+            hidden_layers=hidden_layers,
+            first_omega_0=first_omega_0,
+            hidden_omega_0=hidden_omega_0,
+            device=dev,
+            rng=pixelated._rng_seed if rng is None else rng,
+        )
+        target = pixelated.obj.detach().to(dev)  # (num_slices, H, W) real phase
+        inr._obj_shape = tuple(int(x) for x in target.shape)  # type: ignore[assignment]
+        if pixelated._sampling is not None:
+            inr.sampling = pixelated.sampling
+        inr.pretrain_target = target
+        return inr
+
+    # region --- properties ---
+    @property
+    def is_implicit(self) -> bool:
+        return True
+
+    @property
+    def name(self) -> str:
+        return "ObjectINR"
+
+    @property
+    def num_slices(self) -> int:
+        return self._num_slices
+
+    @property
+    def model(self) -> "torch.nn.Module":
+        return self._model
+
+    @property
+    def params(self):
+        """optimization parameters"""
+        return self._model.parameters()
+
+    @property
+    def pretrained_weights(self) -> dict[str, torch.Tensor]:
+        return self._pretrained_weights
+
+    def _set_pretrained_weights(self, model: "torch.nn.Module") -> None:
+        self._pretrained_weights = deepcopy(model.state_dict())
+
+    @property
+    def pretrain_target(self) -> torch.Tensor:
+        """Target object (real phase, ``(num_slices, H, W)``) fitted by ``pretrain()``."""
+        return self._pretrain_target
+
+    @pretrain_target.setter
+    def pretrain_target(self, target: torch.Tensor | np.ndarray | None) -> None:
+        if target is None:
+            self._pretrain_target = torch.tensor([], device=self.device)
+            return
+        t = validate_tensor(
+            target,
+            name="pretrain_target",
+            ndim=3,
+            dtype=config.get("dtype_real"),
+            expand_dims=True,
+        )
+        self._pretrain_target = t.to(self.device)
+
+    @property
+    def pretrain_losses(self) -> np.ndarray:
+        return np.array(self._pretrain_losses)
+
+    @property
+    def pretrain_lrs(self) -> np.ndarray:
+        return np.array(self._pretrain_lrs)
+
+    @property
+    def _z_coords(self) -> torch.Tensor:
+        """Normalized z-coordinate of each slice in [-1, 1], shape (num_slices,)."""
+        real_dtype = getattr(torch, config.get("dtype_real"))
+        s = self.num_slices
+        if s == 1:
+            return torch.zeros(1, device=self.device, dtype=real_dtype)
+        thick = self._slice_thicknesses.to(self.device)  # (S-1,)
+        zeros = torch.zeros(1, device=self.device, dtype=real_dtype)
+        z_pos = torch.cat([zeros, torch.cumsum(thick, dim=0)])  # (S,)
+        total = z_pos[-1]
+        if total <= 0:
+            return torch.linspace(-1.0, 1.0, s, device=self.device, dtype=real_dtype)
+        return (z_pos / total) * 2.0 - 1.0
+
+    @property
+    def obj(self):
+        """Materialized full object on the padded grid (real phase for pure_phase).
+
+        Cold-path only (display / logging / serialization); the training loop queries
+        the INR directly via ``forward`` and does not touch this. Cached and invalidated
+        on each ``forward`` call.
+        """
+        if self._obj_cache is None:
+            raw = self._materialize_obj()
+            self._obj_cache = self.apply_hard_constraints(raw, mask=self.mask)
+        return self._obj_cache
 
-#     def pretrain(self, *args, **kwargs):
+    # endregion --- properties ---
 
+    def _invalidate_obj_cache(self) -> None:
+        self._obj_cache = None
 
-#     ### here the forward call will take the batch indices and create the appropriate
-#     ### input (which maybe is just the raw patch indices? tbd) for the implicit input
-#     ### so it will be parallelized inference across the batches rather than inference once
-#     ### and then patching that, like it will be for DIP
+    def _query_phase(self, coords_xy: torch.Tensor) -> torch.Tensor:
+        """Query the INR at normalized (y, x) coords for every slice.
+
+        ``coords_xy`` has shape ``(..., 2)`` (normalized row=y, col=x in [-1, 1]).
+        Returns phase of shape ``(num_slices, ...)``.
+        """
+        s = self.num_slices
+        lead = coords_xy.shape[:-1]
+        xy = coords_xy.reshape(-1, 2)  # (M, 2)
+        m = xy.shape[0]
+        z = self._z_coords  # (S,)
+        z_full = z.view(s, 1, 1).expand(s, m, 1)
+        xy_full = xy.view(1, m, 2).expand(s, m, 2)
+        coords3d = torch.cat([z_full, xy_full], dim=-1).reshape(-1, 3)  # (S*M, 3)
+        phase = self._model(coords3d).squeeze(-1).reshape(s, *lead)
+        return phase
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        """Query the object at continuous per-patch coordinates.
+
+        ``coords`` has shape ``(batch, Hroi, Wroi, 2)``: normalized ``(y, x)`` positions
+        in ``[-1, 1]`` produced by the dataset at the true (fractional) scan positions.
+        Returns complex transmission patches of shape ``(num_slices, batch, Hroi, Wroi)``.
+        """
+        self._invalidate_obj_cache()
+        inside = (coords[..., 0].abs() <= 1.0) & (coords[..., 1].abs() <= 1.0)  # (batch, H, W)
+        phase = self._query_phase(coords)  # (S, batch, H, W)
+        phase = phase * inside.unsqueeze(0)  # off-object -> phase 0 (vacuum)
+        return torch.exp(1.0j * phase)
+
+    def _grid_coords_xy(self) -> torch.Tensor:
+        """Normalized ``(row, col)`` grid over the padded object, shape ``(H, W, 2)``."""
+        if self._obj_shape is None:
+            raise ValueError("ObjectINR shape not set, call _initialize_obj() first")
+        real_dtype = getattr(torch, config.get("dtype_real"))
+        _, h, w = (int(x) for x in self._obj_shape)
+        ys = torch.linspace(-1.0, 1.0, h, device=self.device, dtype=real_dtype)
+        xs = torch.linspace(-1.0, 1.0, w, device=self.device, dtype=real_dtype)
+        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
+        return torch.stack([gy, gx], dim=-1)  # (H, W, 2)
+
+    def _materialize_obj(self) -> torch.Tensor:
+        with torch.no_grad():
+            phase = self._query_phase(self._grid_coords_xy())  # (S, H, W)
+        return phase
+
+    def pretrain(
+        self,
+        pretrain_target: torch.Tensor | np.ndarray | None = None,
+        num_iters: int = 200,
+        optimizer_params: "dict | OptimizerType | None" = None,
+        scheduler_params: "dict | SchedulerType | None" = None,
+        loss_fn: Callable | str = "l2",
+        device: str | int | None = None,
+        show: bool = True,
+        normalize_object_plotting: bool = True,
+    ) -> None:
+        """Warm-start the INR by fitting it to a target object (e.g. a pixelated recon).
+
+        Queries the INR on the full normalized grid and regresses it onto ``pretrain_target``
+        (a real ``(num_slices, H, W)`` phase array, typically the pixelated reconstruction set by
+        ``from_pixelated``). The fitted weights become the reset state, so a subsequent
+        ``reconstruct(reset=True)`` resumes from this warm start instead of vacuum. A learning-rate
+        scheduler is recommended (pass ``scheduler_params``) as INRs converge better with one.
+        """
+        if device is not None:
+            dev, _ = config.validate_device(device)
+            self.to(dev)
+        if pretrain_target is not None:
+            self.pretrain_target = pretrain_target
+        if self._pretrain_target is None or self._pretrain_target.numel() == 0:
+            raise ValueError(
+                "No pretrain target set; pass pretrain_target or use from_pixelated()."
+            )
+        if self._obj_shape is None:
+            self._obj_shape = tuple(int(x) for x in self._pretrain_target.shape)  # type: ignore[assignment]
+        if optimizer_params is not None:
+            self.set_optimizer(optimizer_params)
+        if scheduler_params is not None:
+            self.set_scheduler(scheduler_params, num_iters)
+        loss_module = get_loss_module(loss_fn, getattr(torch, config.get("dtype_real")))
+        self._pretrain(
+            num_iters, loss_module, show=show, normalize_object_plotting=normalize_object_plotting
+        )
+        self._set_pretrained_weights(self._model)
+        self._invalidate_obj_cache()
+
+    def _pretrain(
+        self,
+        num_iters: int,
+        loss_fn: Callable,
+        show: bool = False,
+        normalize_object_plotting: bool = True,
+    ) -> None:
+        optimizer = self.optimizer
+        if optimizer is None:
+            raise ValueError("Optimizer not set. Pass optimizer_params to pretrain().")
+        scheduler = self.scheduler
+        coords_xy = self._grid_coords_xy()
+        target = self._pretrain_target.to(self.device)
+        self._model.train()
+        pbar = tqdm(range(num_iters))
+        output = self._query_phase(coords_xy)
+        for _ in pbar:
+            optimizer.zero_grad()
+            output = self._query_phase(coords_xy)  # (S, H, W), differentiable
+            loss: torch.Tensor = loss_fn(output, target)
+            loss.backward()
+            optimizer.step()
+            if scheduler is not None:
+                if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                    scheduler.step(loss.item())
+                else:
+                    scheduler.step()
+            self._pretrain_losses.append(loss.item())
+            self._pretrain_lrs.append(optimizer.param_groups[0]["lr"])
+            pbar.set_description(
+                f"Iter {len(self._pretrain_losses)}/{num_iters}, Loss: {loss.item():.3e}"
+            )
+        if show:
+            self.visualize_pretrain(output.detach(), normalize_object_plotting)
+
+    def visualize_pretrain(
+        self, pred_obj: torch.Tensor, normalize_object_plotting: bool = True
+    ) -> None:
+        """Plot the pretraining loss / learning-rate curves and the pred vs target object."""
+        import matplotlib.gridspec as gridspec
+
+        fig = plt.figure(figsize=(12, 6))
+        gs = gridspec.GridSpec(2, 1, height_ratios=[1, 2], hspace=0.3)
+        ax = fig.add_subplot(gs[0])
+        lines = []
+        lines.extend(
+            ax.semilogy(
+                np.arange(len(self._pretrain_losses)), self._pretrain_losses, c="k", label="loss"
+            )
+        )
+        ax.set_ylabel("Loss", color="k")
+        ax.tick_params(axis="y", which="both", colors="k")
+        ax.spines["left"].set_color("k")
+        ax.set_xlabel("Iterations")
+        nx = ax.twinx()
+        nx.spines["left"].set_visible(False)
+        lines.extend(
+            nx.semilogy(
+                np.arange(len(self._pretrain_lrs)),
+                self._pretrain_lrs,
+                c="tab:orange",
+                label="LR",
+            )
+        )
+        labs = [lin.get_label() for lin in lines]
+        nx.legend(lines, labs, loc="upper center")
+        nx.set_ylabel("LRs")
+
+        gs_bot = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=gs[1])
+        axs_bot = np.array([fig.add_subplot(gs_bot[0, i]) for i in range(2)])
+        target = self._pretrain_target
+        norm = None
+        if normalize_object_plotting:
+            target_mean = target.mean(0).cpu().detach().numpy()
+            target_norm = CustomNormalization(interval_type="quantile", data=target_mean)
+            norm = {
+                "interval_type": "manual",
+                "vmin": target_norm.vmin,
+                "vmax": target_norm.vmax,
+            }
+        show_2d(
+            [
+                pred_obj.mean(0).cpu().detach().numpy(),
+                target.mean(0).cpu().detach().numpy(),
+            ],
+            figax=(fig, axs_bot),
+            title=[f"Pred obj ({self.obj_type})", f"Target obj ({self.obj_type})"],
+            cmap="magma",
+            cbar=True,
+            norm=norm,
+        )
+        plt.suptitle(
+            f"Final loss: {self._pretrain_losses[-1]:.3e} | Iters: {len(self._pretrain_losses)}",
+            fontsize=14,
+            y=0.94,
+        )
+        plt.show()
+
+    def _initialize_obj(
+        self,
+        shape: tuple[int, int, int] | np.ndarray,
+        sampling: tuple[float, float] | np.ndarray | None = None,
+    ) -> None:
+        super()._initialize_obj(shape, sampling)
+        shape_t = tuple(int(x) for x in shape)
+        if shape_t[0] != self.num_slices:
+            raise ValueError(
+                f"shape[0] ({shape_t[0]}) does not match num_slices ({self.num_slices})"
+            )
+        self._obj_shape = shape_t  # type: ignore[assignment]
+        self._invalidate_obj_cache()
+
+    def reset(self) -> None:
+        """Reset the INR weights to their initial (pretrained) state."""
+        self._model.load_state_dict(deepcopy(self._pretrained_weights))
+        self._invalidate_obj_cache()
+
+    def to(self, *args, **kwargs):
+        """Move all relevant tensors to a different device."""
+        super().to(*args, **kwargs)
+        self._model = self._model.to(*args, **kwargs)
+        device = kwargs.get("device", args[0] if args else None)
+        if device is not None:
+            self.device = device
+            self._rng_to_device(device)
+            self.reconnect_optimizer_to_parameters()
+        self._invalidate_obj_cache()
+        return self
+
+    # region --- constraints ---
+    def apply_hard_constraints(
+        self, raw: torch.Tensor, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Project the materialized object (display only).
+
+        Unlike the grid-based ``Raster`` constraints, an INR has nothing to clamp or
+        filter in place; for ``pure_phase`` we only recenter the phase to zero mean so
+        the displayed object matches the pixelated convention.
+        """
+        with torch.no_grad():
+            if self.obj_type == "pure_phase":
+                return raw - raw.mean()
+            return raw
+
+    def apply_soft_constraints(
+        self, obj: torch.Tensor | None = None, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Coordinate-sampled total-variation penalty.
+
+        The ``obj`` argument is accepted for interface parity with the grid-based object
+        models but ignored: TV is evaluated at randomly sampled coordinates so the
+        penalty is differentiable w.r.t. the INR weights without materializing the full
+        grid.
+        """
+        self.reset_soft_constraint_losses()
+        loss = self._get_zero_loss_tensor()
+        w_z = self.constraints.tv_weight_z if self.num_slices > 1 else 0.0
+        w_xy = self.constraints.tv_weight_xy
+        if w_z > 0 or w_xy > 0:
+            tv_loss = self._sampled_tv_loss(w_z, w_xy)
+            loss = loss + tv_loss
+            self.add_soft_constraint_loss("tv_loss", tv_loss)
+        self.accumulate_constraint_losses()
+        return loss
+
+    def _sampled_tv_loss(self, w_z: float, w_xy: float, num_samples: int = 4096) -> torch.Tensor:
+        """Finite-difference TV over (z, y, x) at randomly sampled coordinates."""
+        real_dtype = getattr(torch, config.get("dtype_real"))
+        coords_xy = (
+            torch.rand(
+                num_samples, 2, device=self.device, dtype=real_dtype, generator=self._rng_torch
+            )
+            * 2.0
+            - 1.0
+        )
+        phase = self._query_phase(coords_xy)  # (S, num_samples)
+        loss = self._get_zero_loss_tensor()
+        # finite-difference step ~ one pixel in normalized coords
+        if self._obj_shape is not None:
+            h = 2.0 / max(int(self._obj_shape[-1]), int(self._obj_shape[-2]))
+        else:
+            h = 1e-2
+        if w_xy > 0:
+            for axis in range(2):
+                offset = torch.zeros(2, device=self.device, dtype=real_dtype)
+                offset[axis] = h
+                shifted = self._query_phase(coords_xy + offset)
+                loss = loss + w_xy * torch.mean(torch.abs(shifted - phase))
+            loss = loss / 2
+        if w_z > 0 and self.num_slices > 1:
+            loss = loss + w_z * torch.mean(torch.abs(phase[1:] - phase[:-1]))
+        return loss
+
+    # endregion --- constraints ---
+
+    def backward(self, *args, **kwargs):
+        raise NotImplementedError(
+            f"Analytical gradients are not implemented for {self.name}, use autograd=True"
+        )
 
-# constraints are going to be tricky, specifically the TV and filtering if we want to allow
-# multiscale reconstructions
 
-ObjectModelType = ObjectPixelated | ObjectDIP  # | ObjectImplicit
+ObjectModelType = ObjectPixelated | ObjectDIP | ObjectINR
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 4d9345e73..325a66ba9 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -188,9 +188,14 @@ def _soft_constraints(self) -> torch.Tensor:
         """Calculate soft constraints by calling apply_soft_constraints on each model."""
         total_loss = torch.tensor(0, device=self._single_device, dtype=self._dtype_real)
 
-        obj_loss = self.obj_model.apply_soft_constraints(
-            self.obj_model.obj, mask=self.obj_model.mask
-        )
+        if self.obj_model.is_implicit:
+            # Implicit objects evaluate soft constraints at sampled coordinates and don't
+            # need the materialized grid (which would force full-grid inference each iter).
+            obj_loss = self.obj_model.apply_soft_constraints(None, mask=self.obj_model.mask)
+        else:
+            obj_loss = self.obj_model.apply_soft_constraints(
+                self.obj_model.obj, mask=self.obj_model.mask
+            )
         total_loss += obj_loss
 
         probe_loss = self.probe_model.apply_soft_constraints(self.probe_model.probe)
@@ -384,11 +389,11 @@ def _reconstruct_inner(
                 self.zero_grad_all()
                 batch_indices = batch["index"].to(self._single_device)
                 targets = batch["target"].to(self._single_device, non_blocking=True)
-                patch_indices, _positions_px, positions_px_fractional, descan_shifts = (
+                patch_data, _positions_px, positions_px_fractional, descan_shifts = (
                     self.dset.forward(batch_indices, self.obj_padding_px)
                 )
                 shifted_probes = self.probe_model.forward(positions_px_fractional)
-                obj_patches = self.obj_model.forward(patch_indices)
+                obj_patches = self.obj_model.forward(patch_data)
                 propagated_probes, overlap = self.forward_operator(
                     obj_patches, shifted_probes, descan_shifts
                 )
@@ -411,7 +416,7 @@ def _reconstruct_inner(
                     obj_patches,
                     propagated_probes,
                     overlap,
-                    patch_indices,
+                    patch_data,
                     targets,
                 )
                 if _dist_world_size > 1:
@@ -441,11 +446,11 @@ def _reconstruct_inner(
                     for batch in val_loader:
                         batch_indices = batch["index"].to(self._single_device)
                         targets = batch["target"].to(self._single_device, non_blocking=True)
-                        patch_indices, _positions_px, positions_px_fractional, descan_shifts = (
+                        patch_data, _positions_px, positions_px_fractional, descan_shifts = (
                             self.dset.forward(batch_indices, self.obj_padding_px)
                         )
                         shifted_probes = self.probe_model.forward(positions_px_fractional)
-                        obj_patches = self.obj_model.forward(patch_indices)
+                        obj_patches = self.obj_model.forward(patch_data)
                         _propagated_probes, overlap = self.forward_operator(
                             obj_patches, shifted_probes, descan_shifts
                         )
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 6a0f24daf..0f38c79cb 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -528,6 +528,10 @@ def obj_model(self, model: ObjectModelType | type):
         # Set object shape
         model.to(self._single_device)
         self._obj_model = cast(ObjectModelType, model)
+        # Keep the dataset's forward path (coordinates vs. integer patch_indices) in sync with
+        # the object representation. Implicit objects are queried at continuous coordinates.
+        if hasattr(self, "_dset"):
+            self.dset.implicit_object = model.is_implicit
 
     @property
     def probe_model(self) -> ProbeModelType:

From 0c55030f9b53d68dfcd952e12ff2ea4f8e0f9896 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 29 May 2026 16:31:28 -0700
Subject: [PATCH 27/59] fixing linter errors

---
 src/quantem/diffractive_imaging/object_models.py     | 7 ++++---
 src/quantem/diffractive_imaging/ptychography.py      | 6 +++---
 src/quantem/diffractive_imaging/ptychography_base.py | 4 +++-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 78d713bb0..832fe4a0b 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -361,7 +361,8 @@ def params(self):
         raise NotImplementedError()
 
     @abstractmethod
-    def forward(self, patch_indices: torch.Tensor):
+    def forward(self, patch_indices: torch.Tensor, /):
+        # positional-only: implicit object models accept coordinates here instead of indices
         raise NotImplementedError()
 
     @abstractmethod
@@ -1497,7 +1498,7 @@ def visualize_pretrain(
         plt.show()
 
 
-class ObjectINR(ObjectConstraints):
+class ObjectINR(BaseConstraints[PtychoObjConstraintParams.INR], ObjectBase):
     """Implicit (coordinate-queried) object model.
 
     Wraps an implicit neural representation (INR; an ``HSiren`` by default) that maps
@@ -1632,7 +1633,7 @@ def from_uniform(
         )
         # Zero the final linear layer so the INR outputs phase 0 everywhere (vacuum start).
         with torch.no_grad():
-            final_linear = model.net[-2]
+            final_linear = cast(nn.Linear, model.net[-2])
             final_linear.weight.zero_()
             if final_linear.bias is not None:
                 final_linear.bias.zero_()
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 325a66ba9..0c0d47aa3 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -21,7 +21,7 @@
 from quantem.diffractive_imaging.dataset_models import DatasetModelType
 from quantem.diffractive_imaging.detector_models import DetectorModelType
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
-from quantem.diffractive_imaging.object_models import ObjectModelType, ObjectPixelated
+from quantem.diffractive_imaging.object_models import ObjectINR, ObjectModelType, ObjectPixelated
 from quantem.diffractive_imaging.probe_models import ProbeModelType, ProbeParametric
 from quantem.diffractive_imaging.ptycho_utils import compute_train_val_split
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase
@@ -188,10 +188,10 @@ def _soft_constraints(self) -> torch.Tensor:
         """Calculate soft constraints by calling apply_soft_constraints on each model."""
         total_loss = torch.tensor(0, device=self._single_device, dtype=self._dtype_real)
 
-        if self.obj_model.is_implicit:
+        if isinstance(self.obj_model, ObjectINR):
             # Implicit objects evaluate soft constraints at sampled coordinates and don't
             # need the materialized grid (which would force full-grid inference each iter).
-            obj_loss = self.obj_model.apply_soft_constraints(None, mask=self.obj_model.mask)
+            obj_loss = self.obj_model.apply_soft_constraints(mask=self.obj_model.mask)
         else:
             obj_loss = self.obj_model.apply_soft_constraints(
                 self.obj_model.obj, mask=self.obj_model.mask
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 0f38c79cb..99a5072e8 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -898,7 +898,9 @@ def reset_recon(self) -> None:
         self.probe_model.reset()
         self.dset.reset()
         self.compute_propagator_arrays()
-        self.obj_model.constraints = self.obj_model.DEFAULT_CONSTRAINTS
+        # obj_model and its DEFAULT_CONSTRAINTS are correlated at runtime (each object type pairs
+        # with its own constraint dataclass), which the union type can't express.
+        self.obj_model.constraints = self.obj_model.DEFAULT_CONSTRAINTS  # pyright: ignore[reportAttributeAccessIssue]
         # detector reset if necessary
         self._iter_losses = []
         self._iter_val_losses = []

From 5488cd97d77c47f9b00c77072ce0024df7b40a58 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Mon, 1 Jun 2026 15:48:37 -0700
Subject: [PATCH 28/59] adding loss criterion organization, testing inr with
 s3im and smoothl1

---
 .../diffractive_imaging/dataset_models.py     |  34 +-
 .../diffractive_imaging/ptycho_losses.py      | 182 +++++++
 .../diffractive_imaging/ptychography.py       |  24 +-
 .../diffractive_imaging/ptychography_base.py  |  43 +-
 tests/diffractive_imaging/test_object_inr.py  | 498 ++++++++++++++++++
 5 files changed, 740 insertions(+), 41 deletions(-)
 create mode 100644 src/quantem/diffractive_imaging/ptycho_losses.py
 create mode 100644 tests/diffractive_imaging/test_object_inr.py

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 473eb8db5..dd8fb2492 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -327,27 +327,25 @@ def target_residency(self, value: str) -> None:
             raise ValueError(f"target_residency must be 'device' or 'cpu', got {value!r}")
         self._target_residency = value
 
-    def _set_targets(
-        self,
-        loss_type: Literal[
-            "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
-        ],
-    ):
+    def _set_targets(self, target_space: Literal["amplitude", "intensity"]):
+        """Build the per-position targets in the given measurement space.
+
+        ``target_space`` is the space the data-fidelity criterion compares in (set from the
+        criterion's ``target_space``); the comparison itself lives in the criterion, not here.
+        """
         # When residency is "cpu", build targets on CPU so they can be streamed per-batch
         # (and read by DataLoader workers); otherwise keep them resident on the compute device.
         target_device = "cpu" if self.target_residency == "cpu" else self.device
-        if "amplitude" in loss_type:
-            if self.learn_descan and self.has_optimizer():
-                self._targets = self.amplitudes.clone().to(target_device)
-            else:
-                self._targets = self.centered_amplitudes.clone().to(target_device)
-        elif "intensity" in loss_type or loss_type == "poisson":
-            if self.learn_descan and self.has_optimizer():
-                self._targets = self.intensities.clone().to(target_device)
-            else:
-                self._targets = self.centered_intensities.clone().to(target_device)
+        learn_descan = self.learn_descan and self.has_optimizer()
+        if target_space == "amplitude":
+            source = self.amplitudes if learn_descan else self.centered_amplitudes
+        elif target_space == "intensity":
+            source = self.intensities if learn_descan else self.centered_intensities
         else:
-            raise ValueError(f"Unknown loss type {loss_type}")
+            raise ValueError(
+                f"target_space must be 'amplitude' or 'intensity', got {target_space!r}"
+            )
+        self._targets = source.clone().to(target_device)
 
     @property
     def patch_indices(self) -> torch.Tensor:
@@ -1177,7 +1175,7 @@ def preprocess(
         self._set_initial_scan_positions_px(obj_padding_px)
         self._set_patch_indices(obj_padding_px)
 
-        self._set_targets("l2_amplitude")
+        self._set_targets("amplitude")  # default space; criterion-driven space set in reconstruct
 
         self._preprocessed = True
         return
diff --git a/src/quantem/diffractive_imaging/ptycho_losses.py b/src/quantem/diffractive_imaging/ptycho_losses.py
new file mode 100644
index 000000000..d1ac180ea
--- /dev/null
+++ b/src/quantem/diffractive_imaging/ptycho_losses.py
@@ -0,0 +1,182 @@
+"""Data-fidelity criteria for iterative ptychography.
+
+A *criterion* decouples the two concerns that used to be tangled in ``error_estimate``:
+
+1. the **measurement space** it compares in (``target_space``: ``"amplitude"`` or
+   ``"intensity"``) — this is what ``PtychographyDatasetBase._set_targets`` builds targets in
+   and what predictions are mapped to; and
+2. the **comparison** itself (``__call__``) — L2, L1, smooth-L1, Poisson, S3IM, ....
+
+``error_estimate`` masks the predictions/targets, calls the criterion, and divides the result by
+the mean diffraction intensity. To add a new loss, write a ``DataCriterion`` subclass and register
+it in ``_REGISTRY`` (or pass an instance straight to ``reconstruct(loss_type=...)``).
+"""
+
+from typing import Callable, Literal
+
+import torch
+import torch.nn.functional as F
+
+TargetSpace = Literal["amplitude", "intensity"]
+
+
+def _global_scale(preds: torch.Tensor, n: int) -> float:
+    """Batch -> full-scan normalization: a batch sum is rescaled to a full-dataset-equivalent
+    sum so the loss magnitude is independent of batch size (``B / n``)."""
+    return preds.shape[0] / n
+
+
+class DataCriterion:
+    """Base class for ptychography data-fidelity criteria.
+
+    Subclasses set ``target_space`` and implement ``__call__(preds, targets, n)``, returning a
+    scalar error (before the mean-intensity normalization applied by ``error_estimate``). ``preds``
+    and ``targets`` are already in ``target_space`` and detector-masked; ``n`` is the global scan
+    count (for batch-size-independent scaling).
+    """
+
+    target_space: TargetSpace = "amplitude"
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class L2(DataCriterion):
+    """Sum of squared residuals in ``target_space`` (amplitude or intensity)."""
+
+    def __init__(self, target_space: TargetSpace = "amplitude"):
+        self.target_space = target_space
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        return torch.sum((preds - targets) ** 2) / _global_scale(preds, n)
+
+
+class L1(DataCriterion):
+    """Sum of absolute residuals in ``target_space`` (amplitude or intensity)."""
+
+    def __init__(self, target_space: TargetSpace = "amplitude"):
+        self.target_space = target_space
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        return torch.sum((preds - targets).abs()) / _global_scale(preds, n)
+
+
+class Poisson(DataCriterion):
+    target_space: TargetSpace = "intensity"
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        return torch.sum(preds - targets * torch.log(preds + 1e-6))
+
+
+class AmplitudeSmoothL1(DataCriterion):
+    """Smooth-L1 (Huber) on the amplitude residual: ~L2 for ``|r| < beta``, ~L1 beyond."""
+
+    target_space: TargetSpace = "amplitude"
+
+    def __init__(self, beta: float = 1.0):
+        self.beta = float(beta)
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        return F.smooth_l1_loss(preds, targets, beta=self.beta, reduction="sum") / _global_scale(
+            preds, n
+        )
+
+
+def _gaussian_window(size: int, sigma: float, device, dtype) -> torch.Tensor:
+    coords = torch.arange(size, device=device, dtype=dtype) - (size - 1) / 2
+    g = torch.exp(-(coords**2) / (2 * sigma**2))
+    g = g / g.sum()
+    w2d = g[:, None] * g[None, :]
+    return w2d[None, None]  # (1, 1, size, size)
+
+
+def _ssim(x: torch.Tensor, y: torch.Tensor, window: torch.Tensor) -> torch.Tensor:
+    """Mean SSIM between two single-channel images ``x``, ``y`` of shape ``(1, 1, H, W)``."""
+    c1, c2 = 0.01**2, 0.03**2
+    mu_x = F.conv2d(x, window)
+    mu_y = F.conv2d(y, window)
+    mu_x2, mu_y2, mu_xy = mu_x**2, mu_y**2, mu_x * mu_y
+    sigma_x = F.conv2d(x * x, window) - mu_x2
+    sigma_y = F.conv2d(y * y, window) - mu_y2
+    sigma_xy = F.conv2d(x * y, window) - mu_xy
+    ssim_map = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / (
+        (mu_x2 + mu_y2 + c1) * (sigma_x + sigma_y + c2)
+    )
+    return ssim_map.mean()
+
+
+class AmplitudeS3IM(DataCriterion):
+    """Stochastic Structural SIMilarity loss (Xie et al. 2023), as ``MSE + lambda * (1 - S3IM)``.
+
+    S3IM applies SSIM to *non-local* groups of pixels: the flattened predictions/targets are
+    randomly permuted ``repeats`` times, tiled into a 2D patch, and compared with windowed SSIM;
+    this captures structural relationships a per-pixel loss misses. It is used as an auxiliary
+    term on top of an MSE term (both mean-reduced here, so ``lambda`` ~ O(1) balances them). The
+    SSIM passes make this notably more expensive than L2 — keep ``repeats`` modest.
+    """
+
+    target_space: TargetSpace = "amplitude"
+
+    def __init__(
+        self,
+        lambda_s3im: float = 1.0,
+        repeats: int = 5,
+        patch_height: int = 32,
+        window_size: int = 11,
+        sigma: float = 1.5,
+    ):
+        self.lambda_s3im = float(lambda_s3im)
+        self.repeats = int(repeats)
+        self.patch_height = int(patch_height)
+        self.window_size = int(window_size)
+        self.sigma = float(sigma)
+
+    def _s3im(self, src: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        num = src.numel()
+        idx_list = [torch.arange(num, device=src.device)]
+        for _ in range(self.repeats - 1):
+            idx_list.append(torch.randperm(num, device=src.device))
+        idx = torch.cat(idx_list)
+        ph = self.patch_height
+        usable = (idx.numel() // ph) * ph  # trim so it reshapes to (ph, -1)
+        idx = idx[:usable]
+        src_img = src[idx].reshape(1, 1, ph, -1)
+        tar_img = tar[idx].reshape(1, 1, ph, -1)
+        window = _gaussian_window(self.window_size, self.sigma, src.device, src.dtype)
+        return 1.0 - _ssim(src_img, tar_img, window)
+
+    def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
+        mse = torch.mean((preds - targets) ** 2)
+        s3im = self._s3im(preds.reshape(-1), targets.reshape(-1))
+        return mse + self.lambda_s3im * s3im
+
+
+_REGISTRY: dict[str, Callable[[], DataCriterion]] = {
+    "l2_amplitude": lambda: L2("amplitude"),
+    "l1_amplitude": lambda: L1("amplitude"),
+    "l2_intensity": lambda: L2("intensity"),
+    "l1_intensity": lambda: L1("intensity"),
+    "poisson": Poisson,
+    "smooth_l1_amplitude": AmplitudeSmoothL1,
+    "s3im_amplitude": AmplitudeS3IM,
+}
+
+
+def get_data_criterion(loss_type: "str | DataCriterion") -> DataCriterion:
+    """Resolve a ``loss_type`` to a :class:`DataCriterion`.
+
+    Accepts a registered name (e.g. ``"l2_amplitude"``, ``"smooth_l1_amplitude"``,
+    ``"s3im_amplitude"``) for the default-configured criterion, or a ``DataCriterion`` instance
+    (use this to tune parameters, e.g. ``AmplitudeS3IM(lambda_s3im=0.5, repeats=10)``).
+    """
+    if isinstance(loss_type, DataCriterion):
+        return loss_type
+    if isinstance(loss_type, str):
+        key = loss_type.lower()
+        if key not in _REGISTRY:
+            raise ValueError(
+                f"Unknown loss_type {loss_type!r}; expected one of {sorted(_REGISTRY)} "
+                "or a DataCriterion instance."
+            )
+        return _REGISTRY[key]()
+    raise TypeError(f"loss_type must be a str or DataCriterion, got {type(loss_type)}")
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 0c0d47aa3..f39056215 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -23,6 +23,7 @@
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
 from quantem.diffractive_imaging.object_models import ObjectINR, ObjectModelType, ObjectPixelated
 from quantem.diffractive_imaging.probe_models import ProbeModelType, ProbeParametric
+from quantem.diffractive_imaging.ptycho_losses import DataCriterion
 from quantem.diffractive_imaging.ptycho_utils import compute_train_val_split
 from quantem.diffractive_imaging.ptychography_base import PtychographyBase
 from quantem.diffractive_imaging.ptychography_opt import PtychographyOpt
@@ -222,9 +223,7 @@ def reconstruct(
         store_snapshots_every: int | None = None,
         device: Literal["cpu", "gpu"] | int | list[int] | None = None,
         autograd: bool = True,
-        loss_type: Literal[
-            "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
-        ] = "l2_amplitude",
+        loss_type: "str | DataCriterion" = "l2_amplitude",
         num_workers: int = 0,
     ) -> Self:
         """Run iterative ptychography reconstruction.
@@ -246,6 +245,13 @@ def reconstruct(
         Multi-GPU (``device`` is a list) launches worker processes via ``mp.spawn`` when called
         from a notebook, or uses the existing distributed process group when launched with
         ``torchrun``. Only autograd mode is supported for multi-GPU in this release.
+        
+        ``loss_type`` selects the data-fidelity criterion: a registered name
+        (``"l2_amplitude"`` [default], ``"l1_amplitude"``, ``"l2_intensity"``, ``"l1_intensity"``,
+        ``"poisson"``, ``"smooth_l1_amplitude"``, ``"s3im_amplitude"``) or a ``DataCriterion``
+        instance for custom parameters (e.g. ``AmplitudeS3IM(lambda_s3im=0.5)``). See
+        ``ptycho_losses``.
+
         """
         self._check_preprocessed()
 
@@ -322,9 +328,7 @@ def _reconstruct_inner(
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
         autograd: bool = True,
-        loss_type: Literal[
-            "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
-        ] = "l2_amplitude",
+        loss_type: "str | DataCriterion" = "l2_amplitude",
         num_workers: int = 0,
         _dist_rank: int = 0,
         _dist_world_size: int = 1,
@@ -356,7 +360,8 @@ def _reconstruct_inner(
         if new_scheduler:
             self.set_schedulers(self.scheduler_params, num_iter=num_iters)
 
-        self.dset._set_targets(loss_type)
+        self.criterion = loss_type  # resolve name/instance -> DataCriterion
+        self.dset._set_targets(self._criterion.target_space)
         self.compute_propagator_arrays()  # required to avoid issue if stopped learning probe tilt
 
         # Compute the global scan count once — needed to keep loss scale consistent across world
@@ -403,7 +408,6 @@ def _reconstruct_inner(
                     pred_intensities,
                     batch_indices,
                     targets=targets,
-                    loss_type=loss_type,
                     global_n=global_n,
                 )
 
@@ -459,7 +463,6 @@ def _reconstruct_inner(
                             pred_intensities,
                             batch_indices,
                             targets=targets,
-                            loss_type=loss_type,
                             global_n=global_n,
                         )
                         val_consistency_loss += batch_val_loss.item()
@@ -741,6 +744,9 @@ def save(
         if isinstance(skip, (str, type)):
             skip = [skip]
         skip = list(skip)
+        # The data-fidelity criterion is transient config (re-set each reconstruct); don't
+        # serialize it (avoids a dill fallback and keeps the archive model-agnostic).
+        skip.append("_criterion")
 
         # Always skip raw dataset data unless explicitly requested
         if not save_raw_data:
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 99a5072e8..490bb1321 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -32,6 +32,7 @@
 from quantem.diffractive_imaging.logger_ptychography import LoggerPtychography
 from quantem.diffractive_imaging.object_models import ObjectBase, ObjectModelType
 from quantem.diffractive_imaging.probe_models import ProbeBase, ProbeModelType, ProbePixelated
+from quantem.diffractive_imaging.ptycho_losses import DataCriterion, get_data_criterion
 from quantem.diffractive_imaging.ptycho_utils import (
     AffineTransform,
     center_crop_arr,
@@ -77,6 +78,9 @@ class PtychographyBase(RNGMixin, AutoSerialize):
     """
 
     _token = object()
+    # Default data-fidelity criterion (overridden per-instance from `loss_type` in reconstruct).
+    # Class-level so freshly-built and freshly-loaded objects always resolve a criterion.
+    _criterion: DataCriterion = get_data_criterion("l2_amplitude")
 
     def __init__(  # TODO prevent direct instantiation
         self,
@@ -245,6 +249,19 @@ def dset(self, new_dset: DatasetModelType):
             raise TypeError(f"dset should be a PtychographyDataset, got {type(new_dset)}")
         self._dset = new_dset
 
+    @property
+    def criterion(self) -> DataCriterion:
+        """Active data-fidelity criterion. Assign a registered name or a ``DataCriterion``.
+
+        Transient config (re-set from ``loss_type`` each ``reconstruct``, defaults to L2); not
+        serialized.
+        """
+        return self._criterion
+
+    @criterion.setter
+    def criterion(self, value: "str | DataCriterion") -> None:
+        self._criterion = get_data_criterion(value)
+
     @property
     def detector_model(self) -> DetectorModelType:
         return self._detector_model
@@ -1084,26 +1101,24 @@ def error_estimate(
         pred_intensities: torch.Tensor,
         batch_indices: np.ndarray,
         targets: torch.Tensor,
-        loss_type: Literal[
-            "l2_amplitude", "l1_amplitude", "l2_intensity", "l1_intensity", "poisson"
-        ] = "l2_amplitude",
         global_n: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if "amplitude" in loss_type:
-            preds = torch.sqrt(pred_intensities + 1e-9)  # add eps to avoid diverging gradients
+        """Data-fidelity loss for one batch via the active criterion (``self.criterion``).
+
+        Maps predictions into the criterion's measurement space (amplitude or intensity),
+        applies the detector mask, evaluates the criterion, and normalizes by the mean
+        diffraction intensity. Which comparison is used (L2/L1/smooth-L1/Poisson/S3IM/...) is
+        entirely the criterion's concern; see ``ptycho_losses``.
+        """
+        criterion = self.criterion
+        if criterion.target_space == "amplitude":
+            preds = torch.sqrt(pred_intensities + 1e-9)  # eps avoids diverging gradients at 0
         else:
             preds = pred_intensities
 
-        diff = preds * self.dset.detector_mask - targets * self.dset.detector_mask
+        mask = self.dset.detector_mask
         n = global_n if global_n is not None else self.dset.num_gpts
-        if "l1" in loss_type:
-            error = torch.sum(torch.abs(diff)) / (diff.shape[0] / n)
-        elif "l2" in loss_type:
-            error = torch.sum(torch.abs(diff) ** 2) / (diff.shape[0] / n)
-        elif loss_type == "poisson":
-            error = torch.sum(preds - targets * torch.log(preds + 1e-6))
-        else:
-            raise ValueError(f"Unknown loss type {loss_type}, should be 'l1' or 'l2'")
+        error = criterion(preds * mask, targets * mask, n)
         loss = error / self.dset.mean_diffraction_intensity
         return loss, targets
 
diff --git a/tests/diffractive_imaging/test_object_inr.py b/tests/diffractive_imaging/test_object_inr.py
new file mode 100644
index 000000000..c22c704dc
--- /dev/null
+++ b/tests/diffractive_imaging/test_object_inr.py
@@ -0,0 +1,498 @@
+"""
+Tests for the implicit (INR) object model, ``ObjectINR``.
+
+Covers the object model in isolation (forward/obj shapes, vacuum init, off-object
+masking, multislice z-coordinates, gradient flow), the dataset's implicit-object
+coordinate production, and an end-to-end reconstruction on a smooth synthetic object.
+
+The reconstruction fixtures deliberately differ from ``test_ptychography.py``: that
+fixture scans a torus-wrapped object (edge positions, zero padding), which the
+pixelated path reproduces via ``% obj_shape`` indexing but the INR (vacuum outside the
+object) cannot. Here the ground-truth object is larger than the scanned region and the
+scan is confined to the interior with padding >= roi // 2, so patches never reach the
+object boundary -- the physically realistic, non-toroidal regime both models agree on.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.core.datastructures.dataset4dstem import Dataset4dstem
+from quantem.core.io.serialize import load as autoserialize_load
+from quantem.core.ml import OptimizerParams, SchedulerParams
+from quantem.core.ml.cnn import CNN2d
+from quantem.core.utils.utils import electron_wavelength_angstrom
+from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster
+from quantem.diffractive_imaging.detector_models import DetectorPixelated
+from quantem.diffractive_imaging.object_models import ObjectINR, ObjectPixelated
+from quantem.diffractive_imaging.probe_models import ProbeDIP, ProbeParametric, ProbePixelated
+from quantem.diffractive_imaging.ptychography import Ptychography
+
+if config.NUM_DEVICES > 0:
+    config.set_device("gpu")
+
+N = 40  # detector / roi size (px)
+OGT = 64  # ground-truth object size (px); larger than the scanned region
+PAD = 20  # obj padding (>= roi // 2 so interior patches never hit the boundary)
+Q_MAX = 0.5  # inverse Angstroms
+Q_PROBE = Q_MAX / 2
+PROBE_ENERGY = 300e3  # eV
+C10 = 50.0  # defocus (Angstrom)
+STEP = 2  # scan step (px)
+SCAN_START = 20  # first scan position (px); SCAN_START - roi//2 >= 0
+SCAN_STOP = 44  # exclusive; SCAN_STOP - 1 + roi//2 - 1 < OGT  -> no wrap
+
+# Reconstruction config (validated against the fixture below: pixelated reaches corr~0.89,
+# the INR reaches corr~0.91 at omega_0=5 / lr=1e-2). omega_0=5 suits this smooth object;
+# the default of 10 (and the image-fitting default of 30) stall here.
+_INR_OMEGA = 5.0
+_RECON_LR = 1e-2
+_LOSS_RATIO = 0.3
+_CORR_THRESHOLD = 0.7
+
+
+def _smooth_phase() -> np.ndarray:
+    """A smooth, band-limited phase object (friendly to a SIREN, unlike white noise)."""
+    yy, xx = np.meshgrid(np.arange(OGT), np.arange(OGT), indexing="ij")
+    return (0.7 * np.sin(2 * np.pi * xx / OGT * 4) * np.cos(2 * np.pi * yy / OGT * 3)).astype(
+        np.float32
+    )
+
+
+def _probe_array() -> np.ndarray:
+    sampling = 1 / Q_MAX / 2
+    reciprocal_sampling = 2 * Q_MAX / N
+    qx = qy = np.fft.fftfreq(N, sampling)
+    q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2)
+    aperture = np.sqrt(np.clip((Q_PROBE - q) / reciprocal_sampling + 0.5, 0, 1))
+    chi = q**2 * electron_wavelength_angstrom(PROBE_ENERGY) * np.pi * C10
+    probe_fourier = aperture * np.exp(-1j * chi)
+    probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2))
+    return (np.fft.ifft2(probe_fourier) * N).astype(np.complex64)
+
+
+def _semiangle_mrad() -> float:
+    return electron_wavelength_angstrom(PROBE_ENERGY) * Q_PROBE * 1e3
+
+
+def _build_synthetic_dataset() -> tuple[PtychographyDatasetRaster, np.ndarray, np.ndarray]:
+    """Simulate a non-toroidal 4D-STEM dataset; return (dataset_model, gt_phase, probe)."""
+    phase = _smooth_phase()
+    complex_obj = np.exp(1j * phase)
+    probe = _probe_array()
+    reciprocal_sampling = 2 * Q_MAX / N
+
+    gpos = np.arange(SCAN_START, SCAN_STOP, STEP)
+    xx, yy = np.meshgrid(gpos, gpos, indexing="ij")
+    positions = np.stack((xx.ravel(), yy.ravel()), axis=-1)
+    x0 = positions[:, 0].astype(int)
+    y0 = positions[:, 1].astype(int)
+    x_ind = np.fft.fftfreq(N, d=1 / N).astype(int)
+    y_ind = np.fft.fftfreq(N, d=1 / N).astype(int)
+    row = x0[:, None, None] + x_ind[None, :, None]  # no modulo: interior scan, no wrap
+    col = y0[:, None, None] + y_ind[None, None, :]
+    assert row.min() >= 0 and row.max() < OGT and col.min() >= 0 and col.max() < OGT
+    exit_waves = complex_obj[row, col] * probe
+    intensities = np.abs(np.fft.fft2(exit_waves)) ** 2
+
+    sxy = len(gpos)
+    dset = Dataset4dstem.from_array(
+        array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((sxy, sxy, N, N)),
+        sampling=(STEP, STEP, reciprocal_sampling, reciprocal_sampling),
+        units=("A", "A", "A^-1", "A^-1"),
+    )
+    pdset = PtychographyDatasetRaster.from_dataset4dstem(dset)
+    pdset.learn_scan_positions = False
+    pdset.learn_descan = False
+    pdset.preprocess(
+        com_fit_function="constant",
+        plot_rotation=False,
+        plot_com=False,
+        probe_energy=PROBE_ENERGY,
+    )
+    return pdset, phase, probe
+
+
+def _build_inr_ptycho(
+    num_slices: int = 1, slice_thicknesses=None, first_omega_0: float = _INR_OMEGA
+) -> tuple[Ptychography, np.ndarray]:
+    """Build an ObjectINR ptychography with the exact (frozen) simulation probe."""
+    pdset, gt_phase, probe = _build_synthetic_dataset()
+    obj = ObjectINR.from_uniform(
+        num_slices=num_slices,
+        slice_thicknesses=slice_thicknesses,
+        obj_type="pure_phase",
+        hidden_features=128,
+        first_omega_0=first_omega_0,
+        hidden_omega_0=first_omega_0,
+        rng=0,
+    )
+    probe_model = ProbePixelated.from_array(
+        num_probes=1,
+        probe_params={"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()},
+        probe_array=probe.copy(),
+    )
+    ptycho = Ptychography.from_models(
+        dset=pdset,
+        obj_model=obj,
+        probe_model=probe_model,
+        detector_model=DetectorPixelated(),
+        rng=0,
+        verbose=False,
+    )
+    ptycho.preprocess(obj_padding_px=(PAD, PAD))
+    return ptycho, gt_phase
+
+
+def _corr(a: np.ndarray, b: np.ndarray) -> float:
+    a = a - a.mean()
+    b = b - b.mean()
+    return float((a * b).sum() / np.sqrt((a**2).sum() * (b**2).sum() + 1e-12))
+
+
+def _center_crop(a: np.ndarray, s: int) -> np.ndarray:
+    r0 = (a.shape[0] - s) // 2
+    c0 = (a.shape[1] - s) // 2
+    return a[r0 : r0 + s, c0 : c0 + s]
+
+
+def _best_corr(recon: np.ndarray, gt: np.ndarray, s: int = 32) -> float:
+    """Best |correlation| over small integer shifts (handles ~1px registration offset)."""
+    g = _center_crop(gt, s)
+    best = -1.0
+    for dr in range(-3, 4):
+        for dc in range(-3, 4):
+            r = _center_crop(np.roll(recon, (dr, dc), (0, 1)), s)
+            best = max(best, abs(_corr(r, g)))
+    return best
+
+
+# --------------------------------------------------------------------------- #
+# ObjectINR in isolation
+# --------------------------------------------------------------------------- #
+class TestObjectINRUnit:
+    def test_forward_shape_dtype_and_vacuum(self):
+        obj = ObjectINR.from_uniform(num_slices=1, hidden_features=64, rng=0)
+        obj._initialize_obj((1, 32, 40))
+        coords = torch.rand(5, 8, 8, 2) * 2 - 1
+        patches = obj.forward(coords)
+        assert patches.shape == (1, 5, 8, 8)
+        assert patches.is_complex()
+        # vacuum init (zeroed final layer) -> unit transmission everywhere
+        assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6)
+
+    def test_off_object_is_vacuum(self):
+        obj = ObjectINR.from_uniform(num_slices=1, hidden_features=64, rng=0)
+        obj._initialize_obj((1, 16, 16))
+        # train the final layer a little so the INR is not identically zero
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
+        inside = torch.rand(3, 4, 4, 2) * 2 - 1
+        for _ in range(5):
+            opt.zero_grad()
+            obj.forward(inside).imag.sum().backward()
+            opt.step()
+        # coordinates outside [-1, 1] must map to identity transmission regardless of weights
+        off = torch.full((1, 2, 2, 2), 5.0)
+        out = obj.forward(off)
+        assert torch.allclose(out, torch.ones_like(out), atol=1e-6)
+
+    def test_multislice_z_coordinates(self):
+        obj = ObjectINR.from_uniform(
+            num_slices=3, slice_thicknesses=2.0, hidden_features=32, rng=0
+        )
+        obj._initialize_obj((3, 16, 16))
+        z = obj._z_coords
+        assert z.shape == (3,)
+        # equally spaced slices span [-1, 1]
+        assert torch.allclose(z, torch.tensor([-1.0, 0.0, 1.0]), atol=1e-6)
+        patches = obj.forward(torch.rand(4, 6, 6, 2) * 2 - 1)
+        assert patches.shape == (3, 4, 6, 6)
+
+    def test_single_slice_z_is_zero(self):
+        obj = ObjectINR.from_uniform(num_slices=1, hidden_features=32, rng=0)
+        obj._initialize_obj((1, 8, 8))
+        assert torch.allclose(obj._z_coords, torch.zeros(1))
+
+    def test_obj_materialization(self):
+        obj = ObjectINR.from_uniform(
+            num_slices=2, slice_thicknesses=1.0, hidden_features=32, rng=0
+        )
+        obj._initialize_obj((2, 20, 24))
+        materialized = obj.obj
+        assert materialized.shape == (2, 20, 24)
+        assert not materialized.is_complex()  # pure_phase -> real phase array
+        # vacuum init -> phase 0
+        assert float(materialized.abs().max()) == pytest.approx(0.0, abs=1e-6)
+
+    def test_gradients_flow_and_fit_smooth_phase(self):
+        """The INR should be able to fit a smooth target phase via autograd."""
+        obj = ObjectINR.from_uniform(num_slices=1, hidden_features=64, rng=1)
+        obj._initialize_obj((1, 24, 24))
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-3)
+        ys = torch.linspace(-1, 1, 24)
+        xs = torch.linspace(-1, 1, 24)
+        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
+        target = 0.5 * torch.sin(3 * gy) * torch.cos(2 * gx)
+        coords = torch.stack([gy, gx], dim=-1)[None]  # (1, 24, 24, 2)
+
+        losses = []
+        for it in range(120):
+            opt.zero_grad()
+            pred_phase = obj.forward(coords)[0, 0].angle()
+            loss = ((pred_phase - target) ** 2).mean()
+            loss.backward()
+            if it == 0:
+                # zero-init final layer: after the first step only the final layer has a
+                # gradient. parameters() order is [first-layer weight, bias, ..., final weight,
+                # final bias], so params[0] is the first-layer weight and params[-2] the final.
+                grads = [p.grad for p in obj.model.parameters()]
+                assert grads[0] is not None and float(grads[0].abs().sum()) == 0.0
+                assert grads[-2] is not None and float(grads[-2].abs().sum()) > 0.0
+            opt.step()
+            losses.append(loss.item())
+        assert losses[-1] < 0.05 * losses[0]
+
+    def test_autograd_only_backward_raises(self):
+        obj = ObjectINR.from_uniform(num_slices=1, hidden_features=32, rng=0)
+        with pytest.raises(NotImplementedError):
+            obj.backward()
+
+    def test_complex_obj_type_not_supported(self):
+        with pytest.raises(NotImplementedError):
+            ObjectINR.from_uniform(num_slices=1, obj_type="complex")
+
+    def test_from_pixelated_pretrain(self):
+        """from_pixelated + pretrain warm-starts the INR to reproduce a pixelated object."""
+        h = w = 32
+        ys = torch.linspace(-1, 1, h)
+        xs = torch.linspace(-1, 1, w)
+        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
+        phase = (0.5 * torch.sin(3 * gy) * torch.cos(2 * gx)).float()[None]  # (1, h, w)
+        pix = ObjectPixelated.from_array(initial_obj=phase, obj_type="pure_phase")
+        pix._initialize_obj((1, h, w), sampling=(1.0, 1.0))
+
+        inr = ObjectINR.from_pixelated(
+            pix, hidden_features=128, first_omega_0=5.0, hidden_omega_0=5.0
+        )
+        assert inr.num_slices == pix.num_slices
+        assert tuple(inr.pretrain_target.shape) == tuple(pix.obj.shape)
+
+        inr.pretrain(
+            num_iters=150,
+            optimizer_params=OptimizerParams.Adam(lr=1e-3),
+            scheduler_params=SchedulerParams.Plateau(factor=0.5),
+            show=False,
+        )
+        losses = inr.pretrain_losses
+        assert losses[-1] < 0.1 * losses[0]
+
+        gt = pix.obj[0].detach().cpu().numpy()
+        gt -= gt.mean()
+
+        def _corr_to_pix(arr: np.ndarray) -> float:
+            a = arr - arr.mean()
+            return float((a * gt).sum() / np.sqrt((a**2).sum() * (gt**2).sum() + 1e-12))
+
+        assert _corr_to_pix(inr.obj[0].detach().cpu().numpy()) > 0.95
+        # the pretrained weights are the reset state (reconstruct(reset=True) resumes from them)
+        inr.reset()
+        assert _corr_to_pix(inr.obj[0].detach().cpu().numpy()) > 0.95
+
+
+# --------------------------------------------------------------------------- #
+# Dataset implicit-object coordinate production
+# --------------------------------------------------------------------------- #
+class TestImplicitDatasetCoords:
+    def test_scan_coords_normalization(self):
+        """An integer scan position maps to linspace(-1,1,N) grid nodes; spacing = 1 px."""
+        pdset, _, _ = _build_synthetic_dataset()
+        pdset.implicit_object = True
+        padding = (PAD, PAD)
+        h_full, w_full = pdset._obj_shape_full_2d(padding)
+        # place a clean integer position at the object center
+        r, c = int(h_full) // 2, int(w_full) // 2
+        pdset.scan_positions_px.data[0] = torch.tensor(
+            [float(r), float(c)], device=pdset.scan_positions_px.device
+        )
+        coords = pdset._scan_coords(torch.tensor([0]), padding)[0].detach()  # (Hroi, Wroi, 2)
+
+        # fftfreq offset 0 is the first ROI pixel -> coordinate of pixel (r, c)
+        assert float(coords[0, 0, 0]) == pytest.approx(r / (int(h_full) - 1) * 2 - 1, abs=1e-5)
+        assert float(coords[0, 0, 1]) == pytest.approx(c / (int(w_full) - 1) * 2 - 1, abs=1e-5)
+        # adjacent rows/cols differ by exactly one normalized pixel
+        d_row = float(coords[1, 0, 0] - coords[0, 0, 0])
+        d_col = float(coords[0, 1, 1] - coords[0, 0, 1])
+        assert d_row == pytest.approx(2 / (int(h_full) - 1), abs=1e-5)
+        assert d_col == pytest.approx(2 / (int(w_full) - 1), abs=1e-5)
+
+    def test_forward_returns_coords_and_zero_fractional_when_implicit(self):
+        pdset, _, _ = _build_synthetic_dataset()
+        pdset.implicit_object = True
+        batch = torch.arange(4)
+        patch_data, positions_px, fractional, _descan = pdset.forward(batch, (PAD, PAD))
+        # implicit: patch_data are float coords (batch, Hroi, Wroi, 2), not integer indices
+        assert patch_data.shape == (4, N, N, 2)
+        assert patch_data.dtype.is_floating_point
+        # the probe must not be subpixel-shifted -> fractional is zero
+        assert torch.allclose(fractional, torch.zeros_like(fractional))
+
+    def test_implicit_flag_synced_from_obj_model(self):
+        ptycho, _ = _build_inr_ptycho()
+        assert ptycho.obj_model.is_implicit is True
+        assert ptycho.dset.implicit_object is True
+
+
+# --------------------------------------------------------------------------- #
+# End-to-end reconstruction
+# --------------------------------------------------------------------------- #
+@pytest.mark.slow
+class TestObjectINRReconstruction:
+    def test_loss_decreases_and_recovers_object(self):
+        ptycho, gt_phase = _build_inr_ptycho()
+        ptycho.reconstruct(
+            num_iters=150,
+            optimizer_params={"object": {"name": "adam", "lr": _RECON_LR}},  # probe frozen
+            batch_size=200,
+        )
+        losses = np.array(ptycho._iter_losses)
+        assert losses[-1] < _LOSS_RATIO * losses[0]
+        assert _best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD
+
+    def test_multislice_runs(self):
+        ptycho, _ = _build_inr_ptycho(num_slices=2, slice_thicknesses=20.0)
+        ptycho.reconstruct(
+            num_iters=5,
+            optimizer_params={"object": {"name": "adam", "lr": _RECON_LR}},
+            batch_size=200,
+        )
+        assert ptycho.obj.shape[0] == 2
+
+    def test_data_loss_criteria_run(self):
+        """The pluggable data-fidelity criteria run end-to-end and stay finite."""
+        from quantem.diffractive_imaging.ptycho_losses import AmplitudeS3IM
+
+        ptycho, _ = _build_inr_ptycho()
+        for loss_type in ["l1_amplitude", "smooth_l1_amplitude", AmplitudeS3IM(repeats=3)]:
+            ptycho.reconstruct(
+                num_iters=5,
+                reset=True,
+                optimizer_params={"object": {"name": "adam", "lr": _RECON_LR}},
+                batch_size=200,
+                loss_type=loss_type,
+            )
+            losses = np.array(ptycho._iter_losses)
+            assert np.isfinite(losses).all(), loss_type
+
+    def test_save_load_roundtrip(self, tmp_path):
+        ptycho, _ = _build_inr_ptycho()
+        ptycho.reconstruct(
+            num_iters=20,
+            optimizer_params={"object": {"name": "adam", "lr": _RECON_LR}},
+            batch_size=200,
+        )
+        obj_before = ptycho.obj.copy()
+        path = tmp_path / "inr_ptycho.zip"
+        ptycho.save(path, mode="o", save_raw_data=True)  # persist dset so loaded.dset works
+        loaded = autoserialize_load(path)
+        assert loaded.obj_model.is_implicit is True
+        assert loaded.dset.implicit_object is True
+        np.testing.assert_allclose(loaded.obj, obj_before, rtol=1e-5, atol=1e-6)
+        # continued training still runs after reload
+        loaded.reconstruct(
+            num_iters=5,
+            optimizer_params={"object": {"name": "adam", "lr": _RECON_LR}},
+            batch_size=200,
+        )
+
+
+# --------------------------------------------------------------------------- #
+# ObjectINR composes with every probe model (it is a drop-in object model)
+# --------------------------------------------------------------------------- #
+@pytest.mark.slow
+class TestObjectINRProbeTypes:
+    """ObjectINR is a drop-in object model: it composes with each probe representation.
+
+    For all probe types the implicit object is queried at continuous coordinates (the dataset
+    reports ``implicit_object=True``) and the probe receives a zero fractional shift, so the
+    reconstruction runs end-to-end and the loss stays finite.
+    """
+
+    def _probe_params(self) -> dict:
+        return {"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()}
+
+    def test_inr_runs_with_each_probe_type(self):
+        pdset, _gt, _probe = _build_synthetic_dataset()
+        pp = self._probe_params()
+        probe_builders = {
+            "pixelated": lambda: ProbePixelated.from_params(probe_params=pp),
+            "parametric": lambda: ProbeParametric.from_params(probe_params=pp),
+            "dip": lambda: ProbeDIP.from_model(
+                model=CNN2d(in_channels=1, dtype=torch.complex64, num_layers=3),
+                roi_shape=(N, N),
+                num_probes=1,
+                probe_params=pp,
+            ),
+        }
+        for name, build in probe_builders.items():
+            obj = ObjectINR.from_uniform(
+                num_slices=1,
+                obj_type="pure_phase",
+                hidden_features=128,
+                first_omega_0=_INR_OMEGA,
+                hidden_omega_0=_INR_OMEGA,
+                rng=0,
+            )
+            ptycho = Ptychography.from_models(
+                dset=pdset,
+                obj_model=obj,
+                probe_model=build(),
+                detector_model=DetectorPixelated(),
+                rng=0,
+                verbose=False,
+            )
+            ptycho.preprocess(obj_padding_px=(PAD, PAD))
+            assert ptycho.dset.implicit_object is True, name
+            ptycho.reconstruct(
+                num_iters=10,
+                optimizer_params={
+                    "object": {"name": "adam", "lr": _RECON_LR},
+                    "probe": {"name": "adam", "lr": 1e-3},
+                },
+                batch_size=200,
+            )
+            losses = np.array(ptycho._iter_losses)
+            assert np.isfinite(losses).all(), name
+            assert ptycho.obj.shape[0] == 1, name
+
+
+# --------------------------------------------------------------------------- #
+# Data-fidelity criterion system (ptycho_losses)
+# --------------------------------------------------------------------------- #
+class TestDataCriteria:
+    def test_registry_and_target_spaces(self):
+        from quantem.diffractive_imaging.ptycho_losses import (
+            L2,
+            AmplitudeS3IM,
+            get_data_criterion,
+        )
+
+        assert isinstance(get_data_criterion("l2_amplitude"), L2)
+        assert get_data_criterion("l2_amplitude").target_space == "amplitude"
+        assert get_data_criterion("l2_intensity").target_space == "intensity"
+        assert get_data_criterion("poisson").target_space == "intensity"
+        assert get_data_criterion("s3im_amplitude").target_space == "amplitude"
+        # passing a DataCriterion instance returns it unchanged (tune params this way)
+        crit = AmplitudeS3IM(lambda_s3im=0.5)
+        assert get_data_criterion(crit) is crit
+        with pytest.raises(ValueError):
+            get_data_criterion("not_a_loss")
+
+    def test_criterion_values(self):
+        from quantem.diffractive_imaging.ptycho_losses import L1, L2
+
+        preds = torch.tensor([[1.0, 2.0]])  # B = 1
+        targets = torch.tensor([[1.5, 2.0]])
+        # n == B -> global scale 1; matches the legacy sum-reduced amplitude losses
+        assert torch.isclose(L2()(preds, targets, n=1), torch.tensor(0.25))
+        assert torch.isclose(L1()(preds, targets, n=1), torch.tensor(0.5))

From 990430a6a6517b7a98036640399f5e2a98f6e113 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Mon, 1 Jun 2026 16:07:49 -0700
Subject: [PATCH 29/59] bufix of saving criterion, potential INR to softplus
 final activation by default

---
 .../diffractive_imaging/object_models.py      | 77 +++++++++++++------
 .../diffractive_imaging/ptychography_base.py  | 12 +--
 tests/diffractive_imaging/test_object_inr.py  | 32 ++++++++
 3 files changed, 93 insertions(+), 28 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 832fe4a0b..20d1c65b3 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1502,8 +1502,9 @@ class ObjectINR(BaseConstraints[PtychoObjConstraintParams.INR], ObjectBase):
     """Implicit (coordinate-queried) object model.
 
     Wraps an implicit neural representation (INR; an ``HSiren`` by default) that maps
-    normalized 3D coordinates ``(z, y, x)`` in ``[-1, 1]`` to the object's real phase
-    (``obj_type="pure_phase"``). Rather than gathering grid-aligned patches at integer
+    normalized 3D coordinates ``(z, y, x)`` in ``[-1, 1]`` to a real-valued object — the phase
+    for ``obj_type="pure_phase"`` or the potential for ``obj_type="potential"``, both wrapped to
+    the complex transmission ``exp(1j * value)``. Rather than gathering grid-aligned patches at integer
     scan positions like ``ObjectPixelated``, the paired dataset produces continuous
     per-patch ``(y, x)`` coordinates at the *true* (fractional) scan positions; this
     model augments them with each slice's ``z`` coordinate, queries the INR, and returns
@@ -1549,10 +1550,10 @@ def __init__(
             rng=rng,
             _token=_token,
         )
-        if self.obj_type != "pure_phase":
+        if self.obj_type == "complex":
             raise NotImplementedError(
-                f"ObjectINR currently only supports obj_type='pure_phase', got '{self.obj_type}'. "
-                "Complex/potential INR objects are planned."
+                "ObjectINR does not support obj_type='complex' yet (planned); use 'pure_phase' "
+                "or 'potential' (both real-valued, wrapped to exp(1j * value))."
             )
         if num_slices < 1:
             raise ValueError(f"num_slices must be greater than 0, got {num_slices}")
@@ -1606,13 +1607,21 @@ def from_uniform(
         first_omega_0: float = 10.0,
         hidden_omega_0: float = 10.0,
         obj_type: object_type = "pure_phase",
+        final_activation: str | Callable | None = None,
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
     ) -> "ObjectINR":
         """Create an ObjectINR backed by a default ``HSiren``, initialized to vacuum.
 
-        The HSiren's final layer is zero-initialized so the object starts as a flat,
-        phase-0 (vacuum) transmission, matching ``ObjectPixelated.from_uniform``.
+        The HSiren's final layer is zero-initialized so the object starts uniform (a
+        diffraction-equivalent vacuum), matching ``ObjectPixelated.from_uniform``.
+
+        ``final_activation`` sets the output nonlinearity. When ``None`` (default) it is chosen
+        from ``obj_type``: ``"identity"`` for ``pure_phase``, and ``"softplus"`` for
+        ``potential`` -- a non-negative activation enforces the potential's positivity (min-value)
+        constraint directly at the output (use ``"relu"`` for a hard floor). With the zeroed final
+        layer the potential starts uniform (``softplus(0) = ln 2``), which is just a global phase
+        and hence diffraction-equivalent to vacuum.
 
         Note
         ----
@@ -1622,6 +1631,8 @@ def from_uniform(
         is typically too high here (optimization stalls near vacuum), while objects with fine
         features may want a larger value. Pair omega_0 with the object learning rate.
         """
+        if final_activation is None:
+            final_activation = "softplus" if obj_type == "potential" else "identity"
         model = HSiren(
             in_features=3,
             out_features=1,
@@ -1629,9 +1640,10 @@ def from_uniform(
             hidden_features=hidden_features,
             first_omega_0=first_omega_0,
             hidden_omega_0=hidden_omega_0,
+            final_activation=final_activation,
             dtype=getattr(torch, config.get("dtype_real")),
         )
-        # Zero the final linear layer so the INR outputs phase 0 everywhere (vacuum start).
+        # Zero the final linear layer so the INR output is uniform at init (vacuum / global phase).
         with torch.no_grad():
             final_linear = cast(nn.Linear, model.net[-2])
             final_linear.weight.zero_()
@@ -1650,6 +1662,7 @@ def from_uniform(
     def from_pixelated(
         cls,
         pixelated: "ObjectModelType",
+        model: "torch.nn.Module | None" = None,
         hidden_features: int = 256,
         hidden_layers: int = 3,
         first_omega_0: float = 10.0,
@@ -1663,24 +1676,41 @@ def from_pixelated(
         ``slice_thicknesses``, ``obj_type``, padded shape) and the current pixelated object is
         stored as the pretrain target, so ``pretrain()`` warm-starts the INR to reproduce the
         pixelated reconstruction -- mirroring ``ObjectDIP.from_pixelated`` + ``pretrain``.
+
+        Pass ``model`` to wrap a custom INR ``nn.Module`` directly (mapping ``(N, 3)`` coords to
+        ``(N, 1)``), as with ``ObjectDIP.from_pixelated`` -- handy for testing architectures. When
+        ``model`` is ``None`` a default zero-initialized ``HSiren`` is built from the
+        ``hidden_features`` / ``hidden_layers`` / ``omega_0`` args (with a positivity activation
+        for ``potential``); when a ``model`` is given those args and the activation are its own.
         """
         if not (
             isinstance(pixelated, ObjectPixelated) or "ObjectPixelated" in str(type(pixelated))
         ):
             raise ValueError(f"pixelated must be an ObjectPixelated, got {type(pixelated)}")
         dev = pixelated.device if device is None else device
-        inr = cls.from_uniform(
-            num_slices=pixelated.num_slices,
-            slice_thicknesses=pixelated.slice_thicknesses,
-            obj_type=pixelated.obj_type,
-            hidden_features=hidden_features,
-            hidden_layers=hidden_layers,
-            first_omega_0=first_omega_0,
-            hidden_omega_0=hidden_omega_0,
-            device=dev,
-            rng=pixelated._rng_seed if rng is None else rng,
-        )
-        target = pixelated.obj.detach().to(dev)  # (num_slices, H, W) real phase
+        seed = pixelated._rng_seed if rng is None else rng
+        if model is not None:
+            inr = cls.from_inr(
+                model=model,
+                num_slices=pixelated.num_slices,
+                slice_thicknesses=pixelated.slice_thicknesses,
+                obj_type=pixelated.obj_type,
+                device=dev,
+                rng=seed,
+            )
+        else:
+            inr = cls.from_uniform(
+                num_slices=pixelated.num_slices,
+                slice_thicknesses=pixelated.slice_thicknesses,
+                obj_type=pixelated.obj_type,
+                hidden_features=hidden_features,
+                hidden_layers=hidden_layers,
+                first_omega_0=first_omega_0,
+                hidden_omega_0=hidden_omega_0,
+                device=dev,
+                rng=seed,
+            )
+        target = pixelated.obj.detach().to(dev)  # (num_slices, H, W) real phase / potential
         inr._obj_shape = tuple(int(x) for x in target.shape)  # type: ignore[assignment]
         if pixelated._sampling is not None:
             inr.sampling = pixelated.sampling
@@ -1998,9 +2028,10 @@ def apply_hard_constraints(
     ) -> torch.Tensor:
         """Project the materialized object (display only).
 
-        Unlike the grid-based ``Raster`` constraints, an INR has nothing to clamp or
-        filter in place; for ``pure_phase`` we only recenter the phase to zero mean so
-        the displayed object matches the pixelated convention.
+        Unlike the grid-based ``Raster`` constraints, an INR has nothing to clamp or filter in
+        place. For ``pure_phase`` we recenter the phase to zero mean (a global-phase gauge) so the
+        displayed object matches the pixelated convention; ``potential`` is returned as-is (the INR
+        constraint set has no positivity/baseline fields — potential is left unconstrained).
         """
         with torch.no_grad():
             if self.obj_type == "pure_phase":
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 490bb1321..033506caa 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -78,9 +78,6 @@ class PtychographyBase(RNGMixin, AutoSerialize):
     """
 
     _token = object()
-    # Default data-fidelity criterion (overridden per-instance from `loss_type` in reconstruct).
-    # Class-level so freshly-built and freshly-loaded objects always resolve a criterion.
-    _criterion: DataCriterion = get_data_criterion("l2_amplitude")
 
     def __init__(  # TODO prevent direct instantiation
         self,
@@ -110,6 +107,9 @@ def __init__(  # TODO prevent direct instantiation
         self._obj_model: ObjectModelType = obj_model
         self._probe_model: ProbeModelType = probe_model
         self._detector_model: DetectorModelType = detector_model
+        # Data-fidelity criterion (transient; re-set from `loss_type` in reconstruct). Not
+        # serialized (skipped on save), so the getter lazily re-defaults it on loaded objects.
+        self._criterion: DataCriterion = get_data_criterion("l2_amplitude")
 
         self.verbose = verbose
         self.dset = dset
@@ -253,9 +253,11 @@ def dset(self, new_dset: DatasetModelType):
     def criterion(self) -> DataCriterion:
         """Active data-fidelity criterion. Assign a registered name or a ``DataCriterion``.
 
-        Transient config (re-set from ``loss_type`` each ``reconstruct``, defaults to L2); not
-        serialized.
+        Transient config (re-set from ``loss_type`` each ``reconstruct``); not serialized, so it
+        lazily re-defaults to L2 on a loaded object (where ``__init__`` did not run).
         """
+        if getattr(self, "_criterion", None) is None:
+            self._criterion = get_data_criterion("l2_amplitude")
         return self._criterion
 
     @criterion.setter
diff --git a/tests/diffractive_imaging/test_object_inr.py b/tests/diffractive_imaging/test_object_inr.py
index c22c704dc..76f752a6d 100644
--- a/tests/diffractive_imaging/test_object_inr.py
+++ b/tests/diffractive_imaging/test_object_inr.py
@@ -262,6 +262,38 @@ def test_complex_obj_type_not_supported(self):
         with pytest.raises(NotImplementedError):
             ObjectINR.from_uniform(num_slices=1, obj_type="complex")
 
+    def test_potential_obj_type(self):
+        """`potential` is real-valued like `pure_phase` but uses a non-negative output activation
+        (softplus) to enforce the positivity / min-value constraint."""
+        obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
+        obj._initialize_obj((1, 16, 16))
+        assert obj.obj_type == "potential"
+        assert not obj.dtype.is_complex  # real-valued object
+        patches = obj.forward(torch.rand(3, 4, 4, 2) * 2 - 1)
+        assert patches.shape == (1, 3, 4, 4) and patches.is_complex()
+        # train a few steps so the potential is non-uniform, then check positivity holds
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
+        coords = torch.rand(2, 6, 6, 2) * 2 - 1
+        for _ in range(5):
+            opt.zero_grad()
+            obj.forward(coords).imag.sum().backward()
+            opt.step()
+        materialized = obj.obj
+        assert materialized.shape == (1, 16, 16) and not materialized.is_complex()
+        assert float(materialized.min()) >= 0.0  # softplus enforces non-negative potential
+
+    def test_from_pixelated_with_model(self):
+        """from_pixelated can wrap a directly-passed INR model (like ObjectDIP.from_pixelated)."""
+        from quantem.core.ml.inr import HSiren
+
+        h = w = 16
+        pix = ObjectPixelated.from_array(initial_obj=torch.zeros(1, h, w), obj_type="pure_phase")
+        pix._initialize_obj((1, h, w), sampling=(1.0, 1.0))
+        my_model = HSiren(in_features=3, out_features=1, hidden_features=16, hidden_layers=2)
+        inr = ObjectINR.from_pixelated(pix, model=my_model)
+        assert inr.model is my_model
+        assert tuple(inr.pretrain_target.shape) == (1, h, w)
+
     def test_from_pixelated_pretrain(self):
         """from_pixelated + pretrain warm-starts the INR to reproduce a pixelated object."""
         h = w = 32

From f8e300f41ca8f76b85c6df0eca1a387b15d5a275 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Mon, 1 Jun 2026 16:09:09 -0700
Subject: [PATCH 30/59] comment removal

---
 src/quantem/diffractive_imaging/object_models.py     | 2 +-
 src/quantem/diffractive_imaging/ptychography_base.py | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 20d1c65b3..3c63c285a 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1528,7 +1528,7 @@ class ObjectINR(BaseConstraints[PtychoObjConstraintParams.INR], ObjectBase):
     """
 
     DEFAULT_LRS = {
-        "object": 8e-6,
+        "object": 1e-3,
         "tv_weight_z": 0,
         "tv_weight_xy": 0,
     }
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 033506caa..93d034d51 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -107,8 +107,6 @@ def __init__(  # TODO prevent direct instantiation
         self._obj_model: ObjectModelType = obj_model
         self._probe_model: ProbeModelType = probe_model
         self._detector_model: DetectorModelType = detector_model
-        # Data-fidelity criterion (transient; re-set from `loss_type` in reconstruct). Not
-        # serialized (skipped on save), so the getter lazily re-defaults it on loaded objects.
         self._criterion: DataCriterion = get_data_criterion("l2_amplitude")
 
         self.verbose = verbose

From 586c1d3e9055e24fc6cc0fbd7e138298525a8756 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Mon, 1 Jun 2026 17:11:48 -0700
Subject: [PATCH 31/59] fixing linter error

---
 src/quantem/diffractive_imaging/object_models.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index bbaf7d596..1c0099fcc 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1735,9 +1735,9 @@ def model(self) -> "torch.nn.Module":
         return self._model
 
     @property
-    def params(self):
+    def params(self) -> list[nn.Parameter]:
         """optimization parameters"""
-        return self._model.parameters()
+        return list(self._model.parameters())
 
     @property
     def pretrained_weights(self) -> dict[str, torch.Tensor]:
@@ -1856,8 +1856,8 @@ def pretrain(
         self,
         pretrain_target: torch.Tensor | np.ndarray | None = None,
         num_iters: int = 200,
-        optimizer_params: "dict | OptimizerType | None" = None,
-        scheduler_params: "dict | SchedulerType | None" = None,
+        optimizer_params: "dict | OptimizerParamsType | None" = None,
+        scheduler_params: "dict | SchedulerParamsType | None" = None,
         loss_fn: Callable | str = "l2",
         device: str | int | None = None,
         show: bool = True,

From c26c6bfdf15e3f26541947f4d77b77531c39411c Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Tue, 2 Jun 2026 22:17:44 -0700
Subject: [PATCH 32/59] initial workingish kplanes

---
 src/quantem/core/ml/models/kplanes.py         |  25 +-
 src/quantem/diffractive_imaging/__init__.py   |   1 +
 .../diffractive_imaging/object_models.py      | 288 ++++++++-
 .../diffractive_imaging/ptychography_opt.py   |  22 +-
 .../test_object_tensor_decomp.py              | 590 ++++++++++++++++++
 5 files changed, 922 insertions(+), 4 deletions(-)
 create mode 100644 tests/diffractive_imaging/test_object_tensor_decomp.py

diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py
index cdf552617..5aa1f5e0e 100644
--- a/src/quantem/core/ml/models/kplanes.py
+++ b/src/quantem/core/ml/models/kplanes.py
@@ -169,6 +169,7 @@ class KPlanes(PPLR, TensorDecompositionModel):
     """
     K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241
     """
+
     def __init__(
         self,
         # Grid parameters
@@ -176,7 +177,7 @@ def __init__(
         input_coords_dims: int = 3,
         M_features: int = 32,
         resolution: Sequence[int] = (200, 200, 200),
-        multiscale_res_multipliers: Optional[Sequence[int]] = None,
+        multiscale_res_multipliers: Optional[Sequence[float]] = None,
         concat_features: bool = True,
         density_activation: Callable = lambda x: F.softplus(
             x - 1
@@ -208,7 +209,22 @@ def __init__(
             self.grids.append(plane)
             self.feature_dim += self.M_features
 
-        # Network head
+        # Network head (single linear when not hybrid; small ReLU MLP when hybrid)
+        self._build_sigma_net(use_hybrid_mlp, hybrid_hidden_dim, hybrid_num_layers)
+
+    def _build_sigma_net(
+        self,
+        use_hybrid_mlp: bool,
+        hybrid_hidden_dim: int,
+        hybrid_num_layers: int,
+    ) -> None:
+        """Build the decoder head mapping concatenated grid features -> density.
+
+        ``use_hybrid_mlp=True`` builds a small ReLU MLP; otherwise a single linear
+        "explicit" decoder. Both init the final layer small (``weight ~ N(0, 0.01**2)``,
+        ``bias=0``) so the density starts near zero. Called after ``self.feature_dim``
+        is finalized.
+        """
         if use_hybrid_mlp:
             hybrid_hidden_dim = int(hybrid_hidden_dim)
             hybrid_num_layers = int(hybrid_num_layers)
@@ -233,6 +249,11 @@ def __init__(
             nn.init.zeros_(out.bias)
             layers.append(out)
             self.sigma_net = nn.Sequential(*layers)
+        else:
+            # Single-linear "explicit" decoder. Small init -> density ~ 0 initially.
+            self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True)
+            nn.init.normal_(self.sigma_net.weight, std=0.01)
+            nn.init.zeros_(self.sigma_net.bias)
 
     def get_densities(self, coords: torch.Tensor):
         """Computes and returns densities"""
diff --git a/src/quantem/diffractive_imaging/__init__.py b/src/quantem/diffractive_imaging/__init__.py
index d9058a3fa..9c1520dd3 100644
--- a/src/quantem/diffractive_imaging/__init__.py
+++ b/src/quantem/diffractive_imaging/__init__.py
@@ -8,6 +8,7 @@
     ObjectDIP as ObjectDIP,
     ObjectINR as ObjectINR,
     ObjectPixelated as ObjectPixelated,
+    ObjectTensorDecomp as ObjectTensorDecomp,
     PtychoObjConstraintParams as PtychoObjConstraintParams,
     PtychoObjConstraintsType as PtychoObjConstraintsType,
 )
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 1c0099fcc..1c5336702 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -17,8 +17,10 @@
 from quantem.core.ml.constraints import BaseConstraints, Constraints, parse_constraint_dict
 from quantem.core.ml.inr import HSiren
 from quantem.core.ml.loss_functions import get_loss_module
+from quantem.core.ml.models.kplanes import CPTilted, KPlanes, KPlanesTILTED, KPlanesType
 from quantem.core.ml.optimizer_mixin import (
     OptimizerMixin,
+    OptimizerParams,
     OptimizerParamsType,
     SchedulerParamsType,
 )
@@ -2095,4 +2097,288 @@ def backward(self, *args, **kwargs):
         )
 
 
-ObjectModelType = ObjectPixelated | ObjectDIP | ObjectINR
+class ObjectTensorDecomp(ObjectINR):
+    """Implicit object model backed by a tensor-decomposition network (K-Planes family).
+
+    A thin subclass of :class:`ObjectINR` that swaps the SIREN for a tensor-decomposition model
+    from :mod:`quantem.core.ml.models.kplanes` — plain :class:`KPlanes`, the tilted
+    :class:`KPlanesTILTED` (T learned SO(3) rotations), or the :class:`CPTilted` bottleneck.
+    Because these consume the same ``(N, 3)`` ``(z, y, x)`` coordinates and return ``(N, 1)``,
+    every coordinate-query path of ``ObjectINR`` (``forward``, ``_query_phase``, materialization,
+    sampled-TV soft constraints, ``pretrain``) is reused unchanged; only the optimizer wiring
+    differs.
+
+    These models expose multiple parameter groups (``grids``/``sigma_net``, plus ``so3`` for the
+    tilted variants), so the model uses per-parameter-group learning rates (PPLR):
+    ``optimizer_params`` must be a dict keyed by ``model.param_keys`` (see
+    :meth:`get_optimization_parameters` / :meth:`_normalize_optimizer_params`), e.g.
+    ``{"grids": OptimizerParams.Adam(lr=1e-2), "sigma_net": OptimizerParams.Adam(lr=1e-3)}``.
+
+    The object grid is treated as a 3D ``(z, y, x)`` volume with ``z`` the slice axis: a single
+    slice is queried at ``z = 0`` (the in-plane K-plane already provides the 2D feature grid),
+    multislice spans ``z`` via the slice thicknesses. A model's ``resolution`` is the
+    *feature-plane* resolution, set by the user and decoupled from the padded object grid (which
+    comes from ``_initialize_obj``).
+
+    Notes
+    -----
+    - The ``density_activation`` must be a picklable ``nn.Module`` (``nn.Identity`` for
+      ``pure_phase``, ``nn.Softplus`` for ``potential``); a bare lambda will break ``save()``
+      because AutoSerialize pickles the whole module. ``from_uniform`` sets this automatically;
+      ``from_model`` warns otherwise.
+    - K-Planes converge from scratch, but ``from_pixelated`` + ``pretrain`` can still warm-start
+      the grid to a pixelated reconstruction (cheap, useful for finding hyperparameters quickly).
+    """
+
+    DEFAULT_CONSTRAINTS: PtychoObjConstraintParams.INR = PtychoObjConstraintParams.INR()
+
+    @classmethod
+    def from_model(
+        cls,
+        model: KPlanesType,
+        num_slices: int = 1,
+        slice_thicknesses: float | Sequence | torch.Tensor | None = None,
+        obj_type: object_type = "pure_phase",
+        device: str = "cpu",
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectTensorDecomp":
+        """Wrap a user-built tensor-decomposition model as a ptychography object model.
+
+        ``model`` is a :class:`KPlanes`, :class:`KPlanesTILTED`, or :class:`CPTilted` mapping
+        ``(N, 3)`` ``(z, y, x)`` coordinates to ``(N, 1)`` and exposing the PPLR interface
+        (``param_keys`` / ``get_params``).
+        """
+        if not isinstance(model, (KPlanes, CPTilted)):  # KPlanesTILTED is a KPlanes subclass
+            raise TypeError(
+                f"model must be a KPlanes/KPlanesTILTED/CPTilted instance, got {type(model)}"
+            )
+        activation = getattr(model, "density_activation", None)
+        if activation is not None and not isinstance(activation, nn.Module):
+            warn(
+                "KPlanes.density_activation is a plain callable (e.g. a lambda); saving this "
+                "object will fail because AutoSerialize pickles the whole module. Use an "
+                "nn.Module activation (nn.Identity for pure_phase, nn.Softplus for potential), "
+                "e.g. via ObjectTensorDecomp.from_uniform.",
+                stacklevel=2,
+            )
+        obj = cls(
+            model=model,
+            num_slices=num_slices,
+            slice_thicknesses=slice_thicknesses,
+            obj_type=obj_type,
+            device=device,
+            rng=rng,
+            _token=cls._token,
+        )
+        obj.to(device)
+        return obj
+
+    @classmethod
+    def from_uniform(  # pyright: ignore[reportIncompatibleMethodOverride]  # KPlanes factory, intentionally diverges from ObjectINR.from_uniform
+        cls,
+        num_slices: int = 1,
+        slice_thicknesses: float | Sequence | torch.Tensor | None = None,
+        M_features: int = 16,
+        resolution: Sequence[int] = (64, 64, 64),
+        multiscale_res_multipliers: Sequence[float] | None = (0.25, 0.5, 1.0),
+        use_hybrid_mlp: bool = False,
+        hybrid_hidden_dim: int = 64,
+        hybrid_num_layers: int = 2,
+        tilted: bool = False,
+        T: int = 4,
+        obj_type: object_type = "pure_phase",
+        device: str = "cpu",
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectTensorDecomp":
+        """Build a default K-Planes-backed object, initialized to vacuum.
+
+        ``tilted=False`` builds a plain :class:`KPlanes`; ``tilted=True`` builds a
+        :class:`KPlanesTILTED` with ``T`` learned SO(3) rotations. The decoder's final layer is
+        zeroed so the object starts uniform (a global phase, diffraction-equivalent to vacuum),
+        matching ``ObjectINR.from_uniform``. ``density_activation`` is chosen from ``obj_type``:
+        ``nn.Identity`` for ``pure_phase`` (phase may be negative) and ``nn.Softplus`` for
+        ``potential`` (non-negative). ``resolution`` is the feature-plane resolution ``(z, y, x)``
+        and is independent of the reconstructed object grid; for multislice set ``resolution[0]``
+        to span the slices.
+        """
+        density_activation: nn.Module = nn.Softplus() if obj_type == "potential" else nn.Identity()
+        ms = list(multiscale_res_multipliers) if multiscale_res_multipliers is not None else None
+        model: KPlanesType
+        if tilted:
+            model = KPlanesTILTED(
+                M_features=M_features,
+                resolution=resolution,
+                multiscale_res_multipliers=ms,
+                density_activation=density_activation,
+                T=T,
+                use_hybrid_mlp=use_hybrid_mlp,
+                hybrid_hidden_dim=hybrid_hidden_dim,
+                hybrid_num_layers=hybrid_num_layers,
+            )
+        else:
+            model = KPlanes(
+                M_features=M_features,
+                resolution=resolution,
+                multiscale_res_multipliers=ms,
+                density_activation=density_activation,
+                use_hybrid_mlp=use_hybrid_mlp,
+                hybrid_hidden_dim=hybrid_hidden_dim,
+                hybrid_num_layers=hybrid_num_layers,
+            )
+        # Zero the final decoder layer so the object starts at vacuum (global phase).
+        with torch.no_grad():
+            final_linear = (
+                model.sigma_net[-1]
+                if isinstance(model.sigma_net, nn.Sequential)
+                else model.sigma_net
+            )
+            final_linear = cast(nn.Linear, final_linear)
+            final_linear.weight.zero_()
+            if final_linear.bias is not None:
+                final_linear.bias.zero_()
+        return cls.from_model(
+            model,
+            num_slices=num_slices,
+            slice_thicknesses=slice_thicknesses,
+            obj_type=obj_type,
+            device=device,
+            rng=rng,
+        )
+
+    @property
+    def name(self) -> str:
+        return "ObjectTensorDecomp"
+
+    @property
+    def model(self) -> KPlanesType:
+        return cast(KPlanesType, self._model)
+
+    def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]":
+        """PPLR: one param group per ``model.param_keys`` (hyperparameters baked by set_optimizer)."""
+        model = self.model
+        groups = model.get_params()
+        return {key: list(groups[key]) for key in model.param_keys}
+
+    def _normalize_optimizer_params(self, params):
+        """Require a dict keyed by ``model.param_keys`` (PPLR); reject single-optimizer specs.
+
+        The framework's "disabled" sentinel — a bare ``NoneOptimizer`` or a dict whose values
+        are all ``NoneOptimizer`` (e.g. the ``{"default": NoneOptimizer()}`` set at init / by
+        ``remove_optimizer`` and replayed through ``reset_optimizer`` on ``reconstruct(reset=True)``)
+        — is passed straight to the base normalizer so the optimizer can be cleanly disabled
+        without matching ``param_keys``.
+        """
+        if isinstance(params, OptimizerParams.NoneOptimizer) or (
+            isinstance(params, dict)
+            and len(params) > 0
+            and all(isinstance(v, OptimizerParams.NoneOptimizer) for v in params.values())
+        ):
+            return super()._normalize_optimizer_params(params)
+        if not isinstance(params, dict) or self._is_single_optimizer_dict(params):
+            raise TypeError(
+                f"{type(self).__name__} requires dict[str, OptimizerParamsType] keyed by "
+                f"param_keys {self.model.param_keys}; got {type(params)}"
+            )
+        expected = set(self.model.param_keys)
+        got = set(params.keys())
+        if got != expected:
+            raise ValueError(
+                f"optimizer_params keys must match model.param_keys: got {got}, expected {expected}"
+            )
+        return super()._normalize_optimizer_params(params)
+
+    @classmethod
+    def from_pixelated(  # pyright: ignore[reportIncompatibleMethodOverride]  # K-Planes factory, intentionally diverges from ObjectINR.from_pixelated
+        cls,
+        pixelated: "ObjectModelType",
+        model: KPlanesType | None = None,
+        M_features: int = 16,
+        resolution: Sequence[int] = (128, 128, 128),
+        multiscale_res_multipliers: Sequence[float] | None = (0.25, 0.5, 1.0),
+        use_hybrid_mlp: bool = False,
+        tilted: bool = False,
+        T: int = 4,
+        device: str | None = None,
+        rng: np.random.Generator | int | None = None,
+    ) -> "ObjectTensorDecomp":
+        """Build a K-Planes object matching a pixelated object, with it as the pretrain target.
+
+        Mirrors :meth:`ObjectINR.from_pixelated`: the K-Planes model is built to the pixelated
+        object's geometry (``num_slices``, ``slice_thicknesses``, ``obj_type``, padded shape) and
+        the current pixelated object is stored as the pretrain target, so ``pretrain()``
+        warm-starts the grid to reproduce the pixelated reconstruction. Pass ``model`` to wrap a
+        custom tensor-decomposition ``nn.Module`` directly; otherwise one is built from the
+        ``M_features`` / ``resolution`` / ``tilted`` / ``T`` args.
+        """
+        if not (
+            isinstance(pixelated, ObjectPixelated) or "ObjectPixelated" in str(type(pixelated))
+        ):
+            raise ValueError(f"pixelated must be an ObjectPixelated, got {type(pixelated)}")
+        dev = pixelated.device if device is None else device
+        seed = pixelated._rng_seed if rng is None else rng
+        if model is not None:
+            obj = cls.from_model(
+                model=model,
+                num_slices=pixelated.num_slices,
+                slice_thicknesses=pixelated.slice_thicknesses,
+                obj_type=pixelated.obj_type,
+                device=dev,
+                rng=seed,
+            )
+        else:
+            obj = cls.from_uniform(
+                num_slices=pixelated.num_slices,
+                slice_thicknesses=pixelated.slice_thicknesses,
+                M_features=M_features,
+                resolution=resolution,
+                multiscale_res_multipliers=multiscale_res_multipliers,
+                use_hybrid_mlp=use_hybrid_mlp,
+                tilted=tilted,
+                T=T,
+                obj_type=pixelated.obj_type,
+                device=dev,
+                rng=seed,
+            )
+        target = pixelated.obj.detach().to(dev)  # (num_slices, H, W) real phase / potential
+        obj._obj_shape = tuple(int(x) for x in target.shape)  # type: ignore[assignment]
+        if pixelated._sampling is not None:
+            obj.sampling = pixelated.sampling
+        obj.pretrain_target = target
+        return obj
+
+    def pretrain(
+        self,
+        pretrain_target: torch.Tensor | np.ndarray | None = None,
+        num_iters: int = 200,
+        optimizer_params: "dict | OptimizerParamsType | None" = None,
+        scheduler_params: "dict | SchedulerParamsType | None" = None,
+        loss_fn: Callable | str = "l2",
+        device: str | int | None = None,
+        show: bool = True,
+        normalize_object_plotting: bool = True,
+    ) -> None:
+        """Warm-start the K-Planes grid by regressing it onto a target object (PPLR).
+
+        Same direct grid->target regression as ``ObjectINR.pretrain`` (no forward model), but
+        ``optimizer_params`` is PPLR-keyed. When ``None`` it defaults to per-group Adam
+        (``grids`` lr 1e-2, others 1e-3) so pretraining works out of the box; the fitted weights
+        become the ``reset()`` state.
+        """
+        if optimizer_params is None:
+            optimizer_params = {
+                key: OptimizerParams.Adam(lr=1e-2 if key == "grids" else 1e-3)
+                for key in self.model.param_keys
+            }
+        super().pretrain(
+            pretrain_target=pretrain_target,
+            num_iters=num_iters,
+            optimizer_params=optimizer_params,
+            scheduler_params=scheduler_params,
+            loss_fn=loss_fn,
+            device=device,
+            show=show,
+            normalize_object_plotting=normalize_object_plotting,
+        )
+
+
+ObjectModelType = ObjectPixelated | ObjectDIP | ObjectINR | ObjectTensorDecomp
diff --git a/src/quantem/diffractive_imaging/ptychography_opt.py b/src/quantem/diffractive_imaging/ptychography_opt.py
index a22456f67..9a08ccc37 100644
--- a/src/quantem/diffractive_imaging/ptychography_opt.py
+++ b/src/quantem/diffractive_imaging/ptychography_opt.py
@@ -52,6 +52,22 @@ def _check_key(self, key: str) -> None:
                 f"key to be optimized, {key}, not in allowed keys: {self.OPTIMIZABLE_VALS}"
             )
 
+    @staticmethod
+    def _is_pplr_dict(v: dict) -> bool:
+        """Detect a nested per-parameter-group (PPLR) optimizer spec.
+
+        A PPLR dict maps parameter-group keys (e.g. ``"grids"``/``"sigma_net"`` for
+        ``ObjectTensorDecomp``) to per-group ``OptimizerParamsType`` or dict specs — as opposed
+        to a single-optimizer shorthand like ``{"name": "adam", "lr": 1e-3}``. Such dicts are
+        passed through to the owning model untouched so its ``_normalize_optimizer_params`` can
+        validate the keys against ``param_keys``.
+        """
+        return (
+            len(v) > 0
+            and not OptimizerMixin._is_single_optimizer_dict(v)
+            and all(isinstance(val, (OptimizerParamsType, dict)) for val in v.values())
+        )
+
     def _get_default_lr(self, key: str) -> float:
         """Get default learning rate for a given optimization key."""
         if key == "object":
@@ -92,6 +108,8 @@ def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
         for k, v in _d.items():
             if isinstance(v, OptimizerParamsType):
                 pass  # already a dataclass, pass through
+            elif isinstance(v, dict) and self._is_pplr_dict(v):
+                pass  # nested per-parameter-group (PPLR) spec -> pass through untouched
             elif isinstance(v, dict):
                 if not v:
                     v = replace(self.DEFAULT_OPTIMIZER_TYPE, lr=self._get_default_lr(k))
@@ -101,7 +119,9 @@ def optimizer_params(self, d: dict[str, Any] | list[str] | tuple[str, ...]) -> N
                     if "lr" not in v:
                         v["lr"] = self._get_default_lr(k)
             else:
-                raise TypeError(f"Expected OptimizerParamsType or dict for key '{k}', got {type(v)}")
+                raise TypeError(
+                    f"Expected OptimizerParamsType or dict for key '{k}', got {type(v)}"
+                )
 
             self._models[k].optimizer_params = v  # type: ignore[assignment]
 
diff --git a/tests/diffractive_imaging/test_object_tensor_decomp.py b/tests/diffractive_imaging/test_object_tensor_decomp.py
new file mode 100644
index 000000000..dabf2b6b8
--- /dev/null
+++ b/tests/diffractive_imaging/test_object_tensor_decomp.py
@@ -0,0 +1,590 @@
+"""
+Tests for the tensor-decomposition (K-Planes) object model, ``ObjectTensorDecomp``.
+
+``ObjectTensorDecomp`` subclasses ``ObjectINR`` and swaps the SIREN for a ``KPlanes``
+network, so it reuses every coordinate-query path of the INR and only differs in the
+optimizer wiring (per-parameter-group learning rates, PPLR). These tests cover the
+shared-core ``KPlanes`` head fix, the object model in isolation (forward/vacuum/PPLR
+parameter groups, bad-key rejection), the ptychography PPLR optimizer pass-through, and
+an end-to-end reconstruction on the same non-toroidal synthetic object used by
+``test_object_inr.py`` (fixtures duplicated here, matching that file's own duplication
+from ``test_ptychography.py`` — ``--import-mode=importlib`` makes cross-test imports
+unreliable).
+"""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+from quantem.core import config
+from quantem.core.datastructures.dataset4dstem import Dataset4dstem
+from quantem.core.io.serialize import load as autoserialize_load
+from quantem.core.ml import OptimizerParams, SchedulerParams
+from quantem.core.ml.models.kplanes import CPTilted, KPlanes
+from quantem.core.utils.utils import electron_wavelength_angstrom
+from quantem.diffractive_imaging.dataset_models import PtychographyDatasetRaster
+from quantem.diffractive_imaging.detector_models import DetectorPixelated
+from quantem.diffractive_imaging.object_models import (
+    ObjectModelType,
+    ObjectPixelated,
+    ObjectTensorDecomp,
+)
+from quantem.diffractive_imaging.probe_models import ProbePixelated
+from quantem.diffractive_imaging.ptychography import Ptychography
+
+if config.NUM_DEVICES > 0:
+    config.set_device("gpu")
+
+N = 40  # detector / roi size (px)
+OGT = 64  # ground-truth object size (px); larger than the scanned region
+PAD = 20  # obj padding (>= roi // 2 so interior patches never hit the boundary)
+Q_MAX = 0.5  # inverse Angstroms
+Q_PROBE = Q_MAX / 2
+PROBE_ENERGY = 300e3  # eV
+C10 = 50.0  # defocus (Angstrom)
+STEP = 2  # scan step (px)
+SCAN_START = 20  # first scan position (px); SCAN_START - roi//2 >= 0
+SCAN_STOP = 44  # exclusive; SCAN_STOP - 1 + roi//2 - 1 < OGT  -> no wrap
+
+# K-Planes recon config (validated against the fixture below). grids learn faster than the
+# decoder head; both run under a single cosine-annealing schedule.
+_LR_GRIDS = 1e-2
+_LR_SIGMA = 1e-3
+_LOSS_RATIO = 0.3
+_CORR_THRESHOLD = 0.6
+
+
+def _smooth_phase() -> np.ndarray:
+    """A smooth, band-limited phase object."""
+    yy, xx = np.meshgrid(np.arange(OGT), np.arange(OGT), indexing="ij")
+    return (0.7 * np.sin(2 * np.pi * xx / OGT * 4) * np.cos(2 * np.pi * yy / OGT * 3)).astype(
+        np.float32
+    )
+
+
+def _probe_array() -> np.ndarray:
+    sampling = 1 / Q_MAX / 2
+    reciprocal_sampling = 2 * Q_MAX / N
+    qx = qy = np.fft.fftfreq(N, sampling)
+    q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2)
+    aperture = np.sqrt(np.clip((Q_PROBE - q) / reciprocal_sampling + 0.5, 0, 1))
+    chi = q**2 * electron_wavelength_angstrom(PROBE_ENERGY) * np.pi * C10
+    probe_fourier = aperture * np.exp(-1j * chi)
+    probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2))
+    return (np.fft.ifft2(probe_fourier) * N).astype(np.complex64)
+
+
+def _semiangle_mrad() -> float:
+    return electron_wavelength_angstrom(PROBE_ENERGY) * Q_PROBE * 1e3
+
+
+def _build_synthetic_dataset() -> tuple[PtychographyDatasetRaster, np.ndarray, np.ndarray]:
+    """Simulate a non-toroidal 4D-STEM dataset; return (dataset_model, gt_phase, probe)."""
+    phase = _smooth_phase()
+    complex_obj = np.exp(1j * phase)
+    probe = _probe_array()
+    reciprocal_sampling = 2 * Q_MAX / N
+
+    gpos = np.arange(SCAN_START, SCAN_STOP, STEP)
+    xx, yy = np.meshgrid(gpos, gpos, indexing="ij")
+    positions = np.stack((xx.ravel(), yy.ravel()), axis=-1)
+    x0 = positions[:, 0].astype(int)
+    y0 = positions[:, 1].astype(int)
+    x_ind = np.fft.fftfreq(N, d=1 / N).astype(int)
+    y_ind = np.fft.fftfreq(N, d=1 / N).astype(int)
+    row = x0[:, None, None] + x_ind[None, :, None]  # no modulo: interior scan, no wrap
+    col = y0[:, None, None] + y_ind[None, None, :]
+    assert row.min() >= 0 and row.max() < OGT and col.min() >= 0 and col.max() < OGT
+    exit_waves = complex_obj[row, col] * probe
+    intensities = np.abs(np.fft.fft2(exit_waves)) ** 2
+
+    sxy = len(gpos)
+    dset = Dataset4dstem.from_array(
+        array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((sxy, sxy, N, N)),
+        sampling=(STEP, STEP, reciprocal_sampling, reciprocal_sampling),
+        units=("A", "A", "A^-1", "A^-1"),
+    )
+    pdset = PtychographyDatasetRaster.from_dataset4dstem(dset)
+    pdset.learn_scan_positions = False
+    pdset.learn_descan = False
+    pdset.preprocess(
+        com_fit_function="constant",
+        plot_rotation=False,
+        plot_com=False,
+        probe_energy=PROBE_ENERGY,
+    )
+    return pdset, phase, probe
+
+
+def _build_kplanes_ptycho(
+    num_slices: int = 1,
+    slice_thicknesses=None,
+    M_features: int = 24,
+    resolution: tuple[int, int, int] = (16, 64, 64),
+) -> tuple[Ptychography, np.ndarray]:
+    """Build an ObjectTensorDecomp (K-Planes) ptychography with the exact (frozen) probe."""
+    pdset, gt_phase, probe = _build_synthetic_dataset()
+    obj = ObjectTensorDecomp.from_uniform(
+        num_slices=num_slices,
+        slice_thicknesses=slice_thicknesses,
+        M_features=M_features,
+        resolution=resolution,
+        multiscale_res_multipliers=(0.25, 0.5, 1.0),
+        use_hybrid_mlp=False,
+        obj_type="pure_phase",
+        rng=0,
+    )
+    probe_model = ProbePixelated.from_array(
+        num_probes=1,
+        probe_params={"energy": PROBE_ENERGY, "C10": C10, "semiangle_cutoff": _semiangle_mrad()},
+        probe_array=probe.copy(),
+    )
+    ptycho = Ptychography.from_models(
+        dset=pdset,
+        obj_model=obj,
+        probe_model=probe_model,
+        detector_model=DetectorPixelated(),
+        rng=0,
+        verbose=False,
+    )
+    ptycho.preprocess(obj_padding_px=(PAD, PAD))
+    return ptycho, gt_phase
+
+
+def _corr(a: np.ndarray, b: np.ndarray) -> float:
+    a = a - a.mean()
+    b = b - b.mean()
+    return float((a * b).sum() / np.sqrt((a**2).sum() * (b**2).sum() + 1e-12))
+
+
+def _center_crop(a: np.ndarray, s: int) -> np.ndarray:
+    r0 = (a.shape[0] - s) // 2
+    c0 = (a.shape[1] - s) // 2
+    return a[r0 : r0 + s, c0 : c0 + s]
+
+
+def _best_corr(recon: np.ndarray, gt: np.ndarray, s: int = 32) -> float:
+    """Best |correlation| over small integer shifts (handles ~1px registration offset)."""
+    g = _center_crop(gt, s)
+    best = -1.0
+    for dr in range(-3, 4):
+        for dc in range(-3, 4):
+            r = _center_crop(np.roll(recon, (dr, dc), (0, 1)), s)
+            best = max(best, abs(_corr(r, g)))
+    return best
+
+
+def _pplr_params() -> dict:
+    return {
+        "object": {
+            "grids": OptimizerParams.Adam(lr=_LR_GRIDS),
+            "sigma_net": OptimizerParams.Adam(lr=_LR_SIGMA),
+        }
+    }
+
+
+# --------------------------------------------------------------------------- #
+# Shared-core KPlanes head fix (use_hybrid_mlp=False must build a decoder head)
+# --------------------------------------------------------------------------- #
+class TestKPlanesHead:
+    def test_non_hybrid_builds_head_and_runs(self):
+        m = KPlanes(
+            M_features=8,
+            resolution=(16, 16, 16),
+            multiscale_res_multipliers=[0.5, 1.0],
+            density_activation=nn.Identity(),
+            use_hybrid_mlp=False,
+        )
+        assert isinstance(m.sigma_net, nn.Linear)
+        out = m(torch.rand(20, 3) * 2 - 1)
+        assert out.shape == (20, 1)
+        assert set(m.get_params().keys()) == {"grids", "sigma_net"}
+
+    def test_hybrid_still_builds_mlp_head(self):
+        m = KPlanes(
+            M_features=8,
+            resolution=(16, 16, 16),
+            multiscale_res_multipliers=[1.0],
+            density_activation=nn.Identity(),
+            use_hybrid_mlp=True,
+            hybrid_hidden_dim=32,
+            hybrid_num_layers=2,
+        )
+        assert isinstance(m.sigma_net, nn.Sequential)
+        assert m(torch.rand(7, 3) * 2 - 1).shape == (7, 1)
+
+
+# --------------------------------------------------------------------------- #
+# ObjectTensorDecomp in isolation
+# --------------------------------------------------------------------------- #
+class TestObjectTensorDecompUnit:
+    def _obj(self, **kw):
+        return ObjectTensorDecomp.from_uniform(
+            num_slices=kw.pop("num_slices", 1),
+            slice_thicknesses=kw.pop("slice_thicknesses", None),
+            M_features=8,
+            resolution=kw.pop("resolution", (16, 32, 32)),
+            multiscale_res_multipliers=(0.5, 1.0),
+            rng=0,
+            **kw,
+        )
+
+    def test_is_implicit_and_name(self):
+        obj = self._obj()
+        assert obj.is_implicit is True
+        assert obj.name == "ObjectTensorDecomp"
+        assert isinstance(obj, ObjectModelType)
+
+    def test_forward_shape_dtype_and_vacuum(self):
+        obj = self._obj()
+        obj._initialize_obj((1, 24, 28))
+        coords = torch.rand(5, 8, 8, 2) * 2 - 1
+        patches = obj.forward(coords)
+        assert patches.shape == (1, 5, 8, 8)
+        assert patches.is_complex()
+        # vacuum init (zeroed decoder head) -> unit transmission everywhere
+        assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6)
+
+    def test_off_object_is_vacuum(self):
+        obj = self._obj(resolution=(16, 16, 16))
+        obj._initialize_obj((1, 16, 16))
+        # train the head a little so the model is not identically zero
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
+        inside = torch.rand(3, 4, 4, 2) * 2 - 1
+        for _ in range(5):
+            opt.zero_grad()
+            obj.forward(inside).imag.sum().backward()
+            opt.step()
+        off = torch.full((1, 2, 2, 2), 5.0)
+        out = obj.forward(off)
+        assert torch.allclose(out, torch.ones_like(out), atol=1e-6)
+
+    def test_multislice_z_coordinates(self):
+        obj = self._obj(num_slices=3, slice_thicknesses=2.0, resolution=(16, 16, 16))
+        obj._initialize_obj((3, 16, 16))
+        z = obj._z_coords
+        assert z.shape == (3,)
+        assert torch.allclose(z, torch.tensor([-1.0, 0.0, 1.0]).to(z), atol=1e-6)
+        patches = obj.forward(torch.rand(4, 6, 6, 2) * 2 - 1)
+        assert patches.shape == (3, 4, 6, 6)
+
+    def test_get_optimization_parameters_keys(self):
+        obj = self._obj()
+        groups = obj.get_optimization_parameters()
+        assert set(groups.keys()) == {"grids", "sigma_net"}
+        for tensors in groups.values():
+            assert len(tensors) > 0
+            assert all(isinstance(t, nn.Parameter) and t.is_leaf for t in tensors)
+
+    def test_pplr_optimizer_construction(self):
+        obj = self._obj()
+        obj.set_optimizer(
+            {
+                "grids": OptimizerParams.Adam(lr=_LR_GRIDS),
+                "sigma_net": OptimizerParams.Adam(lr=_LR_SIGMA),
+            }
+        )
+        opt = obj.optimizer
+        assert isinstance(opt, torch.optim.Adam)
+        assert len(opt.param_groups) == 2
+        assert sorted(pg["lr"] for pg in opt.param_groups) == [_LR_SIGMA, _LR_GRIDS]
+
+    def test_optimizer_params_rejects_bad_keys(self):
+        obj = self._obj()
+        # missing a required group
+        with pytest.raises(ValueError):
+            obj.set_optimizer({"grids": OptimizerParams.Adam(lr=_LR_GRIDS)})
+        # single-optimizer shorthand is not a valid PPLR spec here
+        with pytest.raises(TypeError):
+            obj.set_optimizer({"name": "adam", "lr": _LR_GRIDS})
+        # a bare OptimizerParamsType (single optimizer) is also rejected
+        with pytest.raises(TypeError):
+            obj.set_optimizer(OptimizerParams.Adam(lr=_LR_GRIDS))
+
+    def test_pretrain_without_target_raises(self):
+        """pretrain is supported now, but needs a target (use from_pixelated or pass one)."""
+        obj = self._obj()
+        obj._initialize_obj((1, 24, 28))
+        with pytest.raises(ValueError, match="pretrain target"):
+            obj.pretrain(num_iters=1, show=False)
+
+    def test_disable_sentinel_accepted(self):
+        """The framework's "disabled" optimizer sentinel must pass through (used by reset).
+
+        ``reconstruct(reset=True)`` replays ``reset_optimizer`` with the init default
+        ``{"default": NoneOptimizer()}``; that must not be rejected by the PPLR key check.
+        """
+        obj = self._obj()
+        # bare NoneOptimizer and the default-keyed dict both mean "no optimizer"
+        obj.set_optimizer(OptimizerParams.NoneOptimizer())
+        assert obj.optimizer is None
+        obj.set_optimizer({"default": OptimizerParams.NoneOptimizer()})
+        assert obj.optimizer is None
+
+    def test_autograd_only_backward_raises(self):
+        obj = self._obj()
+        with pytest.raises(NotImplementedError):
+            obj.backward()
+
+    def test_from_model_warns_on_lambda_activation(self):
+        # a lambda activation cannot be pickled by AutoSerialize -> from_model must warn
+        model = KPlanes(
+            M_features=4,
+            resolution=(8, 8, 8),
+            multiscale_res_multipliers=[1.0],
+            use_hybrid_mlp=False,  # default density_activation is a lambda
+        )
+        with pytest.warns(UserWarning, match="density_activation"):
+            ObjectTensorDecomp.from_model(model, num_slices=1, obj_type="pure_phase")
+
+    def test_from_model_rejects_non_kplanes(self):
+        with pytest.raises(TypeError):
+            ObjectTensorDecomp.from_model(nn.Linear(3, 1), num_slices=1)
+
+    def test_tilted_from_uniform(self):
+        """tilted=True builds a KPlanesTILTED with an extra `so3` param group."""
+        obj = ObjectTensorDecomp.from_uniform(
+            num_slices=1,
+            M_features=6,
+            resolution=(16, 32, 32),
+            multiscale_res_multipliers=(0.5, 1.0),
+            tilted=True,
+            T=4,
+            obj_type="pure_phase",
+            rng=0,
+        )
+        obj._initialize_obj((1, 24, 28))
+        assert obj.model.tilted is True
+        assert set(obj.get_optimization_parameters().keys()) == {"grids", "sigma_net", "so3"}
+        patches = obj.forward(torch.rand(3, 6, 6, 2) * 2 - 1)
+        assert patches.shape == (1, 3, 6, 6)
+        # vacuum init holds for the tilted decoder too
+        assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6)
+        obj.set_optimizer({k: OptimizerParams.Adam(lr=1e-2) for k in obj.model.param_keys})
+        assert len(obj.optimizer.param_groups) == 3
+
+    def test_from_model_accepts_cptilted(self):
+        """from_model accepts the CPTilted bottleneck (TensorDecompositionModel, not a KPlanes)."""
+        cp = CPTilted(
+            C=4,
+            resolution=(32, 32, 32),
+            multiscale_res_multipliers=[1.0],
+            T=4,
+            density_activation=nn.Identity(),
+        )
+        obj = ObjectTensorDecomp.from_model(cp, num_slices=1, obj_type="pure_phase")
+        assert set(obj.get_optimization_parameters().keys()) == {"grids", "sigma_net", "so3"}
+        assert obj.forward(torch.rand(2, 4, 4, 2) * 2 - 1).shape == (1, 2, 4, 4)
+
+    def test_from_pixelated_pretrain(self):
+        """from_pixelated + pretrain warm-starts the K-Planes grid to a pixelated object."""
+        h = w = 32
+        ys = torch.linspace(-1, 1, h)
+        xs = torch.linspace(-1, 1, w)
+        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
+        phase = (0.5 * torch.sin(3 * gy) * torch.cos(2 * gx)).float()[None]
+        pix = ObjectPixelated.from_array(initial_obj=phase, obj_type="pure_phase")
+        pix._initialize_obj((1, h, w), sampling=(1.0, 1.0))
+
+        kp = ObjectTensorDecomp.from_pixelated(pix, M_features=16, resolution=(48, 48, 48))
+        assert kp.num_slices == pix.num_slices
+        assert tuple(kp.pretrain_target.shape) == tuple(pix.obj.shape)
+
+        kp.pretrain(num_iters=120, show=False)  # default PPLR optimizer
+        losses = kp.pretrain_losses
+        assert losses[-1] < 0.05 * losses[0]
+
+        gt = pix.obj[0].detach().cpu().numpy()
+        gt -= gt.mean()
+
+        def _corr_to_pix(arr):
+            a = arr - arr.mean()
+            return float((a * gt).sum() / np.sqrt((a**2).sum() * (gt**2).sum() + 1e-12))
+
+        assert _corr_to_pix(kp.obj[0].detach().cpu().numpy()) > 0.9
+        # pretrained weights are the reset state
+        kp.reset()
+        assert _corr_to_pix(kp.obj[0].detach().cpu().numpy()) > 0.9
+
+    def test_potential_obj_type_positive(self):
+        obj = self._obj(obj_type="potential", resolution=(16, 16, 16))
+        obj._initialize_obj((1, 16, 16))
+        assert obj.obj_type == "potential"
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
+        coords = torch.rand(2, 6, 6, 2) * 2 - 1
+        for _ in range(5):
+            opt.zero_grad()
+            obj.forward(coords).imag.sum().backward()
+            opt.step()
+        materialized = obj.obj
+        assert materialized.shape == (1, 16, 16) and not materialized.is_complex()
+        assert float(materialized.min()) >= 0.0  # softplus enforces non-negative potential
+
+
+# --------------------------------------------------------------------------- #
+# Implicit-flag sync (ObjectTensorDecomp is implicit, like ObjectINR)
+# --------------------------------------------------------------------------- #
+class TestImplicitSync:
+    def test_implicit_flag_synced_from_obj_model(self):
+        ptycho, _ = _build_kplanes_ptycho()
+        assert ptycho.obj_model.is_implicit is True
+        assert ptycho.dset.implicit_object is True
+
+
+# --------------------------------------------------------------------------- #
+# PPLR optimizer pass-through through the ptychography optimizer layer
+# --------------------------------------------------------------------------- #
+class TestPPLRPassthrough:
+    def test_setter_passes_nested_dict_through_unmutated(self):
+        ptycho, _ = _build_kplanes_ptycho()
+        params = _pplr_params()
+        snapshot = {k: dict(v) for k, v in params.items()}  # shallow copy of inner dicts
+        ptycho.optimizer_params = params
+        ptycho.set_optimizers()
+        # the model received the PPLR groups, not a corrupted single-optimizer dict
+        stored = ptycho.obj_model.optimizer_params
+        assert set(stored.keys()) == {"grids", "sigma_net"}
+        # the ptychography setter must not have mutated the caller's nested dict
+        assert params["object"] == snapshot["object"]
+        assert "name" not in params["object"] and "lr" not in params["object"]
+        # optimizer has the two named groups with the requested LRs
+        opt = ptycho.obj_model.optimizer
+        assert len(opt.param_groups) == 2
+        assert sorted(pg["lr"] for pg in opt.param_groups) == [_LR_SIGMA, _LR_GRIDS]
+
+
+# --------------------------------------------------------------------------- #
+# End-to-end reconstruction
+# --------------------------------------------------------------------------- #
+@pytest.mark.slow
+class TestObjectTensorDecompReconstruction:
+    def test_loss_decreases_and_recovers_object(self):
+        ptycho, gt_phase = _build_kplanes_ptycho()
+        ptycho.reconstruct(
+            num_iters=150,
+            optimizer_params=_pplr_params(),  # probe frozen
+            scheduler_params={"object": SchedulerParams.CosineAnnealing()},
+            batch_size=200,
+        )
+        losses = np.array(ptycho._iter_losses)
+        assert losses[-1] < _LOSS_RATIO * losses[0]
+        assert _best_corr(ptycho.obj[0], gt_phase) > _CORR_THRESHOLD
+
+    def test_reconstruct_with_reset_runs(self):
+        """reset=True replays reset_optimizer with the disabled sentinel before applying PPLR."""
+        ptycho, _ = _build_kplanes_ptycho()
+        ptycho.reconstruct(
+            num_iters=5,
+            reset=True,
+            optimizer_params=_pplr_params(),
+            batch_size=200,
+        )
+        losses = np.array(ptycho._iter_losses)
+        assert np.isfinite(losses).all()
+        # the object optimizer still has its two PPLR groups after the reset cycle
+        assert len(ptycho.obj_model.optimizer.param_groups) == 2
+
+    def test_tilted_reconstruct_runs(self):
+        """A tilted K-Planes object reconstructs end-to-end with a 3-group PPLR optimizer."""
+        pdset, gt_phase, probe = _build_synthetic_dataset()
+        obj = ObjectTensorDecomp.from_uniform(
+            num_slices=1,
+            M_features=12,
+            resolution=(16, 64, 64),
+            multiscale_res_multipliers=(0.25, 0.5, 1.0),
+            tilted=True,
+            T=4,
+            obj_type="pure_phase",
+            rng=0,
+        )
+        probe_model = ProbePixelated.from_array(
+            num_probes=1,
+            probe_params={
+                "energy": PROBE_ENERGY,
+                "C10": C10,
+                "semiangle_cutoff": _semiangle_mrad(),
+            },
+            probe_array=probe.copy(),
+        )
+        ptycho = Ptychography.from_models(
+            dset=pdset,
+            obj_model=obj,
+            probe_model=probe_model,
+            detector_model=DetectorPixelated(),
+            rng=0,
+            verbose=False,
+        )
+        ptycho.preprocess(obj_padding_px=(PAD, PAD))
+        ptycho.reconstruct(
+            num_iters=60,
+            optimizer_params={
+                "object": {
+                    "grids": OptimizerParams.Adam(lr=_LR_GRIDS),
+                    "sigma_net": OptimizerParams.Adam(lr=_LR_SIGMA),
+                    "so3": OptimizerParams.Adam(lr=1e-3),
+                }
+            },
+            batch_size=200,
+        )
+        losses = np.array(ptycho._iter_losses)
+        assert np.isfinite(losses).all()
+        assert losses[-1] < losses[0]
+
+    def test_tv_constraint_reconstruct_runs(self):
+        """A reconstruct with the in-plane TV soft constraint runs and stays finite."""
+        ptycho, _ = _build_kplanes_ptycho()
+        ptycho.reconstruct(
+            num_iters=20,
+            optimizer_params=_pplr_params(),
+            constraints={"object": {"tv_weight_xy": 1e-3}},
+            batch_size=200,
+        )
+        losses = np.array(ptycho._iter_losses)
+        assert np.isfinite(losses).all()
+        assert ptycho.obj_model.constraints.tv_weight_xy == 1e-3
+
+    def test_multislice_runs(self):
+        ptycho, _ = _build_kplanes_ptycho(
+            num_slices=2, slice_thicknesses=20.0, resolution=(8, 64, 64)
+        )
+        ptycho.reconstruct(
+            num_iters=5,
+            optimizer_params=_pplr_params(),
+            batch_size=200,
+        )
+        assert ptycho.obj.shape[0] == 2
+
+    def test_scheduler_scales_pplr_groups(self):
+        ptycho, _ = _build_kplanes_ptycho()
+        ptycho.reconstruct(
+            num_iters=10,
+            optimizer_params=_pplr_params(),
+            scheduler_params={"object": SchedulerParams.CosineAnnealing()},
+            batch_size=200,
+        )
+        # one cosine schedule drives both param groups; both LRs have decayed below their start
+        lrs = sorted(pg["lr"] for pg in ptycho.obj_model.optimizer.param_groups)
+        assert lrs[0] < _LR_SIGMA and lrs[1] < _LR_GRIDS
+
+    def test_save_load_roundtrip(self, tmp_path):
+        ptycho, _ = _build_kplanes_ptycho()
+        ptycho.reconstruct(
+            num_iters=20,
+            optimizer_params=_pplr_params(),
+            batch_size=200,
+        )
+        obj_before = ptycho.obj.copy()
+        path = tmp_path / "kplanes_ptycho.zip"
+        ptycho.save(path, mode="o", save_raw_data=True)  # persist dset so loaded.dset works
+        loaded = autoserialize_load(path)
+        assert loaded.obj_model.is_implicit is True
+        assert loaded.dset.implicit_object is True
+        assert loaded.obj_model.name == "ObjectTensorDecomp"
+        np.testing.assert_allclose(loaded.obj, obj_before, rtol=1e-5, atol=1e-6)
+        # continued training still runs after reload
+        loaded.reconstruct(
+            num_iters=5,
+            optimizer_params=_pplr_params(),
+            batch_size=200,
+        )

From 8c6ea7ee8443d53b49661d6d12e2fb9c72369107 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Wed, 3 Jun 2026 12:17:32 -0700
Subject: [PATCH 33/59] bug fix

---
 src/quantem/diffractive_imaging/object_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 140ca5852..659039c27 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -327,7 +327,7 @@ def mask(self, mask: torch.Tensor | np.ndarray):
             ndim=3,
             expand_dims=True,
         )
-        self._mask = mask.to(self.device).expand(self.num_slices, -1, -1)
+        self._mask = mask.to(self.device).expand(self.num_slices, -1, -1).contiguous()
 
     @property
     @abstractmethod

From d906a100530726379b65737018a47d1e08139ed6 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Wed, 3 Jun 2026 13:54:30 -0700
Subject: [PATCH 34/59] device and lr changes. all ptycholite.

---
 .../diffractive_imaging/object_models.py      |  2 +-
 .../diffractive_imaging/ptychography_lite.py  | 20 ++++++++++++-------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 659039c27..38eabc797 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1045,7 +1045,7 @@ def from_pixelated(
             num_slices=pixelated.num_slices,
             slice_thicknesses=pixelated.slice_thicknesses,
             input_noise_std=input_noise_std,
-            device=pixelated.device,
+            device=device,
             obj_type=pixelated.obj_type,
             rng=pixelated._rng_seed,
         )
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 0bca7a494..41ce2c780 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -299,7 +299,9 @@ def from_ptycholite(
         cls,
         ptycholite: PtychoLite,
         pretrain_iters: int | None = None,
-        pretrain_lr: float = 1e-3,
+        pretrain_object_lr: float | None = None,
+        pretrain_probe_lr: float = 1e-3,
+        pretrain_lr: float | None = None,
         pretrain_probe: bool = True,
         pretrain_object: bool = True,
         normalize_object_plotting: bool = True,
@@ -311,11 +313,15 @@ def from_ptycholite(
         log_prefix: str = "",
         log_images_every: int = 10,
         log_probe_images: bool = False,
-        device: Literal["cpu", "gpu", "cuda"] = "cpu",
+        device: str | int | torch.device = "cpu",
         verbose: int | bool = True,
     ) -> Self:
-        if device == "gpu":
-            device = "cuda"
+        if pretrain_object_lr is None:
+            pretrain_object_lr = 1e-3 if pretrain_lr is None else pretrain_lr
+        elif pretrain_lr is not None and pretrain_lr != pretrain_object_lr:
+            raise ValueError("Got conflicting values for pretrain_lr and pretrain_object_lr.")
+
+        device, _ = config.validate_device(device)
         # Object model
         obj_dip = CNN2d(
             in_channels=ptycholite.obj_model.num_slices,
@@ -352,14 +358,14 @@ def from_ptycholite(
                     num_iters=pretrain_iters,
                     optimizer_params={
                         "name": "adamw",
-                        "lr": pretrain_lr,
+                        "lr": pretrain_object_lr,
                     },
                     scheduler_params={
                         "name": "plateau",
                         "factor": 0.5,
                     },
                     apply_constraints=False,
-                    device=config.get("device"),
+                    device=device,
                     normalize_object_plotting=normalize_object_plotting,
                 )
             if pretrain_probe:
@@ -368,7 +374,7 @@ def from_ptycholite(
                     num_iters=pretrain_iters,
                     optimizer_params={
                         "name": "adamw",
-                        "lr": 1e-3,
+                        "lr": pretrain_probe_lr,
                     },
                     scheduler_params={
                         "name": "plateau",

From e067d6b4d940014eb1f7a000eb5289b835a0aa50 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 3 Jun 2026 14:19:07 -0700
Subject: [PATCH 35/59] fixing linter errors kplanes

---
 src/quantem/core/ml/models/kplanes.py         | 49 ++++++-------------
 .../test_object_tensor_decomp.py              | 12 +++--
 2 files changed, 25 insertions(+), 36 deletions(-)

diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py
index 5aa1f5e0e..84b83dea8 100644
--- a/src/quantem/core/ml/models/kplanes.py
+++ b/src/quantem/core/ml/models/kplanes.py
@@ -3,7 +3,7 @@
 """
 
 import itertools
-from typing import Callable, Optional, Sequence
+from typing import Callable, Literal, Sequence, cast
 
 # import tinycudann as tcnn
 import torch
@@ -177,7 +177,7 @@ def __init__(
         input_coords_dims: int = 3,
         M_features: int = 32,
         resolution: Sequence[int] = (200, 200, 200),
-        multiscale_res_multipliers: Optional[Sequence[float]] = None,
+        multiscale_res_multipliers: Sequence[float] | None = None,
         concat_features: bool = True,
         density_activation: Callable = lambda x: F.softplus(
             x - 1
@@ -191,7 +191,7 @@ def __init__(
         Assume coords are [-1, 1] in each dimension.
         """
         super().__init__()
-        self._td_type = "kplanes"
+        self.td_type = "kplanes"
         self.grid_dimensions = grid_dimensions
         self.input_coords_dims = input_coords_dims
         self.M_features = M_features
@@ -279,26 +279,10 @@ def get_params(self) -> dict[str, list[torch.nn.Parameter]]:
     def param_keys(self) -> list[str]:
         return ["grids", "sigma_net"]
 
-    @property
-    def td_type(self) -> str:
-        return self._td_type
-
-    @td_type.setter
-    def td_type(self, td_type: str):
-        if not isinstance(td_type, str):
-            raise TypeError("td_type must be a string")
-        self._td_type = td_type
-
     @property
     def tilted(self) -> bool:
         return False
 
-    @tilted.setter
-    def tilted(self, tilted: bool):
-        if not isinstance(tilted, bool):
-            raise TypeError("tilted must be a boolean")
-        self._tilted = tilted
-
     @property
     def grids(self) -> torch.nn.ParameterList:
         return self._grids
@@ -416,18 +400,18 @@ def __init__(
         input_coords_dims: int = 3,
         M_features: int = 32,
         resolution: Sequence[int] = (200, 200, 200),
-        multiscale_res_multipliers: Optional[Sequence[float]] = None,
+        multiscale_res_multipliers: Sequence[float] | None = None,
         density_activation: Callable = lambda x: F.softplus(x - 1),
         # TILTED parameters
         T: int = 4,
-        tau_init: str = "random",
+        tau_init: Literal["random", "identity"] = "random",
         # Hybrid MLP parameters
         use_hybrid_mlp: bool = False,
         hybrid_hidden_dim: int = 64,
         hybrid_num_layers: int = 2,
         so3_param_type: str = "r9svd",
     ):
-        self._td_type = "tilted"
+        self.td_type = "tilted"
         if input_coords_dims != 3:
             raise NotImplementedError("KPlanesTILTED is implemented for 3D only.")
         if T < 1:
@@ -552,7 +536,7 @@ def extract_tau_state(self) -> torch.Tensor:
         -------
         torch.Tensor of shape (T, 3, 3)
         """
-        return self.so3.M.detach().cpu().clone()
+        return cast(torch.Tensor, self.so3.M).detach().cpu().clone()
 
     def load_tau_state(self, M: torch.Tensor) -> None:
         """
@@ -572,7 +556,8 @@ def load_tau_state(self, M: torch.Tensor) -> None:
                 f"Make sure T matches between phase 1 and phase 2."
             )
         with torch.no_grad():
-            self.so3.M.copy_(M.to(self.so3.M.device))
+            so3_M = cast(torch.Tensor, self.so3.M)
+            so3_M.copy_(M.to(so3_M.device))
 
     # ------------------------------------------------------------------
     # Pretty print
@@ -586,7 +571,9 @@ def extra_repr(self) -> str:
             f"num_scales={len(self.multiscale_res_multipliers)}"
         )
 
-    def set_so3_param_type(self, so3_param_type: str, init: str = "rand") -> None:
+    def set_so3_param_type(
+        self, so3_param_type: str, init: Literal["random", "identity"] = "random"
+    ) -> None:
         """
         Set the SO3 parameterization type.
 
@@ -681,14 +668,14 @@ def __init__(
         self,
         C: int = 4,  # channels per transform per scale
         resolution: Sequence[int] = (128, 128, 128),
-        multiscale_res_multipliers: Optional[Sequence[int]] = None,
+        multiscale_res_multipliers: Sequence[float] | None = None,
         T: int = 4,
-        tau_init: str = "random",
+        tau_init: Literal["random", "identity"] = "random",
         density_activation: Callable = lambda x: F.softplus(x - 1),
         so3_param_type: str = "r9svd",
     ):
         super().__init__()
-        self._td_type = "cp_tilted"
+        self.td_type = "cp_tilted"
         self.T = T
         self.C = C
         self.multiscale_res_multipliers = list(multiscale_res_multipliers or [1])
@@ -738,12 +725,8 @@ def get_params(self):
     def param_keys(self):
         return ["grids", "sigma_net", "so3"]
 
-    @property
-    def td_type(self) -> str:
-        return self._td_type
-
     def extract_tau_state(self) -> torch.Tensor:
-        return self.so3.M.detach().clone()
+        return cast(torch.Tensor, self.so3.M).detach().clone()
 
     @property
     def tilted(self) -> bool:
diff --git a/tests/diffractive_imaging/test_object_tensor_decomp.py b/tests/diffractive_imaging/test_object_tensor_decomp.py
index dabf2b6b8..b3920a5aa 100644
--- a/tests/diffractive_imaging/test_object_tensor_decomp.py
+++ b/tests/diffractive_imaging/test_object_tensor_decomp.py
@@ -340,7 +340,7 @@ def test_from_model_warns_on_lambda_activation(self):
 
     def test_from_model_rejects_non_kplanes(self):
         with pytest.raises(TypeError):
-            ObjectTensorDecomp.from_model(nn.Linear(3, 1), num_slices=1)
+            ObjectTensorDecomp.from_model(nn.Linear(3, 1), num_slices=1)  # pyright: ignore[reportArgumentType]
 
     def test_tilted_from_uniform(self):
         """tilted=True builds a KPlanesTILTED with an extra `so3` param group."""
@@ -362,6 +362,7 @@ def test_tilted_from_uniform(self):
         # vacuum init holds for the tilted decoder too
         assert torch.allclose(patches, torch.ones_like(patches), atol=1e-6)
         obj.set_optimizer({k: OptimizerParams.Adam(lr=1e-2) for k in obj.model.param_keys})
+        assert obj.optimizer is not None
         assert len(obj.optimizer.param_groups) == 3
 
     def test_from_model_accepts_cptilted(self):
@@ -450,6 +451,7 @@ def test_setter_passes_nested_dict_through_unmutated(self):
         assert "name" not in params["object"] and "lr" not in params["object"]
         # optimizer has the two named groups with the requested LRs
         opt = ptycho.obj_model.optimizer
+        assert opt is not None
         assert len(opt.param_groups) == 2
         assert sorted(pg["lr"] for pg in opt.param_groups) == [_LR_SIGMA, _LR_GRIDS]
 
@@ -483,7 +485,9 @@ def test_reconstruct_with_reset_runs(self):
         losses = np.array(ptycho._iter_losses)
         assert np.isfinite(losses).all()
         # the object optimizer still has its two PPLR groups after the reset cycle
-        assert len(ptycho.obj_model.optimizer.param_groups) == 2
+        opt = ptycho.obj_model.optimizer
+        assert opt is not None
+        assert len(opt.param_groups) == 2
 
     def test_tilted_reconstruct_runs(self):
         """A tilted K-Planes object reconstructs end-to-end with a 3-group PPLR optimizer."""
@@ -564,7 +568,9 @@ def test_scheduler_scales_pplr_groups(self):
             batch_size=200,
         )
         # one cosine schedule drives both param groups; both LRs have decayed below their start
-        lrs = sorted(pg["lr"] for pg in ptycho.obj_model.optimizer.param_groups)
+        opt = ptycho.obj_model.optimizer
+        assert opt is not None
+        lrs = sorted(pg["lr"] for pg in opt.param_groups)
         assert lrs[0] < _LR_SIGMA and lrs[1] < _LR_GRIDS
 
     def test_save_load_roundtrip(self, tmp_path):

From c2bf1bfab5e5596b17b19c6fa8ad000df70a4769 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 3 Jun 2026 17:01:15 -0700
Subject: [PATCH 36/59] adding hard constraints for INR positivity

---
 .../diffractive_imaging/object_models.py      | 106 +++++++++++++-----
 .../diffractive_imaging/ptychography.py       |   8 +-
 tests/diffractive_imaging/test_object_inr.py  |  48 +++++++-
 .../test_object_tensor_decomp.py              |  22 ++--
 4 files changed, 143 insertions(+), 41 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 2db63a874..54ed48969 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -144,31 +144,47 @@ class Raster(Constraints):
 
     @dataclass
     class INR(Constraints):
-        """Constraints for the implicit (``ObjectINR``) object representation.
+        """Constraints for the implicit (``ObjectINR`` / ``ObjectTensorDecomp``) object.
 
-        An INR has no grid to project, so the grid-based hard constraints of
-        ``Raster`` (positivity, filtering, FOV masking) do not apply. The
-        regularizers that do carry over are soft penalties evaluated at sampled
-        coordinates.
+        An implicit object has no grid to clamp/filter in place, so the grid-based hard
+        constraints of ``Raster`` do not apply directly. Instead: positivity is a **soft**
+        penalty evaluated at sampled coordinates (keeps the network output linear so a
+        zero-background ``potential`` fits without the vanishing/dead gradients of a
+        softplus/relu output activation), and the potential baseline is a **display gauge** on
+        the materialized object (a constant potential offset is a global phase, i.e.
+        diffraction-invariant).
 
         Attributes
         ----------
         tv_weight_z : float, default ``0.0``
-            Soft penalty. Weight on the depth-axis (``z``) total-variation term,
-            evaluated via finite differences at randomly sampled coordinates.
-            Multislice (``num_slices > 1``) only.
+            Soft penalty. Depth-axis (``z``) total variation, finite-differenced at randomly
+            sampled coordinates. Multislice (``num_slices > 1``) only.
         tv_weight_xy : float, default ``0.0``
-            Soft penalty. Weight on the in-plane (``y``, ``x``) total-variation
-            term, evaluated via finite differences at randomly sampled coordinates.
+            Soft penalty. In-plane (``y``, ``x``) total variation at sampled coordinates.
+        positivity_weight : float, default ``0.0``
+            Soft penalty. Weight on ``mean(relu(-value))`` at sampled coordinates -- drives the
+            ``potential`` non-negative. ``obj_type="potential"`` only; ignored otherwise. Scale it
+            relative to the data-loss magnitude; increase if negativity persists.
+        fix_potential_baseline : bool, default ``False``
+            ``obj_type="potential"`` only. Subtracts a background offset from the *materialized*
+            object so background sits at zero (then clamps >= 0). Display gauge only -- it does not
+            perturb the reconstruction (the forward queries the network directly), mirroring the
+            ``Raster`` constraint of the same name.
+        fix_potential_baseline_factor : float, default ``1.0``
+            Scales the subtracted baseline offset (``<1`` relaxes the anchoring).
         """
 
         # soft constraints (evaluated at sampled coordinates)
         tv_weight_z: float = 0.0
         tv_weight_xy: float = 0.0
+        positivity_weight: float = 0.0
+        # hard / display constraints (applied to the materialized object only)
+        fix_potential_baseline: bool = False
+        fix_potential_baseline_factor: float = 1.0
         _name: str = "inr"
 
-        soft_constraint_keys = ["tv_weight_z", "tv_weight_xy"]
-        hard_constraint_keys = []
+        soft_constraint_keys = ["tv_weight_z", "tv_weight_xy", "positivity_weight"]
+        hard_constraint_keys = ["fix_potential_baseline", "fix_potential_baseline_factor"]
 
     @classmethod
     def parse_dict(cls, d: dict) -> "PtychoObjConstraintsType":
@@ -1618,12 +1634,12 @@ def from_uniform(
         The HSiren's final layer is zero-initialized so the object starts uniform (a
         diffraction-equivalent vacuum), matching ``ObjectPixelated.from_uniform``.
 
-        ``final_activation`` sets the output nonlinearity. When ``None`` (default) it is chosen
-        from ``obj_type``: ``"identity"`` for ``pure_phase``, and ``"softplus"`` for
-        ``potential`` -- a non-negative activation enforces the potential's positivity (min-value)
-        constraint directly at the output (use ``"relu"`` for a hard floor). With the zeroed final
-        layer the potential starts uniform (``softplus(0) = ln 2``), which is just a global phase
-        and hence diffraction-equivalent to vacuum.
+        ``final_activation`` sets the output nonlinearity. The default (``None``) is ``"identity"``
+        for both ``pure_phase`` and ``potential``: enforcing positivity at the output (softplus /
+        relu) makes a zero-background potential hard to fit (vanishing / dead gradients), so for
+        ``potential`` positivity is instead the soft ``positivity_weight`` constraint. With the
+        zeroed final layer the object starts at 0 (vacuum). Pass ``final_activation="softplus"`` to
+        opt back into output-activation positivity.
 
         Note
         ----
@@ -1634,7 +1650,7 @@ def from_uniform(
         features may want a larger value. Pair omega_0 with the object learning rate.
         """
         if final_activation is None:
-            final_activation = "softplus" if obj_type == "potential" else "identity"
+            final_activation = "identity"
         model = HSiren(
             in_features=3,
             out_features=1,
@@ -2032,12 +2048,23 @@ def apply_hard_constraints(
 
         Unlike the grid-based ``Raster`` constraints, an INR has nothing to clamp or filter in
         place. For ``pure_phase`` we recenter the phase to zero mean (a global-phase gauge) so the
-        displayed object matches the pixelated convention; ``potential`` is returned as-is (the INR
-        constraint set has no positivity/baseline fields — potential is left unconstrained).
+        displayed object matches the pixelated convention. For ``potential``, if
+        ``fix_potential_baseline`` is set, subtract a background offset (mask background mean, else
+        the global min) scaled by ``fix_potential_baseline_factor`` and clamp ``>= 0`` -- a display
+        gauge (a constant potential offset is a global phase, hence diffraction-invariant), so it
+        does not affect the reconstruction. Positivity *during* the reconstruction is the soft
+        ``positivity_weight`` penalty, not a projection here.
         """
         with torch.no_grad():
             if self.obj_type == "pure_phase":
                 return raw - raw.mean()
+            if self.constraints.fix_potential_baseline:
+                if mask is not None and mask.numel() and (mask < 0.5 * mask.max()).any():
+                    offset = raw[mask < 0.5 * mask.max()].mean()
+                else:
+                    offset = raw.amin()
+                offset = offset * self.constraints.fix_potential_baseline_factor
+                return torch.clamp(raw - offset, min=0.0)
             return raw
 
     def apply_soft_constraints(
@@ -2058,9 +2085,31 @@ def apply_soft_constraints(
             tv_loss = self._sampled_tv_loss(w_z, w_xy)
             loss = loss + tv_loss
             self.add_soft_constraint_loss("tv_loss", tv_loss)
+        w_pos = self.constraints.positivity_weight
+        if w_pos > 0 and self.obj_type == "potential":
+            pos_loss = self._sampled_positivity_loss(w_pos)
+            loss = loss + pos_loss
+            self.add_soft_constraint_loss("positivity_loss", pos_loss)
         self.accumulate_constraint_losses()
         return loss
 
+    def _sampled_positivity_loss(self, weight: float, num_samples: int = 4096) -> torch.Tensor:
+        """Differentiable positivity penalty for ``potential``: ``weight * mean(relu(-value))`` at
+        random coordinates. Keeps the network output linear (identity activation), so a
+        zero-background potential fits without the vanishing (softplus) / dead (relu) gradients of
+        an output activation; negative regions get a constant linear restoring force toward 0.
+        """
+        real_dtype = getattr(torch, config.get("dtype_real"))
+        coords_xy = (
+            torch.rand(
+                num_samples, 2, device=self.device, dtype=real_dtype, generator=self._rng_torch
+            )
+            * 2.0
+            - 1.0
+        )
+        value = self._query_phase(coords_xy)  # (S, num_samples) -- the queried potential
+        return weight * torch.relu(-value).mean()
+
     def _sampled_tv_loss(self, w_z: float, w_xy: float, num_samples: int = 4096) -> torch.Tensor:
         """Finite-difference TV over (z, y, x) at randomly sampled coordinates."""
         real_dtype = getattr(torch, config.get("dtype_real"))
@@ -2194,14 +2243,15 @@ def from_uniform(  # pyright: ignore[reportIncompatibleMethodOverride]  # KPlane
 
         ``tilted=False`` builds a plain :class:`KPlanes`; ``tilted=True`` builds a
         :class:`KPlanesTILTED` with ``T`` learned SO(3) rotations. The decoder's final layer is
-        zeroed so the object starts uniform (a global phase, diffraction-equivalent to vacuum),
-        matching ``ObjectINR.from_uniform``. ``density_activation`` is chosen from ``obj_type``:
-        ``nn.Identity`` for ``pure_phase`` (phase may be negative) and ``nn.Softplus`` for
-        ``potential`` (non-negative). ``resolution`` is the feature-plane resolution ``(z, y, x)``
-        and is independent of the reconstructed object grid; for multislice set ``resolution[0]``
-        to span the slices.
+        zeroed so the object starts at 0 (vacuum), matching ``ObjectINR.from_uniform``. The decoder
+        is **identity**-activated for both ``pure_phase`` and ``potential``: a softplus/relu output
+        activation makes a zero-background potential hard to fit (vanishing/dead gradients), so for
+        ``potential`` enforce positivity with the soft ``positivity_weight`` constraint instead (use
+        ``from_model`` with an ``nn.Softplus`` activation to opt back in). ``resolution`` is the
+        feature-plane resolution ``(z, y, x)`` and is independent of the reconstructed object grid;
+        for multislice set ``resolution[0]`` to span the slices.
         """
-        density_activation: nn.Module = nn.Softplus() if obj_type == "potential" else nn.Identity()
+        density_activation: nn.Module = nn.Identity()
         ms = list(multiscale_res_multipliers) if multiscale_res_multipliers is not None else None
         model: KPlanesType
         if tilted:
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index f39056215..5cbe53542 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -85,6 +85,12 @@ def _ddp_ptycho_worker(
             result_path,
         )
 
+    # Synchronize before teardown so every rank finishes 
+    if dist.is_available() and dist.is_initialized():
+        if torch.cuda.is_available():
+            dist.barrier(device_ids=[device_id])
+        else:
+            dist.barrier()
     dist.destroy_process_group()
 
 
@@ -245,7 +251,7 @@ def reconstruct(
         Multi-GPU (``device`` is a list) launches worker processes via ``mp.spawn`` when called
         from a notebook, or uses the existing distributed process group when launched with
         ``torchrun``. Only autograd mode is supported for multi-GPU in this release.
-        
+
         ``loss_type`` selects the data-fidelity criterion: a registered name
         (``"l2_amplitude"`` [default], ``"l1_amplitude"``, ``"l2_intensity"``, ``"l1_intensity"``,
         ``"poisson"``, ``"smooth_l1_amplitude"``, ``"s3im_amplitude"``) or a ``DataCriterion``
diff --git a/tests/diffractive_imaging/test_object_inr.py b/tests/diffractive_imaging/test_object_inr.py
index 76f752a6d..b61f9cb3f 100644
--- a/tests/diffractive_imaging/test_object_inr.py
+++ b/tests/diffractive_imaging/test_object_inr.py
@@ -262,10 +262,16 @@ def test_complex_obj_type_not_supported(self):
         with pytest.raises(NotImplementedError):
             ObjectINR.from_uniform(num_slices=1, obj_type="complex")
 
-    def test_potential_obj_type(self):
-        """`potential` is real-valued like `pure_phase` but uses a non-negative output activation
-        (softplus) to enforce the positivity / min-value constraint."""
-        obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
+    def test_potential_obj_type_softplus_opt_in(self):
+        """`potential` is real-valued; the default activation is now identity, but passing
+        ``final_activation="softplus"`` opts back into output-activation positivity (min >= 0)."""
+        obj = ObjectINR.from_uniform(
+            num_slices=1,
+            obj_type="potential",
+            final_activation="softplus",
+            hidden_features=32,
+            rng=0,
+        )
         obj._initialize_obj((1, 16, 16))
         assert obj.obj_type == "potential"
         assert not obj.dtype.is_complex  # real-valued object
@@ -282,6 +288,40 @@ def test_potential_obj_type(self):
         assert materialized.shape == (1, 16, 16) and not materialized.is_complex()
         assert float(materialized.min()) >= 0.0  # softplus enforces non-negative potential
 
+    def test_potential_identity_default_and_positivity_penalty(self):
+        """Default ``potential`` activation is identity (vacuum is exactly 0); the soft
+        ``positivity_weight`` penalty drives a forced-negative potential non-negative."""
+        obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
+        obj._initialize_obj((1, 24, 24))
+        # identity + zeroed final layer -> vacuum is exactly 0 (not softplus(0) = ln 2)
+        assert float(obj._materialize_obj().abs().max()) == pytest.approx(0.0, abs=1e-6)
+        # force the whole potential negative via the (zero-weight) final-layer bias
+        with torch.no_grad():
+            obj.model.net[-2].bias.fill_(-0.5) # type:ignore 
+        assert float(obj._materialize_obj().min()) == pytest.approx(-0.5, abs=1e-3)
+        obj.constraints = {"positivity_weight": 1.0}
+        assert float(obj._sampled_positivity_loss(1.0)) == pytest.approx(0.5, abs=0.05)
+        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
+        for _ in range(80):
+            opt.zero_grad()
+            obj.apply_soft_constraints().backward()
+            opt.step()
+        assert float(obj._materialize_obj().min()) > -1e-2  # driven non-negative
+
+    def test_fix_potential_baseline_gauge(self):
+        """``fix_potential_baseline`` subtracts the background offset from the materialized
+        potential (display gauge; the reconstruction forward path is unaffected)."""
+        obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
+        obj._initialize_obj((1, 16, 16))
+        with torch.no_grad():
+            obj.model.net[-2].bias.fill_(1.0) # type:ignore  # constant +1 background
+        raw = obj._materialize_obj()
+        assert float(raw.min()) == pytest.approx(1.0, abs=0.2)
+        obj.constraints = {"fix_potential_baseline": True}
+        disp = obj.apply_hard_constraints(raw, mask=obj.mask)
+        assert float(disp.min()) == pytest.approx(0.0, abs=1e-3)  # background pinned to 0
+        assert float(disp.min()) >= 0.0  # clamped non-negative
+
     def test_from_pixelated_with_model(self):
         """from_pixelated can wrap a directly-passed INR model (like ObjectDIP.from_pixelated)."""
         from quantem.core.ml.inr import HSiren
diff --git a/tests/diffractive_imaging/test_object_tensor_decomp.py b/tests/diffractive_imaging/test_object_tensor_decomp.py
index b3920a5aa..9343da750 100644
--- a/tests/diffractive_imaging/test_object_tensor_decomp.py
+++ b/tests/diffractive_imaging/test_object_tensor_decomp.py
@@ -408,19 +408,25 @@ def _corr_to_pix(arr):
         kp.reset()
         assert _corr_to_pix(kp.obj[0].detach().cpu().numpy()) > 0.9
 
-    def test_potential_obj_type_positive(self):
+    def test_potential_identity_default_and_positivity_penalty(self):
+        """Potential K-Planes uses an identity decoder by default (softplus/relu fit zero-background
+        potentials poorly); the inherited soft positivity penalty drives it non-negative instead."""
         obj = self._obj(obj_type="potential", resolution=(16, 16, 16))
         obj._initialize_obj((1, 16, 16))
         assert obj.obj_type == "potential"
-        opt = torch.optim.Adam(obj.model.parameters(), lr=1e-2)
-        coords = torch.rand(2, 6, 6, 2) * 2 - 1
-        for _ in range(5):
+        assert isinstance(obj.model.density_activation, nn.Identity)  # identity, not softplus
+        # force the whole potential negative via the (zero-weight) decoder bias
+        with torch.no_grad():
+            obj.model.sigma_net.bias.fill_(-0.5) # type:ignore 
+        assert float(obj._materialize_obj().min()) == pytest.approx(-0.5, abs=1e-3)
+        obj.constraints = {"positivity_weight": 1.0}
+        assert float(obj._sampled_positivity_loss(1.0)) == pytest.approx(0.5, abs=0.05)
+        opt = torch.optim.Adam(obj.model.parameters(), lr=2e-2)
+        for _ in range(100):
             opt.zero_grad()
-            obj.forward(coords).imag.sum().backward()
+            obj.apply_soft_constraints().backward()
             opt.step()
-        materialized = obj.obj
-        assert materialized.shape == (1, 16, 16) and not materialized.is_complex()
-        assert float(materialized.min()) >= 0.0  # softplus enforces non-negative potential
+        assert float(obj._materialize_obj().min()) > -1e-2  # driven non-negative
 
 
 # --------------------------------------------------------------------------- #

From 6e111d5b2efe586cc2393baa083492c4a1852faa Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Thu, 4 Jun 2026 10:29:25 -0700
Subject: [PATCH 37/59] adding show probe function

---
 .../ptychography_visualizations.py            | 113 ++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index f1e7010de..19482e02f 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -273,6 +273,119 @@ def show_probe(
 
         show_2d(probes, title=titles, scalebar=scalebar, **kwargs)
 
+    def show_probe_top_bottom(
+        self,
+        probe: np.ndarray | None = None,
+        snapshot_iter: int | None = None,
+        probe_index: int | None = 0,
+        quantity: Literal["intensity", "amplitude", "phase"] = "intensity",
+        fftshift: bool = True,
+        return_arrays: bool = False,
+        **kwargs,
+    ):
+        """
+        Show the real-space probe at the top and bottom object surfaces.
+
+        The bottom surface probe is computed by propagating the reconstructed
+        probe through the multislice free-space propagators only. The reconstructed
+        object is not applied.
+
+        Parameters
+        ----------
+        probe : np.ndarray | None, optional
+            Probe array to show. If None, the current reconstructed probe is used.
+        snapshot_iter : int | None, optional
+            Snapshot iteration to show. If None, the final/current probe is shown.
+        probe_index : int | None, optional
+            Probe mode index to show. If None, mixed-state modes are summed
+            incoherently as intensity.
+        quantity : {"intensity", "amplitude", "phase"}, optional
+            Quantity to display for a single probe mode, by default "intensity".
+        fftshift : bool, optional
+            Whether to center the real-space probe for display.
+        return_arrays : bool, optional
+            If True, return the displayed top and bottom arrays in addition to
+            the figure and axes.
+        **kwargs
+            Additional arguments passed to show_2d.
+
+        Returns
+        -------
+        tuple
+            ``(fig, axs)`` by default, or ``(fig, axs, arrays)`` if
+            return_arrays is True.
+        """
+        if probe is None:
+            if snapshot_iter is not None:
+                if snapshot_iter < 0:
+                    snapshot_iter = len(self.snapshots) + snapshot_iter
+                snp = self.get_snapshot_by_iter(snapshot_iter, closest=True, cropped=True)
+                probe = snp["probe"]
+            else:
+                probe = self.probe
+        else:
+            probe = self._to_numpy(probe)
+            if probe.ndim == 2:
+                probe = probe[None, ...]
+
+        self.compute_propagator_arrays()
+
+        top = probe.copy()
+        bottom = top.copy()
+        for prop in self._to_numpy(self.propagators):
+            bottom = np.fft.ifft2(np.fft.fft2(bottom) * prop)
+
+        if probe_index is None:
+            if quantity != "intensity":
+                raise ValueError("probe_index=None is only supported for quantity='intensity'")
+            top_img = np.sum(np.abs(top) ** 2, axis=0)
+            bottom_img = np.sum(np.abs(bottom) ** 2, axis=0)
+            label = "Summed Probe Intensity"
+            cmap = kwargs.pop("cmap", "magma")
+        else:
+            if probe_index < 0 or probe_index >= top.shape[0]:
+                raise ValueError(
+                    f"probe_index must be between 0 and {top.shape[0] - 1}, got {probe_index}"
+                )
+            top_probe = top[probe_index]
+            bottom_probe = bottom[probe_index]
+
+            if quantity == "intensity":
+                top_img = np.abs(top_probe) ** 2
+                bottom_img = np.abs(bottom_probe) ** 2
+                cmap = kwargs.pop("cmap", "magma")
+            elif quantity == "amplitude":
+                top_img = np.abs(top_probe)
+                bottom_img = np.abs(bottom_probe)
+                cmap = kwargs.pop("cmap", "gray")
+            elif quantity == "phase":
+                top_img = np.angle(top_probe)
+                bottom_img = np.angle(bottom_probe)
+                cmap = kwargs.pop("cmap", config.get("viz.phase_cmap"))
+            else:
+                raise ValueError(
+                    f"quantity must be 'intensity', 'amplitude', or 'phase', got {quantity}"
+                )
+            label = f"Probe {probe_index + 1} {quantity.capitalize()}"
+
+        if fftshift:
+            top_img = np.fft.fftshift(top_img)
+            bottom_img = np.fft.fftshift(bottom_img)
+
+        scalebar = [{"sampling": self.sampling[0], "units": "Å"}, None]
+        titles = [f"Top Surface {label}", f"Bottom Surface {label}"]
+
+        fig, axs = show_2d(
+            [top_img, bottom_img],
+            title=titles,
+            cmap=cmap,
+            scalebar=scalebar,
+            **kwargs,
+        )
+        if return_arrays:
+            return fig, axs, {"top": top_img, "bottom": bottom_img}
+        return fig, axs
+
     def show_fourier_probe(self, probe: np.ndarray | None = None):
         """
         Show the Fourier transform of the probe.

From c71ad008fb307088ab0f793538163895f3e9caf2 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 13:34:36 -0700
Subject: [PATCH 38/59] adding alternative upsampling method to CNNs

---
 src/quantem/core/ml/blocks.py                 | 163 ++++++++++++++----
 .../ptychography_visualizations.py            |   2 +-
 2 files changed, 127 insertions(+), 38 deletions(-)

diff --git a/src/quantem/core/ml/blocks.py b/src/quantem/core/ml/blocks.py
index c2e678338..4465d757d 100644
--- a/src/quantem/core/ml/blocks.py
+++ b/src/quantem/core/ml/blocks.py
@@ -1,4 +1,4 @@
-from typing import Callable
+from typing import Callable, Literal
 
 import numpy as np
 import torch
@@ -141,7 +141,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class Upsample2dBlock(nn.Module):
-    """Upsampling block using transposed convolution or interpolation followed by convolution."""
+    """Upsampling block using transposed convolution or interpolation followed by convolution.
+
+    Two upsampling methods are available:
+
+    - ``"transpose"`` (default): transposed convolution (for ``scale_factor==2``,
+      otherwise interpolation) followed by a 1x1 convolution. This is the
+      original behavior.
+    - ``"resize"``: interpolation followed by a 3x3 convolution. Decoupling the
+      upsampling from the learned weights avoids the uneven kernel overlap that
+      produces checkerboard artifacts with transposed convolutions.
+    """
 
     def __init__(
         self,
@@ -151,6 +161,7 @@ def __init__(
         dtype: "torch.dtype" = torch.float32,
         scale_factor: int = 2,
         mode: str = "bilinear",
+        method: Literal["transpose", "resize"] = "transpose",
     ):
         """Initialize Upsample2dBlock.
 
@@ -168,41 +179,74 @@ def __init__(
             Factor by which to scale the input, by default 2
         mode : str, optional
             Interpolation mode, either "bilinear" or "nearest", by default "bilinear"
+        method : str, optional
+            Upsampling method, either "transpose" (transposed convolution) or
+            "resize" (interpolation followed by convolution). The "resize"
+            method reduces checkerboard artifacts. By default "transpose".
         """
         super().__init__()
         assert mode in ["bilinear", "nearest"], "Mode must be 'bilinear' or 'nearest'."
+        assert method in ("transpose", "resize"), "method must be 'transpose' or 'resize'."
         self.scale_factor = scale_factor
         self.mode = mode
         self.use_batchnorm = use_batchnorm
         self.dtype = dtype
-        self.upsample2x = nn.ConvTranspose2d(
-            input_channels,
-            input_channels,
-            kernel_size=3,
-            stride=2,
-            padding=(1, 1),
-            output_padding=(1, 1),
-            dtype=self.dtype,
-        )
-        self.conv = nn.Conv2d(
-            input_channels,
-            output_channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            dtype=self.dtype,
-            padding_mode="circular",
-        )
+        self.method = method
+
+        if method == "transpose":
+            self.upsample2x = nn.ConvTranspose2d(
+                input_channels,
+                input_channels,
+                kernel_size=3,
+                stride=2,
+                padding=(1, 1),
+                output_padding=(1, 1),
+                dtype=self.dtype,
+            )
+            self.conv = nn.Conv2d(
+                input_channels,
+                output_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                dtype=self.dtype,
+                padding_mode="circular",
+            )
+        else:
+            # Resize-conv: interpolate first, then a 3x3 conv smooths the result.
+            self.upsample2x = None
+            self.conv = nn.Conv2d(
+                input_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                dtype=self.dtype,
+                padding_mode="circular",
+            )
+
         if self.dtype.is_complex:
             self.bn = ComplexBatchNorm2D(output_channels)
         else:
             self.bn = nn.BatchNorm2d(output_channels)
 
+    def _interpolate(self, x: torch.Tensor) -> torch.Tensor:
+        # F.interpolate does not support complex tensors, so handle parts separately.
+        if x.is_complex():
+            real = F.interpolate(x.real, scale_factor=self.scale_factor, mode=self.mode)
+            imag = F.interpolate(x.imag, scale_factor=self.scale_factor, mode=self.mode)
+            return torch.complex(real, imag)
+        return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.scale_factor == 2:
-            x = self.upsample2x(x)
+        if self.method == "transpose":
+            if self.scale_factor == 2:
+                assert self.upsample2x is not None
+                x = self.upsample2x(x)
+            else:
+                x = self._interpolate(x)
         else:
-            x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
+            x = self._interpolate(x)
 
         x = self.conv(x)
         if self.use_batchnorm:
@@ -399,7 +443,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class Upsample3dBlock(nn.Module):
-    """3D upsampling block using transposed convolution followed by 1x1x1 convolution."""
+    """3D upsampling block.
+
+    Two upsampling methods are available:
+
+    - ``"transpose"`` (default): transposed convolution followed by a 1x1x1
+      convolution. This is the original behavior.
+    - ``"resize"``: nearest/trilinear interpolation followed by a 3x3x3
+      convolution. Decoupling the upsampling from the learned weights avoids
+      the uneven kernel overlap that produces checkerboard artifacts with
+      transposed convolutions.
+    """
 
     def __init__(
         self,
@@ -408,7 +462,8 @@ def __init__(
         use_batchnorm: bool = False,
         dtype: torch.dtype = torch.float32,
         scale_factor: int = 2,
-        mode: str = "trilinear",
+        mode: str = "nearest",
+        method: Literal["transpose", "resize"] = "transpose",
     ) -> None:
         """Initialize Upsample3dBlock.
 
@@ -425,29 +480,63 @@ def __init__(
         scale_factor : int, optional
             Factor by which to scale the input, by default 2
         mode : str, optional
-            Interpolation mode, by default "trilinear"
+            Interpolation mode used when ``method="resize"``, by default "trilinear"
+        method : str, optional
+            Upsampling method, either "transpose" (transposed convolution) or
+            "resize" (interpolation followed by convolution). The "resize"
+            method reduces checkerboard artifacts. By default "transpose".
         """
         super().__init__()
+        assert method in ("transpose", "resize"), "method must be 'transpose' or 'resize'."
         self.dtype = dtype
         self.use_batchnorm = use_batchnorm
-        self.upsample = nn.ConvTranspose3d(
-            input_channels,
-            input_channels,
-            kernel_size=3,
-            stride=2,
-            padding=1,
-            output_padding=1,
-            dtype=dtype,
-        )
-        self.conv = nn.Conv3d(input_channels, output_channels, kernel_size=1, dtype=dtype)
+        self.method = method
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+        if method == "transpose":
+            self.upsample = nn.ConvTranspose3d(
+                input_channels,
+                input_channels,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                output_padding=1,
+                dtype=dtype,
+            )
+            self.conv = nn.Conv3d(input_channels, output_channels, kernel_size=1, dtype=dtype)
+        else:
+            # Resize-conv: interpolate first, then a 3x3x3 conv smooths the result.
+            self.upsample = None
+            self.conv = nn.Conv3d(
+                input_channels,
+                output_channels,
+                kernel_size=3,
+                padding=1,
+                dtype=dtype,
+                padding_mode="circular",
+            )
+
         self.bn = (
             ComplexBatchNorm3D(output_channels)
             if dtype.is_complex
             else nn.BatchNorm3d(output_channels)
         )
 
+    def _interpolate(self, x: torch.Tensor) -> torch.Tensor:
+        # F.interpolate does not support complex tensors, so handle parts separately.
+        if x.is_complex():
+            real = F.interpolate(x.real, scale_factor=self.scale_factor, mode=self.mode)
+            imag = F.interpolate(x.imag, scale_factor=self.scale_factor, mode=self.mode)
+            return torch.complex(real, imag)
+        return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.upsample(x)
+        if self.method == "transpose":
+            assert self.upsample is not None
+            x = self.upsample(x)
+        else:
+            x = self._interpolate(x)
         x = self.conv(x)
         if self.use_batchnorm:
             x = self.bn(x)
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index f1e7010de..3b790bcb7 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -442,7 +442,7 @@ def show_obj_slices(
         fig, axs = show_2d(
             objs,
             title=titles,
-            cmap=config.get("viz.phase_cmap"),
+            cmap=kwargs.pop("cmap", config.get("viz.phase_cmap")),
             norm=norm,
             cbar=cbar,
             scalebar=scalebars,

From 9b1aec9d4b551d4d98d6e2cfbde3616f2e65ed49 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 13:38:30 -0700
Subject: [PATCH 39/59] fixing linter errors in ptycho viz

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

diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index 3b790bcb7..0c98f9c1a 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -1,9 +1,10 @@
 import warnings
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
+import torch
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from scipy.signal.windows import tukey
 
@@ -955,8 +956,8 @@ def show_scan_positions(
 
     def show_updated_scan_positions(
         self,
-        scan_positions_px: np.ndarray | None = None,
-        initial_scan_positions_px: np.ndarray | None = None,
+        scan_positions_px: np.ndarray | torch.Tensor | None = None,
+        initial_scan_positions_px: np.ndarray | torch.Tensor | None = None,
         scale_arrows: float = 1.0,
         plot_arrow_freq: int | None = None,
         plot_cropped_rotated_fov: bool = True,
@@ -1011,9 +1012,9 @@ def show_updated_scan_positions(
         )
 
         if scan_positions_px.ndim == 3:
-            scan_positions_px = scan_positions_px.mean(axis=0)
+            scan_positions_px = cast(np.ndarray, scan_positions_px.mean(axis=0))
         if initial_scan_positions_px.ndim == 3:
-            initial_scan_positions_px = initial_scan_positions_px.mean(axis=0)
+            initial_scan_positions_px = cast(np.ndarray, initial_scan_positions_px.mean(axis=0))
 
         if scan_positions_px.ndim != 2 or scan_positions_px.shape[-1] != 2:
             raise ValueError(

From e2f084de82f0a182b1a5e22f7f477b4ecaeb6d76 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 15:59:54 -0700
Subject: [PATCH 40/59] fixing bug of constraints not being passed to multiple
 gpus

---
 src/quantem/core/ml/blocks.py                 |  4 +--
 .../diffractive_imaging/object_models.py      |  1 +
 .../diffractive_imaging/ptychography.py       | 27 +++++++++++++++++--
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/quantem/core/ml/blocks.py b/src/quantem/core/ml/blocks.py
index 4465d757d..c243d2c99 100644
--- a/src/quantem/core/ml/blocks.py
+++ b/src/quantem/core/ml/blocks.py
@@ -239,7 +239,7 @@ def _interpolate(self, x: torch.Tensor) -> torch.Tensor:
         return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.method == "transpose":
+        if getattr(self, "method", "transpose") == "transpose":
             if self.scale_factor == 2:
                 assert self.upsample2x is not None
                 x = self.upsample2x(x)
@@ -532,7 +532,7 @@ def _interpolate(self, x: torch.Tensor) -> torch.Tensor:
         return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.method == "transpose":
+        if getattr(self, "method", "transpose") == "transpose":
             assert self.upsample is not None
             x = self.upsample(x)
         else:
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 54ed48969..3ff7099ec 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -511,6 +511,7 @@ def _apply_hard_potential(
             else:
                 offset = obj.min()
             offset = offset.detach()
+            # offset = max(0, offset.detach()) # TODO figure out instability
             offset = offset * c.fix_potential_baseline_factor
         else:
             offset = 0
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 5cbe53542..67727613e 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -69,14 +69,21 @@ def _ddp_ptycho_worker(
     if rank == 0:
         obj_opt = ptycho.optimizers.get("object")
         probe_opt = ptycho.optimizers.get("probe")
+        dset_opt = ptycho.optimizers.get("dataset")
         torch.save(
             {
                 "obj_state": {k: v.cpu() for k, v in ptycho.obj_model.state_dict().items()},
                 "probe_state": {k: v.cpu() for k, v in ptycho.probe_model.state_dict().items()},
+                # Dataset learnable params (scan positions / descan) are optimized and all-reduced
+                # in the workers; ship them back so the main process keeps the refinement.
+                "dset_scan_positions_px": ptycho.dset._scan_positions_px.detach().cpu(),
+                "dset_descan_shifts": ptycho.dset._descan_shifts.detach().cpu(),
                 "obj_optimizer_params": ptycho.obj_model._optimizer_params,
                 "probe_optimizer_params": ptycho.probe_model._optimizer_params,
+                "dset_optimizer_params": ptycho.dset._optimizer_params,
                 "obj_optimizer_state": obj_opt.state_dict() if obj_opt is not None else None,
                 "probe_optimizer_state": probe_opt.state_dict() if probe_opt is not None else None,
+                "dset_optimizer_state": dset_opt.state_dict() if dset_opt is not None else None,
                 "iter_losses": ptycho._iter_losses,
                 "iter_val_losses": ptycho._iter_val_losses,
                 "iter_lrs": ptycho._iter_lrs,
@@ -85,7 +92,7 @@ def _ddp_ptycho_worker(
             result_path,
         )
 
-    # Synchronize before teardown so every rank finishes 
+    # Synchronize before teardown so every rank finishes
     if dist.is_available() and dist.is_initialized():
         if torch.cuda.is_available():
             dist.barrier(device_ids=[device_id])
@@ -261,6 +268,9 @@ def reconstruct(
         """
         self._check_preprocessed()
 
+        if constraints:
+            self.constraints = constraints
+
         # Determine effective device list: explicit arg takes priority, else fall back to stored.
         devices_to_use = (
             device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
@@ -564,11 +574,20 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
         self.probe_model.load_state_dict(result["probe_state"])
         self.to(restore_device)
 
+        # --- dataset learnable params (scan positions / descan), refined in the workers ---
+        dset_positions = result.get("dset_scan_positions_px")
+        if dset_positions is not None:
+            self.dset._scan_positions_px.data = dset_positions.to(restore_device)
+        dset_descan = result.get("dset_descan_shifts")
+        if dset_descan is not None:
+            self.dset._descan_shifts.data = dset_descan.to(restore_device)
+
         # --- restore optimizer params (worker may have set/changed them) so that future
         #     spawns (e.g. reset=True without optimizer_params) can re-init the optimizer ---
         for model, key in (
             (self.obj_model, "obj_optimizer_params"),
             (self.probe_model, "probe_optimizer_params"),
+            (self.dset, "dset_optimizer_params"),
         ):
             saved = result.get(key)
             if saved is not None:
@@ -579,7 +598,11 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
         self.set_optimizers()
 
         # --- optimizer states (params and device must be set before loading) ---
-        for name, key in (("object", "obj_optimizer_state"), ("probe", "probe_optimizer_state")):
+        for name, key in (
+            ("object", "obj_optimizer_state"),
+            ("probe", "probe_optimizer_state"),
+            ("dataset", "dset_optimizer_state"),
+        ):
             opt_state = result.get(key)
             opt = self.optimizers.get(name)
             if opt_state is not None and opt is not None:

From 5027a722e2441e0c6be194408163f8cc95206022 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 16:08:33 -0700
Subject: [PATCH 41/59] fixing bug in poisson loss scaling for batch sizes

---
 src/quantem/diffractive_imaging/ptycho_losses.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/quantem/diffractive_imaging/ptycho_losses.py b/src/quantem/diffractive_imaging/ptycho_losses.py
index d1ac180ea..cd45970ae 100644
--- a/src/quantem/diffractive_imaging/ptycho_losses.py
+++ b/src/quantem/diffractive_imaging/ptycho_losses.py
@@ -62,10 +62,13 @@ def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.
 
 
 class Poisson(DataCriterion):
+    """Poisson negative log-likelihood in intensity space (up to a pred-independent constant)."""
+
     target_space: TargetSpace = "intensity"
 
     def __call__(self, preds: torch.Tensor, targets: torch.Tensor, n: int) -> torch.Tensor:
-        return torch.sum(preds - targets * torch.log(preds + 1e-6))
+        nll = torch.sum(preds - targets * torch.log(preds + 1e-6))
+        return nll / _global_scale(preds, n)
 
 
 class AmplitudeSmoothL1(DataCriterion):
@@ -113,6 +116,12 @@ class AmplitudeS3IM(DataCriterion):
     this captures structural relationships a per-pixel loss misses. It is used as an auxiliary
     term on top of an MSE term (both mean-reduced here, so ``lambda`` ~ O(1) balances them). The
     SSIM passes make this notably more expensive than L2 — keep ``repeats`` modest.
+
+    Note: both terms are **mean**-reduced. A mean is already batch-size independent (no
+    ``_global_scale`` rescale needed), so this criterion is well-behaved across batch sizes and
+    multi-GPU. It does, however, sit on a different absolute scale than the sum-based criteria
+    (``L2``/``L1``/Poisson, which rescale to a full-dataset sum), so learning rates do **not**
+    transfer between ``s3im_amplitude`` and those losses — retune the LR when switching.
     """
 
     target_space: TargetSpace = "amplitude"

From 4045da2647b2c86ada096c32296d4118eeda7037 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 16:22:04 -0700
Subject: [PATCH 42/59] cleaning up some type hints

---
 src/quantem/diffractive_imaging/probe_models.py      | 2 +-
 src/quantem/diffractive_imaging/ptychography.py      | 4 +---
 src/quantem/diffractive_imaging/ptychography_base.py | 1 -
 src/quantem/diffractive_imaging/ptychography_lite.py | 6 +++---
 4 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index 9cae9f7b0..b89aa1fe2 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -236,7 +236,7 @@ def probe_params(self) -> dict[str, Any]:
         return self._probe_params
 
     @probe_params.setter
-    def probe_params(self, params: dict[str, Any] = {}):
+    def probe_params(self, params: dict[str, Any]):
         validate_dict_keys(
             params,
             [*self.DEFAULT_PROBE_PARAMS.keys(), *POLAR_SYMBOLS, *POLAR_ALIASES.keys()],
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 67727613e..812a50449 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -234,7 +234,7 @@ def reconstruct(
         batch_size: int | None = None,
         store_snapshots: bool | None = None,
         store_snapshots_every: int | None = None,
-        device: Literal["cpu", "gpu"] | int | list[int] | None = None,
+        device: str | int | list[int] | None = None,
         autograd: bool = True,
         loss_type: "str | DataCriterion" = "l2_amplitude",
         num_workers: int = 0,
@@ -422,7 +422,6 @@ def _reconstruct_inner(
 
                 batch_consistency_loss, targets = self.error_estimate(
                     pred_intensities,
-                    batch_indices,
                     targets=targets,
                     global_n=global_n,
                 )
@@ -477,7 +476,6 @@ def _reconstruct_inner(
                         pred_intensities = self.detector_model.forward(overlap)
                         batch_val_loss, _ = self.error_estimate(
                             pred_intensities,
-                            batch_indices,
                             targets=targets,
                             global_n=global_n,
                         )
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 8daacdfb9..1c830f127 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -1103,7 +1103,6 @@ def forward_operator(
     def error_estimate(
         self,
         pred_intensities: torch.Tensor,
-        batch_indices: np.ndarray,
         targets: torch.Tensor,
         global_n: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 41ce2c780..eea9fc71f 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -50,7 +50,7 @@ def from_dataset(
         log_prefix: str = "",
         log_images_every: int = 10,
         log_probe_images: bool = False,
-        device: Literal["cpu", "gpu"] = "cpu",
+        device: str | int = "cpu",
         verbose: int | bool = True,
         rng: np.random.Generator | int | None = None,
     ) -> Self:
@@ -180,7 +180,7 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         new_optimizers: bool = False,  # not sure what the default should be
         constraints: dict[str, Any] | None = None,
         store_iterations_every: int | None = None,
-        device: "Literal['cpu', 'gpu'] | int | list[int] | None" = None,
+        device: "str | int | list[int] | None" = None,
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose
@@ -427,7 +427,7 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         new_optimizers: bool = False,  # not sure what the default should be
         constraints: dict[str, Any] | None = None,
         store_iterations_every: int | None = None,
-        device: Literal["cpu", "gpu"] | int | list[int] | None = None,
+        device: str | int | list[int] | None = None,
         verbose: int | bool = True,
     ) -> Self:
         self.verbose = verbose

From a8d7a53c555708bdc2bdabd7e5338b36118dc85d Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 18:25:57 -0700
Subject: [PATCH 43/59] descan_shifts_zero name change, and adding shrinkage vs
 clamp option to positivity

---
 .../diffractive_imaging/dataset_models.py     | 11 ++-
 .../diffractive_imaging/object_models.py      | 95 +++++++++++++++----
 .../diffractive_imaging/probe_models.py       |  4 +-
 .../diffractive_imaging/ptychography.py       |  2 +
 4 files changed, 88 insertions(+), 24 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 8c29a43fd..689430a0d 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -52,10 +52,11 @@ class Raster(Constraints):
 
         Attributes
         ----------
-        descan_shifts_constant : bool, default ``False``
+        descan_shifts_zero : bool, default ``False``
             Forces all descan shifts to zero after each update. Useful when you
             want to keep the descan optimizer in the parameter group but freeze
-            its effect.
+            its effect. (Distinct from ``learn_descan=False``, which holds descan
+            at its fitted value rather than zeroing it.)
         center_scan_positions : bool, default ``False``
             Shifts all scan positions uniformly so their mean sits at the object
             center after each update. Prevents the reconstruction from
@@ -72,7 +73,7 @@ class Raster(Constraints):
         """
 
         # hard constraints
-        descan_shifts_constant: bool = False
+        descan_shifts_zero: bool = False
         center_scan_positions: bool = False
         clip_scan_positions: bool = True
         # soft constraints
@@ -81,7 +82,7 @@ class Raster(Constraints):
 
         soft_constraint_keys = ["descan_tv_weight"]
         hard_constraint_keys = [
-            "descan_shifts_constant",
+            "descan_shifts_zero",
             "center_scan_positions",
             "clip_scan_positions",
         ]
@@ -782,7 +783,7 @@ def apply_descan_constraints(
         self,
         descan: torch.Tensor,
     ) -> torch.Tensor:
-        if self.constraints.descan_shifts_constant:
+        if self.constraints.descan_shifts_zero:
             descan = torch.zeros_like(descan)
         return descan
 
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 3ff7099ec..fd5a54862 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -114,6 +114,14 @@ class Raster(Constraints):
 
         # hard constraints
         positivity: bool = True
+        # positivity_mode selects HOW positivity is enforced for obj_type="potential":
+        #   "clamp"  -- straight-through clamp(obj, 0) in the forward (default). Cheap display
+        #               gauge: it does NOT move the _obj parameter, only how it is shown/used.
+        #   "shrink" -- proximal per-slice shrinkage on the _obj parameter post-step: subtract
+        #               fix_potential_baseline_factor * (per-slice background) then clamp >= 0.
+        #               Per-slice keeps it on the diffraction-invariant gauge (loss-neutral for
+        #               multislice) and self-limits as the background -> 0. ObjectPixelated only.
+        positivity_mode: Literal["clamp", "shrink"] = "clamp"
         fix_potential_baseline: bool = False
         fix_potential_baseline_factor: float = 1.0
         identical_slices: bool = False
@@ -132,6 +140,7 @@ class Raster(Constraints):
         soft_constraint_keys = ["tv_weight_z", "tv_weight_xy", "surface_zero_weight"]
         hard_constraint_keys = [
             "positivity",
+            "positivity_mode",
             "fix_potential_baseline",
             "fix_potential_baseline_factor",
             "identical_slices",
@@ -391,6 +400,15 @@ def forward(self, patch_indices: torch.Tensor, /):
     def reset(self):
         raise NotImplementedError()
 
+    def project_parameters(self) -> None:
+        """In-place hard projection of the underlying parameters after an optimizer step.
+
+        No-op by default. ``ObjectPixelated`` overrides this to enforce
+        ``positivity_mode="shrink"`` (proximal per-slice background shrinkage on the potential
+        grid). Called once per optimizer step from the reconstruction loop.
+        """
+        return
+
     @abstractmethod
     def _initialize_obj(
         self,
@@ -501,24 +519,38 @@ def _apply_hard_potential(
         c: PtychoObjConstraintParams.Raster,
         mask: torch.Tensor | None,
     ) -> torch.Tensor:
-        if c.fix_potential_baseline:
-            if mask is not None:
-                background = mask < 0.5 * mask.max()
-                if background.any():
-                    offset = obj[background].mean()
-                else:
-                    offset = obj.min()
-            else:
-                offset = obj.min()
-            offset = offset.detach()
-            # offset = max(0, offset.detach()) # TODO figure out instability
-            offset = offset * c.fix_potential_baseline_factor
-        else:
-            offset = 0
-
+        # "shrink" manages the _obj parameter directly in project_parameters() (post-step), so the
+        # forward just passes the (already non-negative) parameter through.
+        if c.positivity_mode == "shrink":
+            return obj
+        offset = self._potential_baseline_offset(obj, c, mask)
+        obj = obj - offset
+        # "clamp" clamps here; the apply_hard_constraints wrapper makes it a straight-through op
+        # (the forward/display is non-negative but the _obj parameter is left untouched).
         if c.positivity:
-            return torch.clamp(obj - offset, min=0.0)
-        return obj - offset
+            return torch.clamp(obj, min=0.0)
+        return obj
+
+    def _potential_baseline_offset(
+        self,
+        obj: torch.Tensor,
+        c: PtychoObjConstraintParams.Raster,
+        mask: torch.Tensor | None,
+    ) -> torch.Tensor | float:
+        """Background offset subtracted by ``fix_potential_baseline`` (``0`` when disabled).
+
+        Estimated from the FOV-mask background mean (else the global min), detached, and scaled by
+        ``fix_potential_baseline_factor``.
+        """
+        if not c.fix_potential_baseline:
+            return 0.0
+        # mask is an empty tensor (not None) when no FOV mask is set, so guard on numel().
+        if mask is not None and mask.numel() and (mask < 0.5 * mask.max()).any():
+            offset = obj[mask < 0.5 * mask.max()].mean()
+        else:
+            offset = obj.min()
+        offset = offset.detach()
+        return offset * c.fix_potential_baseline_factor
 
     def _apply_shared_hard(
         self,
@@ -895,6 +927,35 @@ def params(self) -> list[nn.Parameter]:
         """optimization parameters"""
         return [self._obj]
 
+    def project_parameters(self) -> None:
+        """Post-step proximal shrinkage of the potential grid parameter, for
+        ``positivity_mode="shrink"`` (``obj_type="potential"`` with ``positivity=True``); a no-op
+        otherwise.
+
+        Subtracts a per-slice background offset (``fix_potential_baseline_factor * per-slice
+        background``) from ``_obj`` then clamps ``>= 0``. Unlike the straight-through ``clamp`` mode
+        this moves ``_obj`` itself, so the reconstruction (not just the display) is shrunk toward a
+        zero background. Per-slice keeps it on the diffraction-invariant gauge (loss-neutral for
+        multislice), and the offset shrinks with the background so it self-limits at zero.
+        """
+        c = self.constraints
+        if not (self.obj_type == "potential" and c.positivity and c.positivity_mode == "shrink"):
+            return
+        with torch.no_grad():
+            bg = self._per_slice_background(self._obj, self.mask)
+            offset = (c.fix_potential_baseline_factor * bg.clamp_min(0.0)).view(-1, 1, 1)
+            self._obj.data = (self._obj.data - offset).clamp_min(0.0)
+
+    def _per_slice_background(self, obj: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """Per-slice background level, shape ``[num_slices]``: the FOV-mask background mean when a
+        mask is set, else a robust per-slice 10th percentile."""
+        s = obj.shape[0]
+        flat = obj.reshape(s, -1)
+        if mask is not None and mask.numel() and (mask < 0.5 * mask.max()).any():
+            bg = (mask < 0.5 * mask.max()).reshape(s, -1).to(obj.dtype)
+            return (flat * bg).sum(1) / bg.sum(1).clamp_min(1.0)
+        return torch.quantile(flat, 0.1, dim=1)
+
     @property
     def initial_obj(self):
         return self._initial_obj
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index b89aa1fe2..d5f8eea92 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -873,7 +873,7 @@ def vacuum_probe_intensity(self, vp: np.ndarray | torch.Tensor | Dataset4dstem |
         elif isinstance(vp, np.ndarray):
             vp2 = vp.astype(config.get("dtype_real"))
         elif isinstance(vp, (Dataset4dstem, Dataset2d)):
-            vp2 = vp.array
+            vp2 = cast(np.ndarray, vp.array) # TODO when finished Dataset->torch fix here
         elif isinstance(vp, torch.Tensor):
             vp2 = vp.cpu().detach().numpy()
         else:
@@ -1067,7 +1067,7 @@ def vacuum_probe_intensity(self, vp: np.ndarray | Dataset4dstem | None):
         elif isinstance(vp, np.ndarray):
             vp2 = vp.astype(config.get("dtype_real"))
         elif isinstance(vp, (Dataset4dstem, Dataset2d)):
-            vp2 = vp.array
+            vp2 = cast(np.ndarray, vp.array) # TODO when finished Dataset->torch fix here
         else:
             raise NotImplementedError(f"Unknown vacuum probe type: {type(vp)}")
 
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 812a50449..b4a25cc7b 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -441,6 +441,8 @@ def _reconstruct_inner(
                 if _dist_world_size > 1:
                     self._all_reduce_gradients()
                 self.step_optimizers()
+                # Post-step parameter projection (only for positivity_mode="shrink")
+                self.obj_model.project_parameters()
                 consistency_loss += batch_consistency_loss.item()
                 total_loss += batch_loss.item()
 

From 50f4948cae015095f72d0ad04208e29d0450f9d3 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 10 Jun 2026 19:24:47 -0700
Subject: [PATCH 44/59] fixing docstring

---
 src/quantem/diffractive_imaging/object_models.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index fd5a54862..0301be886 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -75,6 +75,13 @@ class Raster(Constraints):
             Only consulted when ``obj_type="potential"``; for ``"complex"`` /
             ``"pure_phase"`` the amplitude is clamped to ``[0, 1]`` (or fixed to 1)
             regardless of this flag.
+        positivity_mode: Literal["clamp", "shrink"], default ``"clamp"``
+            How to enforce positivity. "clamp" clamps the object to be non-negative after each 
+            update, does not move the parameter, only how it is shown/used.
+            "shrink" subtracts a background offset from the object so background regions sit at 
+            zero, is applied to the parameter after the update step. 
+            If an FOV mask is set the offset is the mean of the background 
+            (``mask < 0.5 * mask.max()``); otherwise it's ``obj.min()``.
         fix_potential_baseline : bool, default ``False``
             ``obj_type="potential"`` only. Subtracts an offset from the object so
             background regions sit at zero. If an FOV mask is set the offset is
@@ -114,13 +121,6 @@ class Raster(Constraints):
 
         # hard constraints
         positivity: bool = True
-        # positivity_mode selects HOW positivity is enforced for obj_type="potential":
-        #   "clamp"  -- straight-through clamp(obj, 0) in the forward (default). Cheap display
-        #               gauge: it does NOT move the _obj parameter, only how it is shown/used.
-        #   "shrink" -- proximal per-slice shrinkage on the _obj parameter post-step: subtract
-        #               fix_potential_baseline_factor * (per-slice background) then clamp >= 0.
-        #               Per-slice keeps it on the diffraction-invariant gauge (loss-neutral for
-        #               multislice) and self-limits as the background -> 0. ObjectPixelated only.
         positivity_mode: Literal["clamp", "shrink"] = "clamp"
         fix_potential_baseline: bool = False
         fix_potential_baseline_factor: float = 1.0

From 80a7b191c172132e6ddfe4ffa21d484bb354b338 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Thu, 11 Jun 2026 16:18:43 -0700
Subject: [PATCH 45/59] converting to degrees

---
 .../direct_ptycho_utils.py                    |  7 +++-
 .../direct_ptychography.py                    | 37 +++++++++++++------
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/quantem/diffractive_imaging/direct_ptycho_utils.py b/src/quantem/diffractive_imaging/direct_ptycho_utils.py
index fcc766011..fb40f9a07 100644
--- a/src/quantem/diffractive_imaging/direct_ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/direct_ptycho_utils.py
@@ -495,7 +495,10 @@ def fit_aberrations_from_shifts(
     gpts: tuple[int, int],
     sampling: tuple[float, float],
 ) -> dict[str, float]:
-    """ """
+    """Fit low-order aberrations from lateral shifts.
+
+    Returns ``rotation_angle`` in degrees.
+    """
     device = shifts_ang.device
 
     # Get spatial frequencies at BF positions
@@ -534,7 +537,7 @@ def fit_aberrations_from_shifts(
         "C10": C10.item(),
         "C12": C12.item(),
         "phi12": phi12.item(),
-        "rotation_angle": rotation_rad.item(),
+        "rotation_angle": torch.rad2deg(rotation_rad).item(),
     }
 
 
diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index ecbeece0d..87538fe34 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -54,6 +54,12 @@
 optuna.logging.set_verbosity(optuna.logging.WARNING)
 
 
+def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None:
+    if rotation_angle is None:
+        return None
+    return math.radians(float(rotation_angle))
+
+
 @dataclass
 class OptimizationParameter:
     low: float
@@ -156,13 +162,13 @@ def add(name: str, value):
             if self.initial_aberrations:
                 add("initial_aberrations", self.initial_aberrations)
             if self.initial_rotation_angle is not None:
-                add("initial_rotation_angle", self.initial_rotation_angle)
+                add("initial_rotation_angle_deg", self.initial_rotation_angle)
 
         elif which == "optimized":
             if self.optimized_aberrations:
                 add("optimized_aberrations", self.optimized_aberrations)
             if self.optimized_rotation_angle is not None:
-                add("optimized_rotation_angle", self.optimized_rotation_angle)
+                add("optimized_rotation_angle_deg", self.optimized_rotation_angle)
 
         elif which == "current":
             current_abers = self.current_aberrations(override_aberration_coefs)
@@ -171,17 +177,17 @@ def add(name: str, value):
             if current_abers:
                 add("current_aberrations", current_abers)
             if current_rot is not None:
-                add("current_rotation_angle", current_rot)
+                add("current_rotation_angle_deg", current_rot)
 
         elif which == "all":
             if self.initial_aberrations:
                 add("initial_aberrations", self.initial_aberrations)
             if self.initial_rotation_angle is not None:
-                add("initial_rotation_angle", self.initial_rotation_angle)
+                add("initial_rotation_angle_deg", self.initial_rotation_angle)
             if self.optimized_aberrations:
                 add("optimized_aberrations", self.optimized_aberrations)
             if self.optimized_rotation_angle is not None:
-                add("optimized_rotation_angle", self.optimized_rotation_angle)
+                add("optimized_rotation_angle_deg", self.optimized_rotation_angle)
 
         else:
             raise ValueError(
@@ -335,7 +341,7 @@ def from_dataset4d(
 
         if rotation_angle is None:
             origin.estimate_detector_rotation()
-            rotation_angle = origin.detector_rotation_deg / 180 * math.pi
+            rotation_angle = origin.detector_rotation_deg
 
         # shift to origin
         origin.shift_origin_to(
@@ -494,6 +500,7 @@ def bf_mask(self, value: torch.Tensor):
 
     @property
     def rotation_angle(self) -> float:
+        """Current detector rotation angle in degrees."""
         return self.hyperparameter_state.current_rotation_angle()
 
     @property
@@ -679,7 +686,7 @@ def reconstruct(
         upsampling_factor : int, optional
             Factor by which to upsample the reconstruction
         override_rotation_angle : float, optional
-            Rotation angle for coordinate system
+            Rotation angle for coordinate system, in degrees
         max_batch_size : int, optional
             Maximum batch size for processing
         deconvolution_kernel : str, one of ['ssb', 'obf', 'mf','prlx','icom']
@@ -742,7 +749,10 @@ def reconstruct(
 
         # Get k-space grid
         kxa, kya = spatial_frequencies(
-            self.gpts, self.sampling, rotation_angle=rotation_angle, device=self.device
+            self.gpts,
+            self.sampling,
+            rotation_angle=_rotation_degrees_to_radians(rotation_angle),
+            device=self.device,
         )
         k, phi = polar_coordinates(kxa, kya)
 
@@ -917,7 +927,7 @@ def optimize_hyperparameters(
         aberration_coefs : dict[str, float|OptimizationParameter]
             Dict of aberration names to either fixed values or optimization ranges.
         rotation_angle : float|OptimizationParameter
-            Fixed rotation or optimization range.
+            Fixed rotation or optimization range, in degrees.
         n_trials : int
             Number of Optuna trials.
         sampler : optuna.samplers.BaseSampler, optional
@@ -1107,7 +1117,10 @@ def _return_lateral_shifts(
     ):
         # Get initial shifts
         kxa, kya = spatial_frequencies(
-            self.gpts, self.sampling, rotation_angle=rotation_angle, device=self.device
+            self.gpts,
+            self.sampling,
+            rotation_angle=_rotation_degrees_to_radians(rotation_angle),
+            device=self.device,
         )
         k, phi = polar_coordinates(kxa, kya)
 
@@ -1335,7 +1348,7 @@ def _fit_hyperparameters_least_squares_inner(
         kxa, kya = spatial_frequencies(
             self.gpts,
             self.sampling,
-            rotation_angle=rotation_angle,
+            rotation_angle=_rotation_degrees_to_radians(rotation_angle),
             device=device,
         )
 
@@ -1511,7 +1524,7 @@ def fit_hyperparameters_least_squares(
 
         Args:
             aberration_coefs: Initial aberration coefficients to deconvolve
-            rotation_angle: Rotation angle for basis functions
+            rotation_angle: Rotation angle for basis functions, in degrees
             cartesian_basis: Aberration basis to fit. Can be:
                 - str: preset name like "low_order"
                 - list[str]: explicit list like ["C10", "C12_a", "C12_b", ...]

From 61ca85200d73ab4fd185e92150c6a91cf6df9bae Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Thu, 11 Jun 2026 18:08:58 -0700
Subject: [PATCH 46/59] correcting stig rotation

---
 .../direct_ptychography.py                    | 41 ++++++++++++++-----
 1 file changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index 87538fe34..9ebc6f4a7 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -60,6 +60,19 @@ def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None:
     return math.radians(float(rotation_angle))
 
 
+def _uses_aberration_orientation_convention(key: str) -> bool:
+    return key.startswith("phi") or key.endswith("_b")
+
+
+def _aberration_coefs_to_direct_convention(
+    aberration_coefs: dict[str, float | torch.Tensor],
+) -> dict[str, float | torch.Tensor]:
+    return {
+        key: -value if _uses_aberration_orientation_convention(key) else value
+        for key, value in aberration_coefs.items()
+    }
+
+
 @dataclass
 class OptimizationParameter:
     low: float
@@ -726,6 +739,8 @@ def reconstruct(
             aberration_coefs = state.current_aberrations(override_aberration_coefs)
             rotation_angle = state.current_rotation_angle(override_rotation_angle)
 
+        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
+
         if upsampling_factor is None:
             upsampling_factor = 1
         upsampling_factor = math.ceil(upsampling_factor)
@@ -761,7 +776,7 @@ def reconstruct(
             dx, dy = aberration_surface_cartesian_gradients(
                 k * self.wavelength,
                 phi,
-                aberration_coefs=aberration_coefs,
+                aberration_coefs=direct_aberration_coefs,
             )
             grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1)
 
@@ -770,7 +785,7 @@ def reconstruct(
                     q * self.wavelength,
                     theta,
                     self.wavelength,
-                    aberration_coefs=aberration_coefs,
+                    aberration_coefs=direct_aberration_coefs,
                 )
                 sign_sin_chi_q = torch.sign(torch.sin(chi_q))
             else:
@@ -786,7 +801,7 @@ def reconstruct(
             self.semiangle_cutoff,
             self.angular_sampling,
             self.wavelength,
-            aberration_coefs=aberration_coefs,
+            aberration_coefs=direct_aberration_coefs,
         )
         BF_weights = cmplx_probe_k[bf_mask].abs().square().sum()
 
@@ -830,7 +845,7 @@ def reconstruct(
                 cmplx_probe_k,
                 grad_k,
                 sign_sin_chi_q,
-                aberration_coefs,
+                direct_aberration_coefs,
                 batch_idx,
             )
             if power is None:
@@ -1115,6 +1130,8 @@ def _return_lateral_shifts(
         aberration_coefs,
         bf_mask,
     ):
+        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
+
         # Get initial shifts
         kxa, kya = spatial_frequencies(
             self.gpts,
@@ -1127,7 +1144,7 @@ def _return_lateral_shifts(
         dx, dy = aberration_surface_cartesian_gradients(
             k * self.wavelength,
             phi,
-            aberration_coefs=aberration_coefs,
+            aberration_coefs=direct_aberration_coefs,
         )
         grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1)
         lateral_shifts = grad_k / 2 / np.pi
@@ -1269,8 +1286,8 @@ def fit_hyperparameters_cross_correlation(
 
         self.corrected_stack = vbf_stack
 
-        fitted_aberration_coefs = fit_results.copy()
-        fitted_rotation_angle = fitted_aberration_coefs.pop("rotation_angle", None)
+        fitted_rotation_angle = fit_results.pop("rotation_angle", None)
+        fitted_aberration_coefs = _aberration_coefs_to_direct_convention(fit_results)
 
         state.optimized_aberrations = validate_aberration_coefficients(fitted_aberration_coefs)
         state.optimized_rotation_angle = fitted_rotation_angle
@@ -1330,6 +1347,7 @@ def _fit_hyperparameters_least_squares_inner(
 
         device = self.device
         wavelength = self.wavelength
+        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
 
         # ---------------------------------------------------------
         # Select strongest spatial frequencies
@@ -1381,9 +1399,9 @@ def _fit_hyperparameters_least_squares_inner(
             km * wavelength, phim, self.semiangle_cutoff, self.angular_sampling, soft_edges=True
         )
 
-        chi0 = aberration_surface(k0 * wavelength, phi0, wavelength, aberration_coefs)
-        chip = aberration_surface(kp * wavelength, phip, wavelength, aberration_coefs)
-        chim = aberration_surface(km * wavelength, phim, wavelength, aberration_coefs)
+        chi0 = aberration_surface(k0 * wavelength, phi0, wavelength, direct_aberration_coefs)
+        chip = aberration_surface(kp * wavelength, phip, wavelength, direct_aberration_coefs)
+        chim = aberration_surface(km * wavelength, phim, wavelength, direct_aberration_coefs)
 
         # ---------------------------------------------------------
         # Overlap-only masks
@@ -1502,7 +1520,8 @@ def ap_to_mask(ap, eps=1e-6):
 
         delta_cartesian = {name: sol[i] for i, name in enumerate(cartesian_basis)}
 
-        return merge_aberration_coefficients(aberration_coefs, delta_cartesian)
+        direct_merged = merge_aberration_coefficients(direct_aberration_coefs, delta_cartesian)
+        return _aberration_coefs_to_direct_convention(direct_merged)
 
     def fit_hyperparameters_least_squares(
         self,

From 44481bcc91db7ea0173ab1422f3b74ddf5e262fc Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Thu, 11 Jun 2026 18:14:57 -0700
Subject: [PATCH 47/59] basic plotting

---
 .../direct_ptychography.py                    | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index 9ebc6f4a7..60e9700b4 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -19,6 +19,7 @@
     validate_int,
     validate_tensor,
 )
+from quantem.core.visualization import show_2d
 from quantem.diffractive_imaging.complex_probe import (
     aberration_surface,
     aberration_surface_cartesian_basis,
@@ -925,6 +926,50 @@ def obj(self) -> np.ndarray:
         obj = to_numpy(self.corrected_bf)
         return obj
 
+    def visualize(
+        self,
+        return_fig: bool = False,
+        **kwargs,
+    ):
+        """
+        Show the reconstructed object and its Hann-windowed Fourier transform.
+
+        Parameters
+        ----------
+        cbar : bool, optional
+            Whether to show colorbars, by default True.
+        return_fig : bool, optional
+            If True, return ``(fig, axs)``.
+        fft_norm : str | dict, optional
+            Normalization passed to ``show_2d`` for the object FFT.
+        **kwargs
+            Additional arguments passed to ``show_2d``.
+        """
+        if self.corrected_bf is None:
+            raise RuntimeError("Run reconstruct() before visualize().")
+
+        obj = self.obj
+        window = np.hanning(obj.shape[-2])[:, None] * np.hanning(obj.shape[-1])[None, :]
+        obj_fft = np.fft.fftshift(np.fft.fft2(obj * window))
+
+        obj_scalebar = {"sampling": self.scan_sampling[0], "units": "Å"}
+        fft_sampling = 1 / (self.scan_sampling[0] * obj.shape[-2])
+        fft_scalebar = {"sampling": fft_sampling, "units": r"$\mathrm{A^{-1}}$"}
+
+        fig, axs = show_2d(
+            [obj, np.abs(obj_fft)],
+            title=["Object", "Object FFT"],
+            scalebar=[obj_scalebar, fft_scalebar],
+            cmap=["magma", "magma"],
+            norm=[None, None],
+            **kwargs,
+        )
+        axs[1].set_aspect(obj.shape[-1] / obj.shape[-2])
+
+        if return_fig:
+            return fig, axs
+        return None
+
     def optimize_hyperparameters(
         self,
         aberration_coefs: dict[str, float | OptimizationParameter] | None = None,

From 5fa1c6119753d009c245a35bdb53e5f7b3ed3a8b Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 11 Jun 2026 18:16:40 -0700
Subject: [PATCH 48/59] fixing bug in com determination when DP sum = 0

---
 src/quantem/diffractive_imaging/dataset_models.py | 4 ++++
 src/quantem/diffractive_imaging/ptycho_utils.py   | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 689430a0d..a2776fdec 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1264,8 +1264,12 @@ def _set_intensities_com(
             com_measured_c = np.sum(intensities_mask * kcm[None, None], axis=(-2, -1))
 
             intensities_sum = np.sum(intensities_mask, axis=(-2, -1))
+            intensities_mask = intensities_sum == 0
+            intensities_sum[intensities_mask] = 1
             com_measured_r /= intensities_sum
             com_measured_c /= intensities_sum
+            com_measured_r[intensities_mask] = np.nan
+            com_measured_c[intensities_mask] = np.nan
 
         else:
             shape_r, shape_c = intensities.shape[:2]
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index f3c097261..29a690475 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -432,7 +432,9 @@ def fit_origin(
         qr0_meas_masked = qr0_meas[mask]
         qc0_meas_masked = qc0_meas[mask]
         mask1D = mask.reshape(1, np.prod(shape))
-        rc_masked = np.vstack((r1D * mask1D, c1D * mask1D))
+        rc_masked = np.vstack((r1D[mask1D], c1D[mask1D]))
+        # old failed for zero-valued DPs
+        # rc_masked = np.vstack((r1D * mask1D, c1D * mask1D))
 
         popt_r, _ = curve_fit(f, rc_masked, qr0_meas_masked)
         popt_c, _ = curve_fit(f, rc_masked, qc0_meas_masked)

From b79ab8d08e4822a01fb4bf7027ab5d06593bae11 Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Fri, 12 Jun 2026 11:43:19 -0700
Subject: [PATCH 49/59] Revert "correcting stig rotation"

This reverts commit 61ca85200d73ab4fd185e92150c6a91cf6df9bae.
---
 .../direct_ptychography.py                    | 41 +++++--------------
 1 file changed, 11 insertions(+), 30 deletions(-)

diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index 60e9700b4..3f3bca7ac 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -61,19 +61,6 @@ def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None:
     return math.radians(float(rotation_angle))
 
 
-def _uses_aberration_orientation_convention(key: str) -> bool:
-    return key.startswith("phi") or key.endswith("_b")
-
-
-def _aberration_coefs_to_direct_convention(
-    aberration_coefs: dict[str, float | torch.Tensor],
-) -> dict[str, float | torch.Tensor]:
-    return {
-        key: -value if _uses_aberration_orientation_convention(key) else value
-        for key, value in aberration_coefs.items()
-    }
-
-
 @dataclass
 class OptimizationParameter:
     low: float
@@ -740,8 +727,6 @@ def reconstruct(
             aberration_coefs = state.current_aberrations(override_aberration_coefs)
             rotation_angle = state.current_rotation_angle(override_rotation_angle)
 
-        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
-
         if upsampling_factor is None:
             upsampling_factor = 1
         upsampling_factor = math.ceil(upsampling_factor)
@@ -777,7 +762,7 @@ def reconstruct(
             dx, dy = aberration_surface_cartesian_gradients(
                 k * self.wavelength,
                 phi,
-                aberration_coefs=direct_aberration_coefs,
+                aberration_coefs=aberration_coefs,
             )
             grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1)
 
@@ -786,7 +771,7 @@ def reconstruct(
                     q * self.wavelength,
                     theta,
                     self.wavelength,
-                    aberration_coefs=direct_aberration_coefs,
+                    aberration_coefs=aberration_coefs,
                 )
                 sign_sin_chi_q = torch.sign(torch.sin(chi_q))
             else:
@@ -802,7 +787,7 @@ def reconstruct(
             self.semiangle_cutoff,
             self.angular_sampling,
             self.wavelength,
-            aberration_coefs=direct_aberration_coefs,
+            aberration_coefs=aberration_coefs,
         )
         BF_weights = cmplx_probe_k[bf_mask].abs().square().sum()
 
@@ -846,7 +831,7 @@ def reconstruct(
                 cmplx_probe_k,
                 grad_k,
                 sign_sin_chi_q,
-                direct_aberration_coefs,
+                aberration_coefs,
                 batch_idx,
             )
             if power is None:
@@ -1175,8 +1160,6 @@ def _return_lateral_shifts(
         aberration_coefs,
         bf_mask,
     ):
-        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
-
         # Get initial shifts
         kxa, kya = spatial_frequencies(
             self.gpts,
@@ -1189,7 +1172,7 @@ def _return_lateral_shifts(
         dx, dy = aberration_surface_cartesian_gradients(
             k * self.wavelength,
             phi,
-            aberration_coefs=direct_aberration_coefs,
+            aberration_coefs=aberration_coefs,
         )
         grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1)
         lateral_shifts = grad_k / 2 / np.pi
@@ -1331,8 +1314,8 @@ def fit_hyperparameters_cross_correlation(
 
         self.corrected_stack = vbf_stack
 
-        fitted_rotation_angle = fit_results.pop("rotation_angle", None)
-        fitted_aberration_coefs = _aberration_coefs_to_direct_convention(fit_results)
+        fitted_aberration_coefs = fit_results.copy()
+        fitted_rotation_angle = fitted_aberration_coefs.pop("rotation_angle", None)
 
         state.optimized_aberrations = validate_aberration_coefficients(fitted_aberration_coefs)
         state.optimized_rotation_angle = fitted_rotation_angle
@@ -1392,7 +1375,6 @@ def _fit_hyperparameters_least_squares_inner(
 
         device = self.device
         wavelength = self.wavelength
-        direct_aberration_coefs = _aberration_coefs_to_direct_convention(aberration_coefs)
 
         # ---------------------------------------------------------
         # Select strongest spatial frequencies
@@ -1444,9 +1426,9 @@ def _fit_hyperparameters_least_squares_inner(
             km * wavelength, phim, self.semiangle_cutoff, self.angular_sampling, soft_edges=True
         )
 
-        chi0 = aberration_surface(k0 * wavelength, phi0, wavelength, direct_aberration_coefs)
-        chip = aberration_surface(kp * wavelength, phip, wavelength, direct_aberration_coefs)
-        chim = aberration_surface(km * wavelength, phim, wavelength, direct_aberration_coefs)
+        chi0 = aberration_surface(k0 * wavelength, phi0, wavelength, aberration_coefs)
+        chip = aberration_surface(kp * wavelength, phip, wavelength, aberration_coefs)
+        chim = aberration_surface(km * wavelength, phim, wavelength, aberration_coefs)
 
         # ---------------------------------------------------------
         # Overlap-only masks
@@ -1565,8 +1547,7 @@ def ap_to_mask(ap, eps=1e-6):
 
         delta_cartesian = {name: sol[i] for i, name in enumerate(cartesian_basis)}
 
-        direct_merged = merge_aberration_coefficients(direct_aberration_coefs, delta_cartesian)
-        return _aberration_coefs_to_direct_convention(direct_merged)
+        return merge_aberration_coefficients(aberration_coefs, delta_cartesian)
 
     def fit_hyperparameters_least_squares(
         self,

From e7cd6521439ffae3971a54c7acd3d62139baf76b Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 12 Jun 2026 16:24:04 -0700
Subject: [PATCH 50/59] adding position masking to iterative ptycho

---
 .../diffractive_imaging/dataset_models.py     | 110 ++++++++++++------
 .../diffractive_imaging/ptycho_utils.py       |   8 +-
 .../diffractive_imaging/ptychography.py       |   6 +-
 .../diffractive_imaging/ptychography_base.py  |   4 +-
 4 files changed, 85 insertions(+), 43 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index a2776fdec..fc1456bb4 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -173,6 +173,7 @@ def __init__(
         )
         self.register_buffer("_last_patch_positions_px", torch.zeros(self.num_gpts, 2))
         self.register_buffer("_detector_mask", torch.ones(*self.roi_shape))
+        self.positions_mask = torch.ones(self.num_gpts, dtype=torch.bool)
         self.detector_mask = detector_mask
         self._constraints = {}
         self._probe_energy = None
@@ -251,7 +252,7 @@ def descan_shifts(self, shifts: torch.Tensor | np.ndarray) -> None:
             shifts,
             name="descan_shifts",
             dtype=getattr(torch, config.get("dtype_real")),
-            shape=(self.num_gpts, 2),
+            shape=(self.num_positions, 2),
         )
         self._descan_shifts.data = shifts.to(self.device)
 
@@ -273,7 +274,7 @@ def scan_positions_px(self, positions: torch.Tensor | np.ndarray) -> None:
             positions,
             name="scan_positions_px",
             dtype=getattr(torch, config.get("dtype_real")),
-            shape=(self.num_gpts, 2),
+            shape=(self.num_positions, 2),
         )
         self._scan_positions_px.data = positions.to(self.device)
 
@@ -321,7 +322,7 @@ def initial_descan_shifts(self, shifts: torch.Tensor | np.ndarray) -> None:
             shifts,
             name="initial_descan_shifts",
             dtype=getattr(torch, config.get("dtype_real")),
-            shape=(self.num_gpts, 2),
+            shape=(self.num_positions, 2),
         )
         self._initial_descan_shifts = shifts
 
@@ -336,7 +337,7 @@ def initial_scan_positions_px(self, positions: torch.Tensor | np.ndarray) -> Non
             positions,
             name="initial_scan_positions_px",
             dtype=getattr(torch, config.get("dtype_real")),
-            shape=(self.num_gpts, 2),
+            shape=(self.num_positions, 2),
         )
         self._initial_scan_positions_px = positions
 
@@ -386,7 +387,7 @@ def patch_indices(self) -> torch.Tensor:
 
     # region --- torch.utils.data.Dataset interface ---
     def __len__(self) -> int:
-        return self.num_gpts
+        return self.num_positions
 
     def __getitem__(self, idx: int) -> dict[str, Any]:
         """Return one sample for the DataLoader.
@@ -424,7 +425,7 @@ def centered_amplitudes(self, arr: "np.ndarray | torch.Tensor") -> None:
             name="centered_amplitudes",
             dtype=getattr(torch, config.get("dtype_real")),
             ndim=3,
-            shape=(self.num_gpts, *self.roi_shape),
+            shape=(self.num_positions, *self.roi_shape),
         )
         self._centered_amplitudes = arr
 
@@ -440,7 +441,7 @@ def amplitudes(self, arr: "np.ndarray | torch.Tensor") -> None:
             name="amplitudes",
             dtype=getattr(torch, config.get("dtype_real")),
             ndim=3,
-            shape=(self.num_gpts, *self.roi_shape),
+            shape=(self.num_positions, *self.roi_shape),
         )
         self._amplitudes = arr
 
@@ -458,7 +459,7 @@ def centered_intensities(self, arr: "np.ndarray | torch.Tensor") -> None:
             name="centered_intensities",
             dtype=getattr(torch, config.get("dtype_real")),
             ndim=3,
-            shape=(self.num_gpts, *self.roi_shape),
+            shape=(self.num_positions, *self.roi_shape),
         )
         self._centered_intensities = arr
 
@@ -474,10 +475,28 @@ def intensities(self, arr: "np.ndarray | torch.Tensor") -> None:
             name="intensities",
             dtype=getattr(torch, config.get("dtype_real")),
             ndim=3,
-            shape=(self.num_gpts, *self.roi_shape),
+            shape=(self.num_positions, *self.roi_shape),
         )
         self._intensities = arr
 
+    @property
+    def positions_mask(self) -> torch.Tensor:
+        """1D boolean mask (length ``num_gpts``) of scan positions kept in the reconstruction.
+        Initialized to all-True. Masked-out (False) positions are dropped from the targets,
+        scan positions, descan shifts, and patch indices during preprocessing.
+        """
+        return self._positions_mask
+
+    @positions_mask.setter
+    def positions_mask(self, mask: torch.Tensor | np.ndarray) -> None:
+        # accept a 2D real-space mask (gpts) or any array-like; flatten to the raster order
+        mask = validate_tensor(mask, name="positions_mask", dtype=torch.bool).reshape(-1)
+        if mask.shape[0] != self.num_gpts:
+            raise ValueError(
+                f"positions_mask must have {self.num_gpts} elements, got {mask.shape[0]}"
+            )
+        self._positions_mask = mask
+
     @property
     def verbose(self) -> int:
         return self._verbose
@@ -571,6 +590,16 @@ def roi_shape(self) -> np.ndarray:
     def num_gpts(self) -> int:
         return int(self.dset.shape[0])
 
+    @property
+    def num_positions(self) -> int:
+        """Number of active scan positions, i.e. ``positions_mask.sum()``.
+
+        Equals ``num_gpts`` until a sparser ``positions_mask`` is applied. This is the length
+        of ``targets``/``scan_positions_px``/``descan_shifts`` and the index range the
+        reconstruction loops (DataLoader, train/val split) operate over.
+        """
+        return int(self.positions_mask.sum())
+
     @property
     def detector_sampling(self) -> np.ndarray:
         """Detector sampling in reciprocal space. Units of A^-1"""
@@ -1095,6 +1124,9 @@ def _set_initial_scan_positions_px(
         if obj_padding_px is None:
             obj_padding_px = np.array([0, 0])
 
+        if positions_mask is not None:
+            self.positions_mask = positions_mask
+
         nr, nc = self.gpts
         Sr, Sc = self._scan_sampling
         r = np.arange(nr) * Sr
@@ -1102,9 +1134,9 @@ def _set_initial_scan_positions_px(
 
         r, c = np.meshgrid(r, c, indexing="ij")
 
-        if positions_mask is not None:
-            r = r[positions_mask]
-            c = c[positions_mask]
+        pos_mask_2d = self.positions_mask.reshape(nr, nc).cpu().numpy()
+        r = r[pos_mask_2d]
+        c = c[pos_mask_2d]
 
         positions = np.stack((r.ravel(), c.ravel()), axis=-1).astype(config.get("dtype_real"))
 
@@ -1128,7 +1160,6 @@ def _set_initial_scan_positions_px(
         # top-left padding
         positions[:, 0] += obj_padding_px[0]
         positions[:, 1] += obj_padding_px[1]
-
         self.scan_positions_px = positions
         self.initial_scan_positions_px = self.scan_positions_px.data.clone()
         return
@@ -1145,7 +1176,11 @@ def preprocess(
         plot_com: str | bool = True,
         vectorized: bool = True,
         probe_energy: float | None = None,
+        positions_mask: np.ndarray | None = None,
     ):
+        if positions_mask is not None:
+            self.positions_mask = positions_mask
+
         # Store preprocessing parameters for serialization and reloading
         self._preprocessing_params = {
             "com_fit_function": com_fit_function,
@@ -1157,6 +1192,7 @@ def preprocess(
             "plot_rotation": False,
             "plot_com": False,
             "vectorized": vectorized,
+            "positions_mask": self.positions_mask.cpu().numpy(),
         }
 
         if probe_energy is not None:
@@ -1298,7 +1334,8 @@ def _set_intensities_com(
             com_fit_r = com_fit_r * self.roi_shape[0] / 2
             com_fit_c = com_fit_c * self.roi_shape[1] / 2
         else:
-            finite_mask = np.isfinite(com_measured_r)
+            pos_mask_2d = self.positions_mask.reshape(self.gpts[0], self.gpts[1]).cpu().numpy()
+            finite_mask = np.isfinite(com_measured_r) & pos_mask_2d
             com_fit_r, com_fit_c, _com_res_r, _com_res_c = fit_origin(
                 data=(com_measured_r, com_measured_c),
                 fit_function=fit_function,
@@ -1354,7 +1391,7 @@ def calculate_curl(com_r: np.ndarray, com_c: np.ndarray) -> float:
             """Calculate curl of CoM gradient vector field"""
             grad_r_c = com_r[1:-1, 2:] - com_r[1:-1, :-2]  # dVh/dw
             grad_c_r = com_c[2:, 1:-1] - com_c[:-2, 1:-1]  # dVw/dh
-            return float(np.mean(np.abs(grad_c_r - grad_r_c)))
+            return float(np.nanmean(np.abs(grad_c_r - grad_r_c)))
 
         def calculate_curl_for_angles(
             angles_rad: np.ndarray,
@@ -1373,7 +1410,7 @@ def calculate_curl_for_angles(
 
             grad_r_c = rot_r[:, 1:-1, 2:] - rot_r[:, 1:-1, :-2]
             grad_c_r = rot_c[:, 2:, 1:-1] - rot_c[:, :-2, 1:-1]
-            return np.mean(np.abs(grad_c_r - grad_r_c), axis=(-2, -1))
+            return np.nanmean(np.abs(grad_c_r - grad_r_c), axis=(-2, -1))
 
         def plot_curl_results(
             angles_deg: np.ndarray,
@@ -1438,6 +1475,10 @@ def plot_com_images(
         rotation_angles_deg = np.asarray(rotation_angles_deg)
         rotation_angles_rad = np.deg2rad(rotation_angles_deg)
 
+        pos_mask_2d = self.positions_mask.reshape(self.gpts[0], self.gpts[1]).cpu().numpy()
+        com_normalized_masked = self.com_normalized.copy()
+        com_normalized_masked[:, ~pos_mask_2d] = np.nan
+
         # Case 1: Known rotation
         if force_com_rotation is not None:
             _rotation_best_rad = np.deg2rad(force_com_rotation)
@@ -1452,12 +1493,12 @@ def plot_com_images(
             else:
                 # Calculate curl for both transpose options
                 rot_r, rot_c = rotate_com_vectors(
-                    self.com_normalized, _rotation_best_rad, transpose=False
+                    com_normalized_masked, _rotation_best_rad, transpose=False
                 )
                 rotation_curl = calculate_curl(rot_r, rot_c)
 
                 rot_r, rot_c = rotate_com_vectors(
-                    self.com_normalized, _rotation_best_rad, transpose=True
+                    com_normalized_masked, _rotation_best_rad, transpose=True
                 )
                 rotation_curl_transpose = calculate_curl(rot_r, rot_c)
 
@@ -1477,7 +1518,7 @@ def plot_com_images(
                 # Calculate curl for all angles with known transpose
                 curl_values = calculate_curl_for_angles(
                     rotation_angles_rad,
-                    self.com_normalized,
+                    com_normalized_masked,
                     transpose=_rotation_best_transpose,
                 )
 
@@ -1501,13 +1542,13 @@ def plot_com_images(
                 # Calculate curl for both transpose options
                 rotation_curl = calculate_curl_for_angles(
                     rotation_angles_rad,
-                    self.com_normalized,
+                    com_normalized_masked,
                     transpose=False,
                 )
 
                 rotation_curl_transpose = calculate_curl_for_angles(
                     rotation_angles_rad,
-                    self.com_normalized,
+                    com_normalized_masked,
                     transpose=True,
                 )
 
@@ -1576,6 +1617,9 @@ def _normalize_diffraction_intensities(
         dtype = config.get("dtype_real")
         diff_intensities = self.intensities_4d.copy().astype(dtype)
         com_fit = self.com_fit
+        if positions_mask is not None:
+            self.positions_mask = positions_mask
+        positions_mask_2d = self.positions_mask.reshape(self.gpts[0], self.gpts[1]).cpu().numpy()
 
         # Aggressive cropping for when off-centered high scattering angle data was recorded
         if crop_patterns:
@@ -1613,9 +1657,8 @@ def _normalize_diffraction_intensities(
             unit="probe position",
             disable=not self._verbose,
         ):
-            if positions_mask is not None:
-                if not positions_mask[Rr, Rc]:
-                    continue
+            if not positions_mask_2d[Rr, Rc]:
+                continue
 
             intensity = np.maximum(diff_intensities[Rr, Rc], 0)
             intensities[Rr, Rc] = intensity
@@ -1637,16 +1680,10 @@ def _normalize_diffraction_intensities(
             centered_amplitudes[Rr, Rc] = shift_amplitude
             centered_intensities[Rr, Rc] = shift_amplitude**2
 
-        if positions_mask is not None:
-            amplitudes = amplitudes[positions_mask]
-            centered_amplitudes = centered_amplitudes[positions_mask]
-            intensities = intensities[positions_mask]
-            centered_intensities = centered_intensities[positions_mask]
-        else:
-            amplitudes = amplitudes.reshape((-1, *self.roi_shape))
-            centered_amplitudes = centered_amplitudes.reshape((-1, *self.roi_shape))
-            intensities = intensities.reshape((-1, *self.roi_shape))
-            centered_intensities = centered_intensities.reshape((-1, *self.roi_shape))
+        amplitudes = amplitudes[positions_mask_2d]
+        centered_amplitudes = centered_amplitudes[positions_mask_2d]
+        intensities = intensities[positions_mask_2d]
+        centered_intensities = centered_intensities[positions_mask_2d]
 
         if crop_patterns:
             amplitudes = amplitudes[:, pattern_crop_mask].reshape((-1, *pattern_crop_mask_shape))
@@ -1665,10 +1702,11 @@ def _normalize_diffraction_intensities(
         self.amplitudes = amplitudes
         self.centered_intensities = centered_intensities
         self.intensities = intensities
-        descan_shifts = -1 * np.stack((com_fit[0].flatten(), com_fit[1].flatten()))
         descan_shifts = -1 * com_fit.reshape((2, -1))  # (2, rr*rc)
         descan_shifts += self.roi_shape[:, None] / 2
-        self.descan_shifts = descan_shifts.T
+        descan_shifts = descan_shifts.T  # (rr*rc, 2)
+        descan_shifts = descan_shifts[positions_mask_2d.ravel()]
+        self.descan_shifts = descan_shifts
         self.initial_descan_shifts = self.descan_shifts.data.clone()
 
         self.mean_diffraction_intensity = mean_intensity
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index 29a690475..ea277533b 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -413,8 +413,12 @@ def fit_origin(
     elif fit_function == "bezier_two":
         f = _bezier_two
     elif fit_function == "constant":
-        qr0_fit = np.mean(qr0_meas) * np.ones_like(qr0_meas)
-        qc0_fit = np.mean(qc0_meas) * np.ones_like(qc0_meas)
+        # only average over the masked-in (and finite) positions; otherwise NaN/masked-out
+        # positions (e.g. zeroed diffraction patterns) would poison the mean
+        qr0_sel = qr0_meas[mask] if mask is not None else qr0_meas
+        qc0_sel = qc0_meas[mask] if mask is not None else qc0_meas
+        qr0_fit = np.nanmean(qr0_sel) * np.ones_like(qr0_meas)
+        qc0_fit = np.nanmean(qc0_sel) * np.ones_like(qc0_meas)
         qr0_residuals = qr0_meas - qr0_fit
         qc0_residuals = qc0_meas - qc0_fit
         return qr0_fit, qc0_fit, qr0_residuals, qc0_residuals
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index b4a25cc7b..a0890f5a9 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -380,11 +380,11 @@ def _reconstruct_inner(
         self.dset._set_targets(self._criterion.target_space)
         self.compute_propagator_arrays()  # required to avoid issue if stopped learning probe tilt
 
-        # Compute the global scan count once — needed to keep loss scale consistent across world
-        global_n = self.dset.num_gpts
+        # Compute the global scan count once — needed to keep loss scale consistent across world.
+        global_n = self.dset.num_positions
 
         train_indices, val_indices = compute_train_val_split(
-            self.dset.num_gpts,
+            self.dset.num_positions,
             self.val_ratio,
             self.val_mode,
             self.rng,
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 1c830f127..0c8482a5f 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -217,7 +217,7 @@ def _set_obj_fov_mask(self, gaussian_sigma: float = 2.0, batch_size=None):
 
     def _get_probe_overlap(self, max_batch_size: int | None = None) -> np.ndarray:
         prb = self.probe_model.probe[0]
-        num_dps = int(np.prod(self.gpts))
+        num_dps = self.dset.num_positions
         shifted_probes = prb.expand(num_dps, *self.roi_shape)
 
         batch_size = num_dps if max_batch_size is None else int(max_batch_size)
@@ -1120,7 +1120,7 @@ def error_estimate(
             preds = pred_intensities
 
         mask = self.dset.detector_mask
-        n = global_n if global_n is not None else self.dset.num_gpts
+        n = global_n if global_n is not None else self.dset.num_positions
         error = criterion(preds * mask, targets * mask, n)
         loss = error / self.dset.mean_diffraction_intensity
         return loss, targets

From 2a4972e7fd5a87eb68c71295d4b1cf0245d460ad Mon Sep 17 00:00:00 2001
From: smribet <stephanie.ribet@gmail.com>
Date: Sat, 13 Jun 2026 05:48:45 -0700
Subject: [PATCH 51/59] small vis change

---
 src/quantem/diffractive_imaging/direct_ptychography.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index 3f3bca7ac..65dfa0a11 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -945,8 +945,6 @@ def visualize(
             [obj, np.abs(obj_fft)],
             title=["Object", "Object FFT"],
             scalebar=[obj_scalebar, fft_scalebar],
-            cmap=["magma", "magma"],
-            norm=[None, None],
             **kwargs,
         )
         axs[1].set_aspect(obj.shape[-1] / obj.shape[-2])

From 6e15f34c9a23dc1ffe5e6ace681ae69b6c1a06b3 Mon Sep 17 00:00:00 2001
From: Georgios Varnavides <g.varnavides@tudelft.nl>
Date: Mon, 15 Jun 2026 15:57:32 +0200
Subject: [PATCH 52/59] clean up

---
 .../direct_ptycho_utils.py                    |  6 +++
 .../direct_ptychography.py                    | 49 +++++++++++--------
 2 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/src/quantem/diffractive_imaging/direct_ptycho_utils.py b/src/quantem/diffractive_imaging/direct_ptycho_utils.py
index fb40f9a07..ce08bfeba 100644
--- a/src/quantem/diffractive_imaging/direct_ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/direct_ptycho_utils.py
@@ -37,6 +37,12 @@
 # fmt: on
 
 
+def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None:
+    if rotation_angle is None:
+        return None
+    return math.radians(float(rotation_angle))
+
+
 def create_edge_window(shape, edge_blend_pixels, device="cpu"):
     """
     Create a smooth edge window that transitions from 0 at edges to 1 in center.
diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py
index 65dfa0a11..912b4bd2e 100644
--- a/src/quantem/diffractive_imaging/direct_ptychography.py
+++ b/src/quantem/diffractive_imaging/direct_ptychography.py
@@ -45,6 +45,7 @@
 from quantem.diffractive_imaging.direct_ptycho_utils import (
     ABERRATION_PRESETS,
     _crop_corner_centered_mask,
+    _rotation_degrees_to_radians,
     align_vbf_stack_multiscale,
     create_edge_window,
     fit_aberrations_from_shifts,
@@ -55,12 +56,6 @@
 optuna.logging.set_verbosity(optuna.logging.WARNING)
 
 
-def _rotation_degrees_to_radians(rotation_angle: float | None) -> float | None:
-    if rotation_angle is None:
-        return None
-    return math.radians(float(rotation_angle))
-
-
 @dataclass
 class OptimizationParameter:
     low: float
@@ -914,6 +909,8 @@ def obj(self) -> np.ndarray:
     def visualize(
         self,
         return_fig: bool = False,
+        show_obj_fft: bool = True,
+        apply_hanning_window: bool = False,
         **kwargs,
     ):
         """
@@ -934,20 +931,32 @@ def visualize(
             raise RuntimeError("Run reconstruct() before visualize().")
 
         obj = self.obj
-        window = np.hanning(obj.shape[-2])[:, None] * np.hanning(obj.shape[-1])[None, :]
-        obj_fft = np.fft.fftshift(np.fft.fft2(obj * window))
-
-        obj_scalebar = {"sampling": self.scan_sampling[0], "units": "Å"}
-        fft_sampling = 1 / (self.scan_sampling[0] * obj.shape[-2])
-        fft_scalebar = {"sampling": fft_sampling, "units": r"$\mathrm{A^{-1}}$"}
-
-        fig, axs = show_2d(
-            [obj, np.abs(obj_fft)],
-            title=["Object", "Object FFT"],
-            scalebar=[obj_scalebar, fft_scalebar],
-            **kwargs,
-        )
-        axs[1].set_aspect(obj.shape[-1] / obj.shape[-2])
+        obj_scalebar = {"sampling": self.scan_sampling[1], "units": "Å"}
+
+        if show_obj_fft:
+            if apply_hanning_window:
+                window = np.hanning(obj.shape[-2])[:, None] * np.hanning(obj.shape[-1])[None, :]
+                obj_fft = np.fft.fftshift(np.abs(np.fft.fft2(obj * window)))
+            else:
+                obj_fft = np.fft.fftshift(np.abs(np.fft.fft2(obj)))
+
+            fft_sampling = 1 / (self.scan_sampling[1] * obj.shape[-1])
+            fft_scalebar = {"sampling": fft_sampling, "units": r"$\mathrm{A^{-1}}$"}
+
+            fig, axs = show_2d(
+                [obj, obj_fft],
+                title=["Object phase", "Object phase FFT"],
+                scalebar=[obj_scalebar, fft_scalebar],
+                **kwargs,
+            )
+            axs[1].set_aspect(obj.shape[-1] / obj.shape[-2])
+        else:
+            fig, axs = show_2d(
+                obj,
+                title="Object phase",
+                scalebar=obj_scalebar,
+                **kwargs,
+            )
 
         if return_fig:
             return fig, axs

From 54618f936fca45e8072a0f45b7837603f5e27ccf Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Wed, 17 Jun 2026 10:03:07 -0700
Subject: [PATCH 53/59] fixing device handling coverage in models, passing
 kwaargs in show_fourier_probe

---
 src/quantem/core/utils/rng.py                                 | 4 +++-
 src/quantem/diffractive_imaging/object_models.py              | 2 +-
 src/quantem/diffractive_imaging/probe_models.py               | 3 +--
 .../diffractive_imaging/ptychography_visualizations.py        | 4 ++--
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/quantem/core/utils/rng.py b/src/quantem/core/utils/rng.py
index acf0342f1..8a4b5efab 100644
--- a/src/quantem/core/utils/rng.py
+++ b/src/quantem/core/utils/rng.py
@@ -3,6 +3,8 @@
 import numpy as np
 import torch
 
+from quantem.core import config
+
 DeviceType = Union[str, "torch.device", int]
 
 
@@ -26,7 +28,7 @@ def __init__(
         *args,
         **kwargs,
     ):
-        self._device = device
+        self._device = config.validate_device(device)[0]
         self.rng = rng
 
     @property
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 0301be886..e0e6eab8b 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -1111,7 +1111,7 @@ def from_model(
             rng=rng,
             _token=cls._token,
         )
-        obj_model.model = model.to(device)
+        obj_model.model = model.to(obj_model.device)
         obj_model.model_input = model_input
         obj_model._set_pretrained_weights(model)
 
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index d5f8eea92..8a034ad65 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -156,14 +156,13 @@ def __init__(
     ):
         if _token is not self._token:
             raise RuntimeError("Use a factory method to instantiate this class.")
-
         # Initialize nn.Module first
         nn.Module.__init__(self)
         RNGMixin.__init__(self, rng=rng, device=device)
         OptimizerMixin.__init__(self)
 
         self.num_probes = num_probes
-        self._device = device
+        self.device = device
         self._probe_params = self.DEFAULT_PROBE_PARAMS
         self._max_aberrations_order = max_aberrations_order
         self.probe_params = probe_params
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index e49f9bc7d..dcb10ea3f 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -387,7 +387,7 @@ def show_probe_top_bottom(
             return fig, axs, {"top": top_img, "bottom": bottom_img}
         return fig, axs
 
-    def show_fourier_probe(self, probe: np.ndarray | None = None):
+    def show_fourier_probe(self, probe: np.ndarray | None = None, **kwargs):
         """
         Show the Fourier transform of the probe.
         Parameters
@@ -421,7 +421,7 @@ def show_fourier_probe(self, probe: np.ndarray | None = None):
             ]
         else:
             titles = "Fourier Probe"
-        show_2d(probes, title=titles, scalebar=scalebar)
+        show_2d(probes, title=titles, scalebar=scalebar, **kwargs)
 
     def show_obj_and_probe(
         self,

From 63e06bc0585cdf6f7af8b35333836fc9bb187b23 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Thu, 2 Jul 2026 13:51:41 -0700
Subject: [PATCH 54/59] fixing bug combining large aspect ratio scans with
 large com_rotations

---
 .../diffractive_imaging/dataset_models.py     | 13 ++-
 .../diffractive_imaging/ptychography_base.py  | 30 +------
 .../diffractive_imaging/test_ptychography.py  | 80 +++++++++++++++++++
 3 files changed, 91 insertions(+), 32 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index fc1456bb4..59b80fc31 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -644,10 +644,12 @@ def reciprocal_units(self) -> list[str]:
 
     @property
     def _obj_shape_crop_2d(self) -> np.ndarray:
-        """All object shapes are 2D"""
+        """Allocated/expected object crop shape, in the object array's own axis order."""
         shp = np.floor(self.fov / self.obj_sampling)
         shp += shp % 2
         shp = shp.astype("int")
+        if self.com_transpose: # match the axis-flipped scan positions
+            shp = shp[::-1]
         return shp
 
     @property
@@ -1095,7 +1097,10 @@ def fov(self) -> np.ndarray:
 
     @property
     def upsample_factor(self) -> float:
-        return (self._obj_shape_crop_2d / self.gpts).mean()
+        # pair each crop axis with its own scan-grid axis so the factor stays
+        # transpose-invariant (_obj_shape_crop_2d is transposed when com_transpose)
+        gpts = self.gpts[::-1] if self.com_transpose else self.gpts
+        return (self._obj_shape_crop_2d / gpts).mean()
 
     # endregion --- properties ---
 
@@ -1149,8 +1154,8 @@ def _set_initial_scan_positions_px(
             positions = np.flip(positions, axis=1)
             sampling = sampling[::-1]
 
-        # ensure positive
-        m: np.ndarray = np.min(positions, axis=0).clip(-np.inf, 0)
+        # anchor the used positions to the object origin
+        m: np.ndarray = np.min(positions, axis=0)
         positions -= m
 
         # finally, switch to pixels
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 0c8482a5f..532b4d5da 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -34,7 +34,6 @@
 from quantem.diffractive_imaging.probe_models import ProbeBase, ProbeModelType, ProbePixelated
 from quantem.diffractive_imaging.ptycho_losses import DataCriterion, get_data_criterion
 from quantem.diffractive_imaging.ptycho_utils import (
-    AffineTransform,
     center_crop_arr,
     fourier_translation_operator,
     sum_patches,
@@ -857,45 +856,20 @@ def _crop_rotate_obj_fov(
         transpose: bool | None = None,
         padding: np.ndarray | tuple[int, int] | None = None,
     ) -> np.ndarray:
-        """
-        Crops and rotated object to FOV bounded by current pixel positions.
-        """
+        """Un-rotates and un-transposes the object and crops it to the reconstruction FOV."""
         array = self._to_numpy(array).copy()
         com_rotation_rad = (
             self.dset.com_rotation_rad if com_rotation_rad is None else com_rotation_rad
         )
         transpose = self.dset.com_transpose if transpose is None else transpose
-        padding = np.array(padding) if padding is not None else self.obj_padding_px
 
         angle = com_rotation_rad if transpose else -1 * com_rotation_rad
 
-        if positions_px is None:
-            positions = self.dset.initial_scan_positions_px.cpu().detach().numpy()
-            # if using learned positions potentially need to pad the object in center_crop_arr
-            # positions = self.dset.scan_positions_px.cpu().detach().numpy()
-        else:
-            positions = positions_px
-
-        tf = AffineTransform(angle=angle)
-        rotated_points = tf(positions, origin=positions.mean(0))
-        rotated_points += 1e-9  # avoid pixel perfect errors
-
-        min_r, min_c = np.floor(np.min(rotated_points, axis=0)).astype("int")
-        min_r = max(min_r, 0)
-        min_c = max(min_c, 0)
-        max_r, max_c = np.ceil(np.max(rotated_points, axis=0)).astype("int")
-        max_r = min(max_r, array.shape[-2])
-        max_c = min(max_c, array.shape[-1])
-        # print(f"{min_r = }, {min_c = }, {max_r = }, {max_c = }")
-
-        rotated_array = ndi.rotate(
-            array, np.rad2deg(-angle), order=1, reshape=False, axes=(-2, -1)
-        )[..., min_r:max_r, min_c:max_c]
+        rotated_array = ndi.rotate(array, np.rad2deg(-angle), order=1, reshape=True, axes=(-2, -1))
 
         if transpose:
             rotated_array = rotated_array.swapaxes(-2, -1)
 
-        # fixing that is sometimes 1 pixel off
         cropped = center_crop_arr(rotated_array, tuple(self.obj_shape_crop), pad_if_needed=False)
 
         return cropped
diff --git a/tests/diffractive_imaging/test_ptychography.py b/tests/diffractive_imaging/test_ptychography.py
index 6d9c16ce6..854c0fe5b 100644
--- a/tests/diffractive_imaging/test_ptychography.py
+++ b/tests/diffractive_imaging/test_ptychography.py
@@ -392,6 +392,86 @@ def test_save_load_roundtrip(self, ptycho_dataset, tmp_path):
         assert reloaded.target_residency == "cpu"
 
 
+def _build_aspect_ratio_ptycho(complex_obj, probe_array, gpts, com_rotation, transpose):
+    """Build a Ptychography on a (possibly non-square) scan grid with a forced rotation
+    and transpose. Reconstruction quality is irrelevant here — the fixture only exercises
+    scan-position placement and object sizing."""
+    scan_x, scan_y = gpts
+    x = np.arange(0.0, scan_x * SCAN_STEP_SIZE, SCAN_STEP_SIZE)
+    y = np.arange(0.0, scan_y * SCAN_STEP_SIZE, SCAN_STEP_SIZE)
+    xx, yy = np.meshgrid(x, y, indexing="ij")
+    positions = np.stack((xx.ravel(), yy.ravel()), axis=-1)
+    reciprocal_sampling = 2 * Q_MAX / N
+
+    sim_row, sim_col = return_patch_indices(positions, (N, N), (N, N))
+    _, _, _, intensities = simulate_intensities(complex_obj, probe_array, sim_row, sim_col)
+
+    dset = Dataset4dstem.from_array(
+        array=np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((scan_x, scan_y, N, N)),
+        sampling=(SCAN_STEP_SIZE, SCAN_STEP_SIZE, reciprocal_sampling, reciprocal_sampling),
+        units=("A", "A", "A^-1", "A^-1"),
+    )
+    pdset = PtychographyDatasetRaster.from_dataset4dstem(dset)
+    pdset.preprocess(
+        com_fit_function="constant",
+        plot_rotation=False,
+        plot_com=False,
+        probe_energy=PROBE_ENERGY,
+        force_com_rotation=com_rotation,
+        force_com_transpose=transpose,
+    )
+
+    probe_params = {
+        "energy": PROBE_ENERGY,
+        "C10": C10,
+        "semiangle_cutoff": electron_wavelength_angstrom(PROBE_ENERGY) * 1e3,
+    }
+    ptycho = Ptychography.from_models(
+        dset=pdset,
+        obj_model=ObjectPixelated.from_uniform(
+            num_slices=1, obj_type="complex", slice_thicknesses=1
+        ),
+        probe_model=ProbePixelated.from_array(
+            num_probes=1, probe_params=probe_params, probe_array=probe_array
+        ),
+        detector_model=DetectorPixelated(),
+        rng=42,
+    )
+    ptycho.preprocess(obj_padding_px=(0, 0))
+    return ptycho
+
+
+class TestAspectRatioRotatedFOV:
+    """Regression for the non-square / rotated / transposed FOV bug.
+
+    A high aspect-ratio scan grid combined with a large com_rotation (and/or transpose)
+    used to (a) place scan positions outside the object array — wrapping the object and
+    pinning positions to the FOV edge — and (b) crash ``obj_cropped`` inside
+    ``_crop_rotate_obj_fov`` because the un-rotated FOV could not fit the object frame.
+    """
+
+    @pytest.mark.parametrize("transpose", [False, True])
+    def test_positions_inside_object(self, complex_obj, probe_array, transpose):
+        ptycho = _build_aspect_ratio_ptycho(
+            complex_obj, probe_array, gpts=(32, 64), com_rotation=89, transpose=transpose
+        )
+        obj_full = ptycho.dset._obj_shape_full_2d(ptycho.obj_padding_px)
+        pos = ptycho.dset.initial_scan_positions_px.cpu().detach().numpy()
+        # positions must sit strictly inside the object (no wrap / no edge pile-up)
+        assert pos[:, 0].min() >= 0 and pos[:, 1].min() >= 0
+        assert pos[:, 0].max() < obj_full[-2]
+        assert pos[:, 1].max() < obj_full[-1]
+
+    @pytest.mark.parametrize("transpose", [False, True])
+    def test_obj_cropped_shape(self, complex_obj, probe_array, transpose):
+        ptycho = _build_aspect_ratio_ptycho(
+            complex_obj, probe_array, gpts=(32, 64), com_rotation=89, transpose=transpose
+        )
+        # obj_cropped must not raise and returns the non-transposed display FOV shape
+        cropped = ptycho.obj_cropped
+        assert tuple(cropped.shape) == tuple(ptycho.obj_shape_crop)
+
+
 @pytest.mark.slow
 class TestPtychographySaveLoadRoundtrip:
     """Reconstruct → save → load → continue training preserves training state.

From 27ba927f18069668039c8073d2c09b38f9360a53 Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Tue, 7 Jul 2026 16:01:00 -0700
Subject: [PATCH 55/59] reducing memory for preprocessing and free port bugfix

---
 src/quantem/core/ml/dist_utils.py             |  36 +++++
 .../diffractive_imaging/dataset_models.py     |  15 +-
 .../diffractive_imaging/ptychography.py       | 139 ++++++++++++------
 .../diffractive_imaging/ptychography_base.py  |  18 ++-
 4 files changed, 155 insertions(+), 53 deletions(-)

diff --git a/src/quantem/core/ml/dist_utils.py b/src/quantem/core/ml/dist_utils.py
index 450c4fead..6a5c561b0 100644
--- a/src/quantem/core/ml/dist_utils.py
+++ b/src/quantem/core/ml/dist_utils.py
@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import os
+import socket
 from typing import Any
 
 import torch
@@ -19,6 +20,40 @@ def is_distributed_launch() -> bool:
     return "RANK" in os.environ
 
 
+def find_free_port() -> str:
+    """Return a currently-free TCP port (as a string) on the loopback interface.
+
+    Used to pick the rendezvous port for the notebook ``mp.spawn`` path instead of a
+    hardcoded ``29500``. A hardcoded port collides across repeated ``reconstruct`` cell
+    re-runs (a run that errors before ``destroy_process_group`` leaves the TCPStore server
+    socket bound), producing "client socket ... failed to connect" / "address already in
+    use" on the next call. Binding to port 0 lets the OS hand back an unused port each time.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return str(s.getsockname()[1])
+
+
+def maybe_configure_fabric_env() -> None:
+    """Set the NCCL/libfabric env for the HPE Slingshot (``hsn``) fabric, if present.
+
+    Perlmutter (and other Slingshot-11 systems) need ``NCCL_SOCKET_IFNAME=hsn`` plus
+    ``FI_CXI_ATS=0`` / ``NCCL_CROSS_NIC=1`` for NCCL to bring up its communicators cleanly;
+    without them multi-GPU init can hang or emit fatal socket errors. Gated on the presence
+    of an ``hsn0`` interface and on each var being unset, so this is a no-op on non-Slingshot
+    systems (e.g. a local GPU workstation) and never overrides an explicit user setting.
+    """
+    if not os.path.isdir("/sys/class/net/hsn0"):
+        return
+    defaults = {
+        "NCCL_SOCKET_IFNAME": "hsn",
+        "FI_CXI_ATS": "0",
+        "NCCL_CROSS_NIC": "1",
+    }
+    for key, value in defaults.items():
+        os.environ.setdefault(key, value)
+
+
 def init_process_group(
     rank: int,
     world_size: int,
@@ -40,6 +75,7 @@ def init_process_group(
     """
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
+    maybe_configure_fabric_env()
     if backend == "nccl":
         device_index = local_device if local_device is not None else rank
         torch.cuda.set_device(device_index)
diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index 59b80fc31..f7d0fbe81 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -1620,7 +1620,7 @@ def _normalize_diffraction_intensities(
         bilinear: bool = False,
     ):
         dtype = config.get("dtype_real")
-        diff_intensities = self.intensities_4d.copy().astype(dtype)
+        diff_intensities = self.intensities_4d.astype(dtype)
         com_fit = self.com_fit
         if positions_mask is not None:
             self.positions_mask = positions_mask
@@ -1651,10 +1651,7 @@ def _normalize_diffraction_intensities(
         mean_amplitude = 0
         centered_amplitudes = np.zeros(diff_intensities.shape, dtype=dtype)
         amplitudes = np.zeros(diff_intensities.shape, dtype=dtype)
-        centered_intensities = np.zeros(diff_intensities.shape, dtype=dtype)
         intensities = np.zeros(diff_intensities.shape, dtype=dtype)
-        ## there is some additional memory overhead in this loop due to numpy array assignment
-        ## but I don't think it's easy to avoid -- ARCM 251212
         for Rr, Rc in tqdmnd(
             range(diff_intensities.shape[0]),
             range(diff_intensities.shape[1]),
@@ -1683,12 +1680,13 @@ def _normalize_diffraction_intensities(
             shift_amplitude = np.fft.fftshift(shift_amplitude)
 
             centered_amplitudes[Rr, Rc] = shift_amplitude
-            centered_intensities[Rr, Rc] = shift_amplitude**2
+
+        # The source 4D copy is no longer needed; free it before the masking copies below 
+        del diff_intensities
 
         amplitudes = amplitudes[positions_mask_2d]
         centered_amplitudes = centered_amplitudes[positions_mask_2d]
         intensities = intensities[positions_mask_2d]
-        centered_intensities = centered_intensities[positions_mask_2d]
 
         if crop_patterns:
             amplitudes = amplitudes[:, pattern_crop_mask].reshape((-1, *pattern_crop_mask_shape))
@@ -1696,9 +1694,8 @@ def _normalize_diffraction_intensities(
                 (-1, *pattern_crop_mask_shape)
             )
             intensities = intensities[:, pattern_crop_mask].reshape((-1, *pattern_crop_mask_shape))
-            centered_intensities = centered_intensities[:, pattern_crop_mask].reshape(
-                (-1, *pattern_crop_mask_shape)
-            )
+
+        centered_intensities = centered_amplitudes**2
 
         mean_intensity /= amplitudes.shape[0]
         mean_amplitude /= amplitudes.shape[0]
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index a0890f5a9..3f4206f7c 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -14,8 +14,10 @@
 
 from quantem.core.io.serialize import load as autoserialize_load
 from quantem.core.ml.dist_utils import (
+    find_free_port,
     init_process_group,
     is_distributed_launch,
+    maybe_configure_fabric_env,
     spawn_distributed_workers,
 )
 from quantem.diffractive_imaging.dataset_models import DatasetModelType
@@ -30,6 +32,20 @@
 from quantem.diffractive_imaging.ptychography_visualizations import PtychographyVisualizations
 
 
+def _cap_worker_cpu_threads(world_size: int) -> None:
+    """Split the visible CPU cores evenly across spawned workers."""
+    if os.environ.get("OMP_NUM_THREADS"):
+        return
+    try:
+        cpus = len(os.sched_getaffinity(0))
+    except AttributeError:  # platforms without sched_getaffinity (e.g. macOS)
+        cpus = os.cpu_count() or 1
+    threads = max(1, cpus // max(world_size, 1))
+    # Env var so BLAS/OpenMP pools initialized later in this process follow suit.
+    os.environ["OMP_NUM_THREADS"] = str(threads)
+    torch.set_num_threads(threads)
+
+
 def _ddp_ptycho_worker(
     rank: int,
     world_size: int,
@@ -37,13 +53,18 @@ def _ddp_ptycho_worker(
     devices: list[int],
     recon_kwargs: dict[str, Any],
     result_path: str,
+    master_port: str,
 ) -> None:
     """Module-level worker for mp.start_processes — must live at module scope to be picklable.
 
     Receives a file path rather than the Ptychography object directly so that no
     large tensors cross the process boundary via pickle (which triggers PyTorch's
     shared-memory tensor mechanism and fails in some Linux environments).
+
+    ``master_port`` is chosen (free) by the parent so all ranks rendezvous on the same
+    port and repeated ``reconstruct`` calls never collide on a stale ``29500``.
     """
+    _cap_worker_cpu_threads(world_size)
     device_id = devices[rank]
     # Bind the CUDA device BEFORE init_process_group so NCCL allocates its
     # communicator buffers on the correct GPU. Without this, NCCL grabs cuda:0
@@ -53,52 +74,68 @@ def _ddp_ptycho_worker(
         rank,
         world_size,
         backend="nccl" if torch.cuda.is_available() else "gloo",
+        master_port=master_port,
         local_device=device_id if torch.cuda.is_available() else None,
     )
 
-    # mmap=True so all workers share one memory-mapped RAM copy of the (potentially large,
-    # CPU-resident) state instead of each duplicating it.
-    ptycho = torch.load(ptycho_path, map_location="cpu", weights_only=False, mmap=True)
-    ptycho.to(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
-
-    if dist.is_available() and dist.is_initialized():
-        ptycho._broadcast_parameters(src=0)
-
-    ptycho._reconstruct_inner(**recon_kwargs, _dist_rank=rank, _dist_world_size=world_size)
-
-    if rank == 0:
-        obj_opt = ptycho.optimizers.get("object")
-        probe_opt = ptycho.optimizers.get("probe")
-        dset_opt = ptycho.optimizers.get("dataset")
-        torch.save(
-            {
-                "obj_state": {k: v.cpu() for k, v in ptycho.obj_model.state_dict().items()},
-                "probe_state": {k: v.cpu() for k, v in ptycho.probe_model.state_dict().items()},
-                # Dataset learnable params (scan positions / descan) are optimized and all-reduced
-                # in the workers; ship them back so the main process keeps the refinement.
-                "dset_scan_positions_px": ptycho.dset._scan_positions_px.detach().cpu(),
-                "dset_descan_shifts": ptycho.dset._descan_shifts.detach().cpu(),
-                "obj_optimizer_params": ptycho.obj_model._optimizer_params,
-                "probe_optimizer_params": ptycho.probe_model._optimizer_params,
-                "dset_optimizer_params": ptycho.dset._optimizer_params,
-                "obj_optimizer_state": obj_opt.state_dict() if obj_opt is not None else None,
-                "probe_optimizer_state": probe_opt.state_dict() if probe_opt is not None else None,
-                "dset_optimizer_state": dset_opt.state_dict() if dset_opt is not None else None,
-                "iter_losses": ptycho._iter_losses,
-                "iter_val_losses": ptycho._iter_val_losses,
-                "iter_lrs": ptycho._iter_lrs,
-                "iter_recon_types": ptycho._iter_recon_types,
-            },
-            result_path,
-        )
+    try:
+        # mmap=True so all workers share one memory-mapped RAM copy of the (potentially large,
+        # CPU-resident) state instead of each duplicating it.
+        ptycho = torch.load(ptycho_path, map_location="cpu", weights_only=False, mmap=True)
+        ptycho.to(f"cuda:{device_id}" if torch.cuda.is_available() else "cpu")
+
+        if dist.is_available() and dist.is_initialized():
+            ptycho._broadcast_parameters(src=0)
+
+        ptycho._reconstruct_inner(**recon_kwargs, _dist_rank=rank, _dist_world_size=world_size)
+
+        if rank == 0:
+            obj_opt = ptycho.optimizers.get("object")
+            probe_opt = ptycho.optimizers.get("probe")
+            dset_opt = ptycho.optimizers.get("dataset")
+            torch.save(
+                {
+                    "obj_state": {k: v.cpu() for k, v in ptycho.obj_model.state_dict().items()},
+                    "probe_state": {
+                        k: v.cpu() for k, v in ptycho.probe_model.state_dict().items()
+                    },
+                    # Dataset learnable params (scan positions / descan) are optimized and
+                    # all-reduced in the workers; ship them back so the main process keeps
+                    # the refinement.
+                    "dset_scan_positions_px": ptycho.dset._scan_positions_px.detach().cpu(),
+                    "dset_descan_shifts": ptycho.dset._descan_shifts.detach().cpu(),
+                    "obj_optimizer_params": ptycho.obj_model._optimizer_params,
+                    "probe_optimizer_params": ptycho.probe_model._optimizer_params,
+                    "dset_optimizer_params": ptycho.dset._optimizer_params,
+                    "obj_optimizer_state": obj_opt.state_dict() if obj_opt is not None else None,
+                    "probe_optimizer_state": (
+                        probe_opt.state_dict() if probe_opt is not None else None
+                    ),
+                    "dset_optimizer_state": dset_opt.state_dict()
+                    if dset_opt is not None
+                    else None,
+                    # Snapshots recorded during the run (rank 0 only) so the notebook multi-GPU
+                    # path returns them just like single-GPU reconstruct does.
+                    "snapshots": ptycho._snapshots,
+                    "iter_losses": ptycho._iter_losses,
+                    "iter_val_losses": ptycho._iter_val_losses,
+                    "iter_lrs": ptycho._iter_lrs,
+                    "iter_recon_types": ptycho._iter_recon_types,
+                },
+                result_path,
+            )
 
-    # Synchronize before teardown so every rank finishes
-    if dist.is_available() and dist.is_initialized():
-        if torch.cuda.is_available():
-            dist.barrier(device_ids=[device_id])
-        else:
-            dist.barrier()
-    dist.destroy_process_group()
+        # Synchronize before teardown so every rank finishes
+        if dist.is_available() and dist.is_initialized():
+            if torch.cuda.is_available():
+                dist.barrier(device_ids=[device_id])
+            else:
+                dist.barrier()
+    finally:
+        # Always tear down the process group — even if the run raised — so the rendezvous
+        # port is released and the next reconstruct() call starts clean.
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
 
 
 class Ptychography(PtychographyOpt, PtychographyVisualizations, PtychographyBase):  # pyright: ignore[reportUnsafeMultipleInheritance]
@@ -259,6 +296,11 @@ def reconstruct(
         from a notebook, or uses the existing distributed process group when launched with
         ``torchrun``. Only autograd mode is supported for multi-GPU in this release.
 
+        ``batch_size`` is GLOBAL: the number of samples contributing to one optimizer step,
+        regardless of GPU count. Under multi-GPU each rank draws ``batch_size // world_size``
+        per step (a warning is emitted when not evenly divisible), so the same ``batch_size``
+        reproduces the same optimization trajectory — and loss curve — on 1 or N GPUs.
+
         ``loss_type`` selects the data-fidelity criterion: a registered name
         (``"l2_amplitude"`` [default], ``"l1_amplitude"``, ``"l2_intensity"``, ``"l1_intensity"``,
         ``"poisson"``, ``"smooth_l1_amplitude"``, ``"s3im_amplitude"``) or a ``DataCriterion``
@@ -303,6 +345,7 @@ def reconstruct(
             if not torch.distributed.is_initialized():
                 # Bind the device BEFORE init_process_group so NCCL allocates
                 # its communicator buffers on the correct GPU.
+                maybe_configure_fabric_env()
                 if torch.cuda.is_available():
                     torch.cuda.set_device(local_rank)
                 torch.distributed.init_process_group(
@@ -556,6 +599,10 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
 
             torch.save(self, ptycho_path, pickle_protocol=4)
 
+            # Pick a free rendezvous port on the parent so every rank agrees on it and
+            # repeated reconstruct() calls in one kernel never collide on a stale 29500.
+            master_port = find_free_port()
+
             # forkserver: workers fork from a clean pre-started server (no inherited
             # CUDA, no Jupyter FDs).  Only plain Python scalars/strings cross the
             # process boundary, so tensor pickling is never triggered.
@@ -566,6 +613,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
                 devices,
                 recon_kwargs,
                 result_path,
+                master_port,
             )
             result = torch.load(result_path, map_location="cpu", weights_only=False)
 
@@ -639,6 +687,13 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
                 self._iter_lrs[k].extend(list(v)[n_before:])
             self._iter_recon_types.extend(result.get("iter_recon_types", [])[n_before:])
 
+        # --- snapshots (recorded on rank 0 in the worker) ---
+        # reset=True gave the worker a fresh _snapshots list; reset=False inherited the old
+        # ones and appended, so the returned list is already the full history either way.
+        returned_snapshots = result.get("snapshots")
+        if returned_snapshots is not None:
+            self._snapshots = returned_snapshots
+
         self._multi_gpu_devices = devices
         return self
 
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 532b4d5da..fa570bc4b 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -219,7 +219,7 @@ def _get_probe_overlap(self, max_batch_size: int | None = None) -> np.ndarray:
         num_dps = self.dset.num_positions
         shifted_probes = prb.expand(num_dps, *self.roi_shape)
 
-        batch_size = num_dps if max_batch_size is None else int(max_batch_size)
+        batch_size = min(num_dps, 4096) if max_batch_size is None else int(max_batch_size)
         probe_overlap = torch.zeros(
             tuple(self.obj_shape_full[-2:]), dtype=self._dtype_real, device=self._single_device
         )
@@ -995,12 +995,26 @@ def _build_dataloaders(
         ``__getitem__`` returns ``{"index": idx, ...}`` for the original dataset index, and
         ``Subset[i]`` calls ``dataset[indices[i]]``, so ``batch["index"]`` is the original
         dataset index under either branch.
+
+        ``self.batch_size`` is the GLOBAL batch: the number of samples contributing to one
+        optimizer step across all ranks. Each rank's DataLoader draws ``batch_size //
+        world_size``, so the same ``batch_size`` gives the same optimization trajectory (and
+        loss curve) on any GPU count.
         """
         pin_memory = self.dset.target_residency == "cpu" and str(self._single_device).startswith(
             "cuda"
         )
+        per_rank_batch = self.batch_size
+        if world_size > 1:
+            per_rank_batch = max(1, self.batch_size // world_size)
+            if self.batch_size % world_size != 0:
+                warn(
+                    f"batch_size={self.batch_size} is not divisible by world_size={world_size}; "
+                    f"each rank uses {per_rank_batch}, so the effective global batch is "
+                    f"{per_rank_batch * world_size}."
+                )
         loader_kwargs: dict[str, Any] = {
-            "batch_size": self.batch_size,
+            "batch_size": per_rank_batch,
             "num_workers": num_workers,
             "pin_memory": pin_memory,
             "drop_last": False,

From 571f4f65be2f508eba879216830e51c1a55084aa Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Tue, 21 Jul 2026 19:16:52 -0700
Subject: [PATCH 56/59] memory optimizations for preprocessing

---
 .../diffractive_imaging/dataset_models.py     | 154 +++++++++++-------
 .../diffractive_imaging/ptychography_base.py  |   8 +-
 2 files changed, 103 insertions(+), 59 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index f7d0fbe81..fd646fb2b 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -165,12 +165,11 @@ def __init__(
         self._initial_scan_positions_px = torch.zeros_like(self._scan_positions_px)
         self._initial_descan_shifts = torch.zeros_like(self._descan_shifts)
 
-        # _targets is a plain attribute (NOT a registered buffer) so that its device can be
-        # managed explicitly per target_residency; AutoSerialize does not serialize it either way.
-        self._targets = torch.zeros(self.num_gpts, *self.roi_shape)
-        self.register_buffer(
-            "_patch_indices", torch.zeros(self.num_gpts, *self.roi_shape, dtype=torch.int32)
-        )
+        # _targets is a plain attribute so that its device can be managed explicitly per
+        # target_residency. Initialized empty (rather than (num_gpts, *roi_shape) buffer) to avoid
+        # a redundant allocation during preprocessing
+        self._targets = torch.zeros(0, dtype=getattr(torch, config.get("dtype_real")))
+        self.register_buffer("_patch_indices", torch.zeros(0, dtype=torch.int32))
         self.register_buffer("_last_patch_positions_px", torch.zeros(self.num_gpts, 2))
         self.register_buffer("_detector_mask", torch.ones(*self.roi_shape))
         self.positions_mask = torch.ones(self.num_gpts, dtype=torch.bool)
@@ -182,6 +181,17 @@ def __init__(
         # fractional positions), so the probe is not subpixel-shifted.
         self._implicit_object = False
 
+        # Measurement arrays. Only ``centered_amplitudes`` is materialized during normalization;
+        self._centered_amplitudes: torch.Tensor | None = None
+        self._amplitudes: torch.Tensor | None = None
+        self._intensities: torch.Tensor | None = None
+        self._centered_intensities: torch.Tensor | None = None
+        self._pattern_crop_mask: np.ndarray | None = None
+        self._pattern_crop_mask_shape: tuple[int, int] = (
+            int(self.roi_shape[0]),
+            int(self.roi_shape[1]),
+        )
+
     def get_optimization_parameters(self) -> "dict[str, list[torch.Tensor]]":
         """Descan and scan-position parameters as separate PPLR groups.
 
@@ -377,7 +387,7 @@ def _set_targets(self, target_space: Literal["amplitude", "intensity"]):
             raise ValueError(
                 f"target_space must be 'amplitude' or 'intensity', got {target_space!r}"
             )
-        self._targets = source.clone().to(target_device)
+        self._targets = source.to(target_device)
 
     @property
     def patch_indices(self) -> torch.Tensor:
@@ -416,6 +426,8 @@ def centered_amplitudes(self) -> torch.Tensor:
         """gives the amplitudes that have had descan corrected and which are centered in the fov
         shaped as (rr*rc, qx, qy)
         """
+        if self._centered_amplitudes is None:
+            raise ValueError("centered_amplitudes is unset; preprocess the dataset first")
         return self._centered_amplitudes
 
     @centered_amplitudes.setter
@@ -429,10 +441,34 @@ def centered_amplitudes(self, arr: "np.ndarray | torch.Tensor") -> None:
         )
         self._centered_amplitudes = arr
 
+    def _recompute_raw_measurement(self, amplitude: bool) -> torch.Tensor:
+        """
+        Recompute the raw (un-centered) masked ``amplitudes``/``intensities`` from
+        ``intensities_4d`` on demand. Values are identical to what
+        ``_normalize_diffraction_intensities`` used to store; kept lazy so the full-size arrays are
+        not held resident.
+
+        Typically will be called only when changing the targets (using ``set_targets``).
+        """
+        dtype = config.get("dtype_real")
+        mask = self.positions_mask.detach().cpu().numpy().ravel()
+        roi = tuple(int(x) for x in self.roi_shape)
+        arr = np.maximum(
+            np.asarray(self.intensities_4d).reshape(self.num_gpts, *roi).astype(dtype), 0
+        )
+        if amplitude:
+            arr = np.sqrt(arr)
+        arr = arr[mask]
+        if self._pattern_crop_mask is not None:
+            arr = arr[:, self._pattern_crop_mask].reshape((-1, *self._pattern_crop_mask_shape))
+        return torch.as_tensor(np.ascontiguousarray(arr), dtype=getattr(torch, dtype))
+
     @property
     def amplitudes(self) -> torch.Tensor:
         """raw intensities converted to amplitudes, as a torch tensor"""
-        return self._amplitudes
+        if self._amplitudes is not None:
+            return self._amplitudes
+        return self._recompute_raw_measurement(amplitude=True)
 
     @amplitudes.setter
     def amplitudes(self, arr: "np.ndarray | torch.Tensor") -> None:
@@ -450,7 +486,9 @@ def centered_intensities(self) -> torch.Tensor:
         """intensities that have had descan corrected and which are centered in the fov
         shaped as (rr*rc, qx, qy)
         """
-        return self._centered_intensities
+        if self._centered_intensities is not None:
+            return self._centered_intensities
+        return self.centered_amplitudes**2
 
     @centered_intensities.setter
     def centered_intensities(self, arr: "np.ndarray | torch.Tensor") -> None:
@@ -466,7 +504,9 @@ def centered_intensities(self, arr: "np.ndarray | torch.Tensor") -> None:
     @property
     def intensities(self) -> torch.Tensor:
         """raw intensities as a torch tensor"""
-        return self._intensities
+        if self._intensities is not None:
+            return self._intensities
+        return self._recompute_raw_measurement(amplitude=False)
 
     @intensities.setter
     def intensities(self, arr: "np.ndarray | torch.Tensor") -> None:
@@ -648,7 +688,7 @@ def _obj_shape_crop_2d(self) -> np.ndarray:
         shp = np.floor(self.fov / self.obj_sampling)
         shp += shp % 2
         shp = shp.astype("int")
-        if self.com_transpose: # match the axis-flipped scan positions
+        if self.com_transpose:  # match the axis-flipped scan positions
             shp = shp[::-1]
         return shp
 
@@ -719,7 +759,11 @@ def _set_patch_indices(self, obj_padding_px: np.ndarray | tuple) -> None:
 
         # Process positions in chunks to reduce memory usage
         chunk_size = min(1000, len(r0))
-        patch_indices_list = []
+        patch_indices = torch.empty(
+            (len(r0), int(self.roi_shape[0]), int(self.roi_shape[1])),
+            dtype=torch.int32,
+            device=r0.device,
+        )
 
         for i in range(0, len(r0), chunk_size):
             end_idx = min(i + chunk_size, len(r0))
@@ -729,10 +773,9 @@ def _set_patch_indices(self, obj_padding_px: np.ndarray | tuple) -> None:
             row_chunk = (r0_chunk[:, None, None] + x_ind[None, :, None]) % obj_shape[-2]
             col_chunk = (c0_chunk[:, None, None] + y_ind[None, None, :]) % obj_shape[-1]
 
-            patch_indices_chunk = (row_chunk * obj_shape[-1] + col_chunk).type(torch.int32)
-            patch_indices_list.append(patch_indices_chunk)
+            patch_indices[i:end_idx] = (row_chunk * obj_shape[-1] + col_chunk).type(torch.int32)
 
-        self._patch_indices = torch.cat(patch_indices_list, dim=0)
+        self._patch_indices = patch_indices
         self._last_patch_positions_px = self.scan_positions_px.detach().clone()
 
     def patch_indices_need_update(self) -> bool:
@@ -1014,8 +1057,10 @@ def intensities_4d(self) -> np.ndarray:
 
     @intensities_4d.setter
     def intensities_4d(self, intensities: np.ndarray) -> None:
+        # Keep raw detector counts in their native dtype
+        dtype = validate_array(intensities, name="intensities_4d", ndim=4).dtype
         self._intensities_4d = validate_array(
-            intensities, name="intensities_4d", ndim=4, dtype=config.get("dtype_real")
+            intensities, name="intensities_4d", ndim=4, dtype=dtype
         )
 
     @property
@@ -1620,7 +1665,8 @@ def _normalize_diffraction_intensities(
         bilinear: bool = False,
     ):
         dtype = config.get("dtype_real")
-        diff_intensities = self.intensities_4d.astype(dtype)
+        np_dtype = np.dtype(dtype)
+        raw = self._intensities_4d
         com_fit = self.com_fit
         if positions_mask is not None:
             self.positions_mask = positions_mask
@@ -1628,13 +1674,9 @@ def _normalize_diffraction_intensities(
 
         # Aggressive cropping for when off-centered high scattering angle data was recorded
         if crop_patterns:
-            crop_r = int(
-                np.minimum(diff_intensities.shape[2] - com_fit[0].max(), com_fit[0].min())
-            )
-            crop_c = int(
-                np.minimum(diff_intensities.shape[3] - com_fit[1].max(), com_fit[1].min())
-            )
-            crop_m = np.minimum(crop_c, crop_r)
+            crop_r = int(np.minimum(raw.shape[2] - com_fit[0].max(), com_fit[0].min()))
+            crop_c = int(np.minimum(raw.shape[3] - com_fit[1].max(), com_fit[1].min()))
+            crop_m = int(np.minimum(crop_c, crop_r))
 
             pattern_crop_mask = np.zeros(self.roi_shape, dtype="bool")
             pattern_crop_mask[:crop_m, :crop_m] = True
@@ -1645,16 +1687,20 @@ def _normalize_diffraction_intensities(
 
         else:
             pattern_crop_mask = None
-            pattern_crop_mask_shape = self.roi_shape
+            pattern_crop_mask_shape = (int(self.roi_shape[0]), int(self.roi_shape[1]))
+
+        # Accumulate into a numpy buffer and convert once at the end
+        n_out = int(positions_mask_2d.sum())
+        centered_amplitudes = np.zeros((n_out, *pattern_crop_mask_shape), dtype=np_dtype)
 
         mean_intensity = 0
         mean_amplitude = 0
-        centered_amplitudes = np.zeros(diff_intensities.shape, dtype=dtype)
-        amplitudes = np.zeros(diff_intensities.shape, dtype=dtype)
-        intensities = np.zeros(diff_intensities.shape, dtype=dtype)
+        out_i = 0
+        current_Rr = -1
+        row = None
         for Rr, Rc in tqdmnd(
-            range(diff_intensities.shape[0]),
-            range(diff_intensities.shape[1]),
+            range(raw.shape[0]),
+            range(raw.shape[1]),
             desc="Normalizing intensities",
             unit="probe position",
             disable=not self._verbose,
@@ -1662,14 +1708,18 @@ def _normalize_diffraction_intensities(
             if not positions_mask_2d[Rr, Rc]:
                 continue
 
-            intensity = np.maximum(diff_intensities[Rr, Rc], 0)
-            intensities[Rr, Rc] = intensity
+            # Cast to float32 one scan-row at a time--avoids a full 4D float32 copy
+            if Rr != current_Rr:
+                current_Rr = Rr
+                row = raw[Rr]
+                if row.dtype != np_dtype:
+                    row = row.astype(dtype)
+            intensity = np.maximum(row[Rc], 0)
             mean_intensity += np.sum(intensity)
+            
             ### shifting amplitude rather than intensity to minimize ringing artifacts
             amplitude = np.maximum(np.sqrt(intensity), 0)
             mean_amplitude += np.sum(amplitude)
-            amplitudes[Rr, Rc] = amplitude
-
             shift_amplitude = shift_array(  # shifting to 0,0 then fftshift
                 amplitude,
                 -(com_fit[0, Rr, Rc] + 0.0),
@@ -1678,32 +1728,22 @@ def _normalize_diffraction_intensities(
             )
             shift_amplitude = np.maximum(shift_amplitude, 0)
             shift_amplitude = np.fft.fftshift(shift_amplitude)
+            if pattern_crop_mask is not None:
+                shift_amplitude = shift_amplitude[pattern_crop_mask].reshape(
+                    pattern_crop_mask_shape
+                )
 
-            centered_amplitudes[Rr, Rc] = shift_amplitude
-
-        # The source 4D copy is no longer needed; free it before the masking copies below 
-        del diff_intensities
-
-        amplitudes = amplitudes[positions_mask_2d]
-        centered_amplitudes = centered_amplitudes[positions_mask_2d]
-        intensities = intensities[positions_mask_2d]
-
-        if crop_patterns:
-            amplitudes = amplitudes[:, pattern_crop_mask].reshape((-1, *pattern_crop_mask_shape))
-            centered_amplitudes = centered_amplitudes[:, pattern_crop_mask].reshape(
-                (-1, *pattern_crop_mask_shape)
-            )
-            intensities = intensities[:, pattern_crop_mask].reshape((-1, *pattern_crop_mask_shape))
-
-        centered_intensities = centered_amplitudes**2
+            centered_amplitudes[out_i] = shift_amplitude
+            out_i += 1
 
-        mean_intensity /= amplitudes.shape[0]
-        mean_amplitude /= amplitudes.shape[0]
+        mean_intensity /= n_out
+        mean_amplitude /= n_out
 
-        self.centered_amplitudes = centered_amplitudes
-        self.amplitudes = amplitudes
-        self.centered_intensities = centered_intensities
-        self.intensities = intensities
+        # Wrap the numpy buffer zero-copy (``torch.from_numpy`` shares storage)
+        self._centered_amplitudes = torch.from_numpy(centered_amplitudes)
+        self._amplitudes = None
+        self._intensities = None
+        self._centered_intensities = None
         descan_shifts = -1 * com_fit.reshape((2, -1))  # (2, rr*rc)
         descan_shifts += self.roi_shape[:, None] / 2
         descan_shifts = descan_shifts.T  # (rr*rc, 2)
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index fa570bc4b..22f7ae1d4 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -133,9 +133,13 @@ def __init__(  # TODO prevent direct instantiation
         if (
             isinstance(probe_model, ProbePixelated)
             and (probe_model.vacuum_probe_intensity is not None)
-            and (dset.amplitudes.shape[1:] != probe_model.vacuum_probe_intensity.shape)
+            # ``centered_amplitudes`` shares amplitudes' shape but is always resident (amplitudes is
+            # recomputed lazily), so use it here to avoid materializing the full raw array.
+            and (dset.centered_amplitudes.shape[1:] != probe_model.vacuum_probe_intensity.shape)
         ):
-            probe_model.rescale_vacuum_probe((dset.amplitudes.shape[1], dset.amplitudes.shape[2]))
+            probe_model.rescale_vacuum_probe(
+                (dset.centered_amplitudes.shape[1], dset.centered_amplitudes.shape[2])
+            )
 
         # Remove centralized optimizer storage - now managed by individual models
         self.probe_model = probe_model

From 19f85de97d0ffbc173949fa7a0119eb3c150c51b Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Fri, 31 Jul 2026 12:40:23 -0700
Subject: [PATCH 57/59] big cleanup of iterative ptycho code, simplifications
 and bugfixes

---
 .../diffractive_imaging/dataset_models.py     |  18 +-
 .../logger_ptychography.py                    |   6 +-
 .../diffractive_imaging/object_models.py      |  44 +--
 .../optimize_hyperparameters.py               | 262 +++++++++---------
 .../diffractive_imaging/probe_models.py       |  91 ++----
 .../diffractive_imaging/ptycho_losses.py      |   7 +-
 .../diffractive_imaging/ptycho_utils.py       |  86 +-----
 .../diffractive_imaging/ptychography.py       |  31 ++-
 .../diffractive_imaging/ptychography_base.py  |  36 +--
 .../diffractive_imaging/ptychography_lite.py  |  55 ++--
 .../ptychography_visualizations.py            | 100 +++----
 11 files changed, 306 insertions(+), 430 deletions(-)

diff --git a/src/quantem/diffractive_imaging/dataset_models.py b/src/quantem/diffractive_imaging/dataset_models.py
index fd646fb2b..a7032df18 100644
--- a/src/quantem/diffractive_imaging/dataset_models.py
+++ b/src/quantem/diffractive_imaging/dataset_models.py
@@ -145,7 +145,7 @@ def __init__(
         self._preprocessed = False
         self._preprocessing_params = {}  # for serialization and reloading
         self._com_rotation_rad = 0  # default
-        self._com_transpose = False  # default
+        self._transpose = False  # default
 
         # scan_positions_px: [num_positions, 2] in pixels
         self._scan_positions_px = nn.Parameter(
@@ -1269,10 +1269,10 @@ def preprocess(
             self.detector_mask = torch.nn.functional.pad(
                 self.detector_mask,
                 (
-                    self.diffraction_padding[0],
-                    self.diffraction_padding[0],
                     self.diffraction_padding[1],
                     self.diffraction_padding[1],
+                    self.diffraction_padding[0],
+                    self.diffraction_padding[0],
                 ),
                 mode="constant",
                 value=0,
@@ -1372,10 +1372,14 @@ def _set_intensities_com(
             ):
                 masked_intensity = intensities[Rr, Rc]
                 if dp_mask is not None:
-                    masked_intensity *= dp_mask
+                    masked_intensity = masked_intensity * dp_mask
                 summed_intensity = masked_intensity.sum()
-                com_measured_r[Rr, Rc] = np.sum(masked_intensity * kcm) / summed_intensity
-                com_measured_c[Rr, Rc] = np.sum(masked_intensity * krm) / summed_intensity
+                if summed_intensity == 0:
+                    com_measured_r[Rr, Rc] = np.nan
+                    com_measured_c[Rr, Rc] = np.nan
+                else:
+                    com_measured_r[Rr, Rc] = np.sum(masked_intensity * krm) / summed_intensity
+                    com_measured_c[Rr, Rc] = np.sum(masked_intensity * kcm) / summed_intensity
 
         if fit_function == "none":
             com_fit_r, com_fit_c = com_measured_r, com_measured_c
@@ -1716,7 +1720,7 @@ def _normalize_diffraction_intensities(
                     row = row.astype(dtype)
             intensity = np.maximum(row[Rc], 0)
             mean_intensity += np.sum(intensity)
-            
+
             ### shifting amplitude rather than intensity to minimize ringing artifacts
             amplitude = np.maximum(np.sqrt(intensity), 0)
             mean_amplitude += np.sum(amplitude)
diff --git a/src/quantem/diffractive_imaging/logger_ptychography.py b/src/quantem/diffractive_imaging/logger_ptychography.py
index 9694e3f6b..5f06a23e0 100644
--- a/src/quantem/diffractive_imaging/logger_ptychography.py
+++ b/src/quantem/diffractive_imaging/logger_ptychography.py
@@ -52,7 +52,7 @@ def object_image(self, volume_obj: ObjectModelType, iter: int, logger_cmap: str
             elif obj_type == "pure_phase":
                 self.log_image(
                     tag="object/phase_zsum",
-                    image=np.angle(obj_sum),
+                    image=obj_sum,
                     step=iter,
                     cmap=self._phase_cmap,
                 )
@@ -74,7 +74,7 @@ def object_image(self, volume_obj: ObjectModelType, iter: int, logger_cmap: str
             print(f"Warning: Failed to log object images at iteration {iter}: {e}")
 
     def probe_image(self, probe_model: ProbeModelType, iter: int, logger_cmap: str = "turbo"):
-        """Log probe images showing both real-space and fourier-space representations (optimized)."""
+        """Log real-space probe images (optimized)."""
         try:
             probe = probe_model.probe
 
@@ -160,7 +160,7 @@ def log_iter(
                     self.log_scalar(f"learning_rate/{param_name}", float(lr_value), iter)
 
             # Images (only when needed)
-            if iter % self.log_images_every == 0 and self.log_images_every > 0:
+            if self.log_images_every > 0 and iter % self.log_images_every == 0:
                 self.object_image(object_model, iter, logger_cmap)
                 if self._log_probe_images:
                     self.probe_image(probe_model, iter, logger_cmap)
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index e0e6eab8b..23783f1a3 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -2,7 +2,7 @@
 from abc import abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Callable, Literal, Self, Sequence, cast
+from typing import Callable, Literal, Sequence, cast
 from warnings import warn
 
 import matplotlib.pyplot as plt
@@ -51,7 +51,8 @@ class PtychoObjConstraintParams:
         Constraints for grid-based object representations (``ObjectPixelated`` and
         ``ObjectDIP`` share this set today).
     INR
-        Placeholder for the upcoming implicit-neural-representation object.
+        Constraints for the implicit-neural-representation objects (``ObjectINR``,
+        ``ObjectTensorDecomp``).
 
     Examples
     --------
@@ -76,12 +77,14 @@ class Raster(Constraints):
             ``"pure_phase"`` the amplitude is clamped to ``[0, 1]`` (or fixed to 1)
             regardless of this flag.
         positivity_mode: Literal["clamp", "shrink"], default ``"clamp"``
-            How to enforce positivity. "clamp" clamps the object to be non-negative after each 
+            How to enforce positivity. "clamp" clamps the object to be non-negative after each
             update, does not move the parameter, only how it is shown/used.
-            "shrink" subtracts a background offset from the object so background regions sit at 
-            zero, is applied to the parameter after the update step. 
-            If an FOV mask is set the offset is the mean of the background 
-            (``mask < 0.5 * mask.max()``); otherwise it's ``obj.min()``.
+            "shrink" subtracts a per-slice background offset from the object so background regions
+            sit at zero, is applied to the parameter after the update step.
+            If an FOV mask is set the offset is the per-slice mean of the background
+            (``mask < 0.5 * mask.max()``); otherwise it's the per-slice 10th percentile. The
+            offset is scaled by ``fix_potential_baseline_factor`` even when
+            ``fix_potential_baseline`` is False.
         fix_potential_baseline : bool, default ``False``
             ``obj_type="potential"`` only. Subtracts an offset from the object so
             background regions sit at zero. If an FOV mask is set the offset is
@@ -475,7 +478,8 @@ def _get_obj_patches(self, obj_array, patch_indices):
 
     def backward(self, *args, **kwargs):
         raise NotImplementedError(
-            f"Analytical gradients are not implemented for {Self}, use autograd=True"
+            f"Analytical gradients are not implemented for {type(self).__name__}, "
+            "use autograd=True"
         )
 
 
@@ -908,8 +912,15 @@ def from_array(
             rng=rng,
             _token=cls._token,
         )
-        obj_model._initial_obj = torch.tensor(
-            initial_obj, dtype=obj_model.dtype, device=obj_model.device
+        initial = torch.as_tensor(initial_obj)
+        if initial.is_complex() and obj_type != "complex":
+            if obj_type == "pure_phase":
+                # Convert legacy complex initial_obj (amp*exp(1j*phase)) to bare phase
+                initial = initial.angle()
+            else:
+                raise ValueError(f"Complex initial_obj is not valid for obj_type '{obj_type}'")
+        obj_model._initial_obj = (
+            initial.clone().detach().to(dtype=obj_model.dtype, device=obj_model.device)
         )
 
         return obj_model
@@ -987,9 +998,6 @@ def _initialize_obj(
                 arr = ph
         elif self._initialize_mode == "array":
             arr = self._initial_obj
-            if self.obj_type == "pure_phase" and arr.is_complex():
-                # Convert legacy complex initial_obj (amp*exp(1j*phase)) to bare phase
-                arr = arr.angle()
         else:
             raise ValueError(f"Invalid initialize mode: {self._initialize_mode}")
 
@@ -1086,8 +1094,6 @@ def __init__(
         self._pretrain_losses = []
         self._pretrain_lrs = []
         self._model_input_noise_std = input_noise_std
-        self._model_input = torch.tensor([])
-        self._pretrain_target = torch.tensor([])
 
     @classmethod
     def from_model(
@@ -1173,9 +1179,9 @@ def dtype(self) -> "torch.dtype":
             return getattr(self.model, "dtype")
         else:
             if self.obj_type in ["complex"]:
-                return config.get("dtype_complex")
+                return getattr(torch, config.get("dtype_complex"))
             else:
-                return config.get("dtype_real")
+                return getattr(torch, config.get("dtype_real"))
 
     @property
     def model(self) -> "torch.nn.Module":
@@ -1760,8 +1766,8 @@ def from_pixelated(
         Pass ``model`` to wrap a custom INR ``nn.Module`` directly (mapping ``(N, 3)`` coords to
         ``(N, 1)``), as with ``ObjectDIP.from_pixelated`` -- handy for testing architectures. When
         ``model`` is ``None`` a default zero-initialized ``HSiren`` is built from the
-        ``hidden_features`` / ``hidden_layers`` / ``omega_0`` args (with a positivity activation
-        for ``potential``); when a ``model`` is given those args and the activation are its own.
+        ``hidden_features`` / ``hidden_layers`` / ``omega_0`` args (with identity output
+        activations); when a ``model`` is given those args and the activation are its own.
         """
         if not (
             isinstance(pixelated, ObjectPixelated) or "ObjectPixelated" in str(type(pixelated))
diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
index 28155a073..99904e203 100644
--- a/src/quantem/diffractive_imaging/optimize_hyperparameters.py
+++ b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
@@ -2,6 +2,7 @@
 
 import copy
 import gc
+import inspect
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Mapping, Optional
 
@@ -90,7 +91,13 @@ def replace_recursive(obj, path=()):
             return type(obj)(replace_recursive(v, (*path, i)) for i, v in enumerate(obj))
         return obj
 
-    return replace_recursive(config)
+    # trial params are named relative to each sub-config, not to the whole config
+    updated = dict(config)
+    for key in ("base_kwargs", "dataset_kwargs", "dataset_preprocess_kwargs"):
+        sub_config = updated.get(key)
+        if sub_config is not None:
+            updated[key] = replace_recursive(sub_config)
+    return updated
 
 
 def _is_dataset_param(param_path):
@@ -210,16 +217,22 @@ def _clone_ptychography_dataset(dset: PtychographyDatasetBase) -> PtychographyDa
     cloned.com_measured = dset.com_measured.copy()
     cloned.com_fit = dset.com_fit.copy()
     cloned.centered_amplitudes = dset.centered_amplitudes.detach().cpu().clone()
-    cloned.amplitudes = dset.amplitudes.detach().cpu().clone()
-    cloned.centered_intensities = dset.centered_intensities.detach().cpu().clone()
-    cloned.intensities = dset.intensities.detach().cpu().clone()
+    # amplitudes / intensities / centered_intensities are derived on demand from
+    # intensities_4d; only carry them over if the source has them materialized
+    for attr in ("_amplitudes", "_intensities", "_centered_intensities"):
+        source = getattr(dset, attr, None)
+        if source is not None:
+            setattr(cloned, attr[1:], source.detach().cpu().clone())
     cloned.detector_mask = detector_mask
     cloned.mean_diffraction_intensity = dset.mean_diffraction_intensity
     if hasattr(dset, "mean_diffraction_amplitude"):
         cloned.mean_diffraction_amplitude = dset.mean_diffraction_amplitude
     cloned._pattern_crop_mask = copy.deepcopy(getattr(dset, "_pattern_crop_mask", None))
-    cloned._pattern_crop_mask_shape = copy.deepcopy(
-        getattr(dset, "_pattern_crop_mask_shape", dset.roi_shape)
+    mask_shape = getattr(dset, "_pattern_crop_mask_shape", None)
+    cloned._pattern_crop_mask_shape = (
+        copy.deepcopy(mask_shape)
+        if mask_shape is not None
+        else (int(dset.roi_shape[0]), int(dset.roi_shape[1]))
     )
     cloned.initial_descan_shifts = dset.initial_descan_shifts.detach().cpu().clone()
     cloned.initial_scan_positions_px = dset.initial_scan_positions_px.detach().cpu().clone()
@@ -232,7 +245,7 @@ def _clone_ptychography_dataset(dset: PtychographyDatasetBase) -> PtychographyDa
     return cloned
 
 
-def _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type):
+def _run_reconstruction_pipeline(recon_obj, resolved_kwargs):
     """Run the reconstruction pipeline for either class."""
     # Preprocess step
     preprocess_kwargs = resolved_kwargs.get("preprocess")
@@ -240,9 +253,13 @@ def _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type):
         recon_obj.preprocess(**preprocess_kwargs)
 
     # Reconstruct step
-    reconstruct_kwargs = resolved_kwargs.get("reconstruct", {})
-    reconstruct_kwargs["verbose"] = False
+    recon_obj.verbose = False
+    reconstruct_kwargs = resolved_kwargs.get("reconstruct")
     if reconstruct_kwargs:
+        reconstruct_kwargs = dict(reconstruct_kwargs)
+        # only PtychoLite.reconstruct takes verbose, and it resets recon_obj.verbose
+        if "verbose" in inspect.signature(recon_obj.reconstruct).parameters:
+            reconstruct_kwargs.setdefault("verbose", False)
         recon_obj.reconstruct(**reconstruct_kwargs)
 
 
@@ -314,7 +331,7 @@ def objective(trial: optuna.trial.Trial) -> float:
             recon_obj = _build_ptychography_instance(constructors, resolved_kwargs)
 
         # 5) Run the reconstruction pipeline
-        _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type)
+        _run_reconstruction_pipeline(recon_obj, resolved_kwargs)
 
         # 6) Extract loss
         if loss_getter is not None:
@@ -347,7 +364,7 @@ def __init__(
         self.study_kwargs = study_kwargs or {}
         self.unit = unit
         self.verbose = verbose
-        self._config = None
+        self._config: Dict[str, Any] | None = None
         self.study = optuna.create_study(direction=direction, **self.study_kwargs)
 
     @classmethod
@@ -438,32 +455,34 @@ def optimize(self) -> "OptimizePtychography":
         if hasattr(self, "_config") and self._config:
             self.study.set_user_attr("config", self._config)
 
-        if not self.verbose:
-            optuna.logging.set_verbosity(optuna.logging.WARNING)
-        else:
-            optuna.logging.set_verbosity(optuna.logging.INFO)
+        prev_verbosity = optuna.logging.get_verbosity()
+        optuna.logging.set_verbosity(
+            optuna.logging.INFO if self.verbose else optuna.logging.WARNING
+        )
 
-        with tqdm(total=self.n_trials, desc="optimizing", unit=self.unit) as pbar:
+        try:
+            with tqdm(total=self.n_trials, desc="optimizing", unit=self.unit) as pbar:
 
-            def _on_trial_end(study_: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
-                pbar.update(1)
+                def _on_trial_end(
+                    study_: optuna.study.Study, trial: optuna.trial.FrozenTrial
+                ) -> None:
+                    pbar.update(1)
 
-                torch.cuda.empty_cache()
-                gc.collect()
-
-            self.study.optimize(
-                self.objective_func,
-                n_trials=self.n_trials,
-                callbacks=[_on_trial_end],
-                show_progress_bar=self.verbose,
-            )
+                    torch.cuda.empty_cache()
+                    gc.collect()
 
-        if not self.verbose:
-            optuna.logging.set_verbosity(optuna.logging.INFO)
+                self.study.optimize(
+                    self.objective_func,
+                    n_trials=self.n_trials,
+                    callbacks=[_on_trial_end],
+                    show_progress_bar=False,
+                )
+        finally:
+            optuna.logging.set_verbosity(prev_verbosity)
 
         return self
 
-    def visualize(self, figsize=(10, 6)):
+    def visualize(self, figsize=None):
         """Visualize optimization results showing parameter values vs loss."""
         if not self.study.trials:
             raise RuntimeError("No trials to plot. Run optimize() first.")
@@ -473,20 +492,22 @@ def visualize(self, figsize=(10, 6)):
         if not trials:
             raise RuntimeError("No completed trials to plot.")
 
-        param_names = list(trials[0].params.keys())
+        # trials may not all sample the same parameters, so take the union
+        param_names = list(dict.fromkeys(name for trial in trials for name in trial.params))
         best_trial = self.study.best_trial
         best_value = best_trial.value
 
         # Special case: 2 parameters - add 2D scatter plot
         if len(param_names) == 2:
-            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+            fig, axes = plt.subplots(1, 3, figsize=figsize or (15, 5))
 
             ax_2d = axes[0]
             param1, param2 = param_names
 
-            param1_values = np.array([trial.params[param1] for trial in trials])
-            param2_values = np.array([trial.params[param2] for trial in trials])
-            losses = np.array([trial.value for trial in trials])
+            pair_trials = [t for t in trials if param1 in t.params and param2 in t.params]
+            param1_values = np.array([trial.params[param1] for trial in pair_trials])
+            param2_values = np.array([trial.params[param2] for trial in pair_trials])
+            losses = np.array([trial.value for trial in pair_trials])
 
             scatter = ax_2d.scatter(
                 param1_values,
@@ -500,18 +521,19 @@ def visualize(self, figsize=(10, 6)):
             )
 
             # Highlight best trial
-            best_param1 = best_trial.params[param1]
-            best_param2 = best_trial.params[param2]
-            ax_2d.scatter(
-                [best_param1],
-                [best_param2],
-                color="red",
-                s=300,
-                marker="*",
-                edgecolors="black",
-                linewidth=2,
-                zorder=5,
-            )
+            best_param1 = best_trial.params.get(param1)
+            best_param2 = best_trial.params.get(param2)
+            if best_param1 is not None and best_param2 is not None:
+                ax_2d.scatter(
+                    [best_param1],
+                    [best_param2],
+                    color="red",
+                    s=300,
+                    marker="*",
+                    edgecolors="black",
+                    linewidth=2,
+                    zorder=5,
+                )
 
             # Colorbar
             cbar = plt.colorbar(scatter, ax=ax_2d)
@@ -530,8 +552,9 @@ def visualize(self, figsize=(10, 6)):
                 ax = axes[idx + 1]
 
                 # Extract data
-                param_values = np.array([trial.params[param_name] for trial in trials])
-                losses = np.array([trial.value for trial in trials])
+                param_trials = [t for t in trials if param_name in t.params]
+                param_values = np.array([trial.params[param_name] for trial in param_trials])
+                losses = np.array([trial.value for trial in param_trials])
 
                 # Scatter plot
                 ax.scatter(
@@ -539,20 +562,23 @@ def visualize(self, figsize=(10, 6)):
                 )
 
                 # Highlight best trial
-                best_param_value = best_trial.params[param_name]
-                ax.scatter(
-                    [best_param_value],
-                    [best_value],
-                    color="red",
-                    s=200,
-                    marker="*",
-                    edgecolors="black",
-                    linewidth=1.5,
-                    zorder=5,
-                )
-
-                # Vertical line at optimal parameter value
-                ax.axvline(best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7)
+                best_param_value = best_trial.params.get(param_name)
+                if best_param_value is not None:
+                    ax.scatter(
+                        [best_param_value],
+                        [best_value],
+                        color="red",
+                        s=200,
+                        marker="*",
+                        edgecolors="black",
+                        linewidth=1.5,
+                        zorder=5,
+                    )
+
+                    # Vertical line at optimal parameter value
+                    ax.axvline(
+                        best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7
+                    )
 
                 # Clean up parameter name for label
                 clean_name = param_name.split(".")[-1]
@@ -571,7 +597,7 @@ def visualize(self, figsize=(10, 6)):
         n_cols = min(3, n_params)  # Max 3 columns
         n_rows = (n_params + n_cols - 1) // n_cols  # Ceiling division
 
-        fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize, squeeze=False)
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize or (10, 6), squeeze=False)
         axes = axes.flatten()
 
         # Plot each parameter
@@ -579,27 +605,29 @@ def visualize(self, figsize=(10, 6)):
             ax = axes[idx]
 
             # Extract data
-            param_values = np.array([trial.params[param_name] for trial in trials])
-            losses = np.array([trial.value for trial in trials])
+            param_trials = [t for t in trials if param_name in t.params]
+            param_values = np.array([trial.params[param_name] for trial in param_trials])
+            losses = np.array([trial.value for trial in param_trials])
 
             # Scatter plot
             ax.scatter(param_values, losses, alpha=0.6, s=50, edgecolors="black", linewidth=0.5)
 
             # Highlight best trial
-            best_param_value = best_trial.params[param_name]
-            ax.scatter(
-                [best_param_value],
-                [best_value],
-                color="red",
-                s=200,
-                marker="*",
-                edgecolors="black",
-                linewidth=1.5,
-                zorder=5,
-            )
+            best_param_value = best_trial.params.get(param_name)
+            if best_param_value is not None:
+                ax.scatter(
+                    [best_param_value],
+                    [best_value],
+                    color="red",
+                    s=200,
+                    marker="*",
+                    edgecolors="black",
+                    linewidth=1.5,
+                    zorder=5,
+                )
 
-            # Vertical line at optimal parameter value
-            ax.axvline(best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7)
+                # Vertical line at optimal parameter value
+                ax.axvline(best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7)
 
             # Clean up parameter name for label
             clean_name = param_name.split(".")[-1]
@@ -684,20 +712,8 @@ def grid_search(self, plot_objects=True, figsize=None, return_results=False):
             else:
                 raise ValueError(f"Invalid parameter spec for {param_name}")
 
-            param_names = list(param_grids.keys())
-            all_combinations = list(product(*param_grids.values()))
-
-        def objective_with_capture(trial):
-            """Modified objective that captures the reconstruction object."""
-            # Call the original objective
-            loss = self.objective_func(trial)
-
-            return loss
-
-        # Enqueue all grid points
-        for combo in all_combinations:
-            params = dict(zip(param_names, combo))
-            self.study.enqueue_trial(params)
+        param_names = list(param_grids.keys())
+        all_combinations = list(product(*param_grids.values()))
 
         # Run trials and capture reconstructions
         print("\nRunning reconstructions...")
@@ -724,7 +740,8 @@ def objective_with_capture(trial):
                 gc.collect()
 
         # Find best
-        best_idx = np.argmin([r["loss"] for r in results])
+        argfn = np.argmax if self.direction == "maximize" else np.argmin
+        best_idx = argfn([r["loss"] for r in results])
         best_result = results[best_idx]
 
         # Plot objects
@@ -749,44 +766,33 @@ def _run_reconstruction_with_params(self, params):
         """
         from quantem.diffractive_imaging.optimize_hyperparameters import _resolve_params_with_trial
 
-        # Create a mock trial that returns our fixed parameters
-        class FixedTrial:
-            def __init__(self, fixed_params):
-                self.params = fixed_params
-                self.number = 0
-
-            def suggest_float(self, name, low, high, **kwargs):
-                return self.params.get(name, (low + high) / 2)
-
-            def suggest_int(self, name, low, high, **kwargs):
-                return int(self.params.get(name, (low + high) // 2))
+        trial = optuna.trial.FixedTrial(params)
 
-            def suggest_categorical(self, name, choices):
-                return self.params.get(name, choices[0])
-
-        trial = FixedTrial(params)
+        config = self._config
+        if config is None:
+            raise RuntimeError("Optimizer is not configured; use a factory method first.")
 
         # Resolve parameters
-        resolved_kwargs = _resolve_params_with_trial(trial, self._config["base_kwargs"])
+        resolved_kwargs = _resolve_params_with_trial(trial, config["base_kwargs"])
 
         # Handle dataset construction if needed
-        if self._config.get("dataset_constructor") is not None:
+        if config.get("dataset_constructor") is not None:
             resolved_dataset_kwargs = _resolve_params_with_trial(
-                trial, self._config.get("dataset_kwargs", {})
+                trial, config.get("dataset_kwargs", {})
             )
-            pdset = self._config["dataset_constructor"](**resolved_dataset_kwargs)
+            pdset = config["dataset_constructor"](**resolved_dataset_kwargs)
 
-            if self._config.get("dataset_preprocess_kwargs") is not None:
+            if config.get("dataset_preprocess_kwargs") is not None:
                 resolved_preprocess_kwargs = _resolve_params_with_trial(
-                    trial, self._config["dataset_preprocess_kwargs"]
+                    trial, config["dataset_preprocess_kwargs"]
                 )
                 pdset.preprocess(**resolved_preprocess_kwargs)
 
             resolved_kwargs.setdefault("init", {})["dset"] = pdset
 
         # Determine reconstruction class
-        reconstruction_class = self._config.get("reconstruction_class", "auto")
-        constructors = self._config["constructors"]
+        reconstruction_class = config.get("reconstruction_class", "auto")
+        constructors = config["constructors"]
 
         if reconstruction_class == "auto":
             main_constructor = constructors.get("ptychography_class")
@@ -824,10 +830,10 @@ def suggest_categorical(self, name, choices):
             _run_reconstruction_pipeline,
         )
 
-        _run_reconstruction_pipeline(recon_obj, resolved_kwargs, class_type)
+        _run_reconstruction_pipeline(recon_obj, resolved_kwargs)
 
         # Extract loss
-        loss_getter = self._config.get("loss_getter")
+        loss_getter = config.get("loss_getter")
         if loss_getter is not None:
             loss = float(loss_getter(recon_obj))
         else:
@@ -856,14 +862,14 @@ def _plot_grid_objects(self, results, param_names, figsize):
 
         # Find best result
         losses = [r["loss"] for r in results]
-        best_idx = np.argmin(losses)
+        best_idx = (np.argmax if self.direction == "maximize" else np.argmin)(losses)
 
         for idx, result in enumerate(results):
             ax = axes[idx]
 
             recon_obj = result["reconstruction"]
 
-            obj = recon_obj._to_numpy(recon_obj.obj_cropped)
+            obj = recon_obj.obj_cropped
             if recon_obj.obj_type == "potential":
                 obj = np.abs(obj).sum(0)
             elif recon_obj.obj_type == "pure_phase":
@@ -872,19 +878,7 @@ def _plot_grid_objects(self, results, param_names, figsize):
             else:
                 obj = np.angle(obj).sum(0)
 
-            if obj is not None:
-                show_2d(obj, cmap="magma", figax=(fig, ax))
-            else:
-                ax.text(
-                    0.5,
-                    0.5,
-                    "No object\navailable",
-                    ha="center",
-                    va="center",
-                    transform=ax.transAxes,
-                    fontsize=10,
-                )
-                ax.set_facecolor("#f0f0f0")
+            show_2d(obj, cmap="magma", figax=(fig, ax))
 
             # Title with parameters and loss
             param_str = ", ".join(
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index 8a034ad65..f9c4e4e2a 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -490,9 +490,6 @@ def _compute_propagator_arrays(
         if probe_energy is None:
             raise ValueError("probe_model energy must be set to compute propagators.")
         wavelength = electron_wavelength_angstrom(probe_energy)
-        propagators = torch.empty(
-            (num_slices - 1, kr.shape[0], kc.shape[0]), dtype=torch.complex64, device=self.device
-        )
 
         theta_r, theta_c = self.probe_tilt
         dz = torch.tensor(slice_thicknesses, device=self.device, dtype=k2.dtype)  # (T,)
@@ -511,9 +508,6 @@ def _compute_propagator_arrays(
 class ProbeConstraints(BaseConstraints[PtychoProbeConstraintParams.Raster], ProbeBase):
     DEFAULT_CONSTRAINTS: PtychoProbeConstraintParams.Raster = PtychoProbeConstraintParams.Raster()
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
     def apply_soft_constraints(self, probe: torch.Tensor) -> torch.Tensor:
         self.reset_soft_constraint_losses()
         loss = self._get_zero_loss_tensor()
@@ -599,28 +593,6 @@ def _probe_orthogonalization_constraint(self, start_probe: torch.Tensor) -> torc
         return orthogonal_probes_sorted
 
 
-#    def _probe_orthogonalization_constraint(self, start_probe: torch.Tensor) -> torch.Tensor:
-#        """
-#        """
-#        n_probes = start_probe.shape[0]
-#
-#        # Gram matrix, G = P @ P.H
-#        P = start_probe.view(n_probes,-1)
-#        G = P @ P.conj().T
-#
-#        # eigen-decomposition of G
-#        _, eigenvecs = torch.linalg.eigh(G)
-#
-#        # rotate probes into orthogonal basis
-#        orthogonal_probes = torch.tensordot(eigenvecs.T, start_probe, dims=1)
-#
-#        # sort by intensity
-#        intensities = torch.sum(torch.abs(orthogonal_probes) ** 2, dim=(-2,-1))
-#        order = torch.argsort(intensities, descending=True)
-#
-#        return orthogonal_probes[order]
-
-
 class ProbePixelated(ProbeConstraints):
     def __init__(
         self,
@@ -835,11 +807,6 @@ def set_initial_probe(
     def reset(self):
         super().reset()
         self.probe = self._initial_probe.clone()
-        self._probe = nn.Parameter(self._initial_probe.clone().to(self.device), requires_grad=True)
-
-    def to(self, *args, **kwargs) -> Self:
-        super().to(*args, **kwargs)
-        return self
 
     @property
     def name(self) -> str:
@@ -872,7 +839,7 @@ def vacuum_probe_intensity(self, vp: np.ndarray | torch.Tensor | Dataset4dstem |
         elif isinstance(vp, np.ndarray):
             vp2 = vp.astype(config.get("dtype_real"))
         elif isinstance(vp, (Dataset4dstem, Dataset2d)):
-            vp2 = cast(np.ndarray, vp.array) # TODO when finished Dataset->torch fix here
+            vp2 = cast(np.ndarray, vp.array)  # TODO when finished Dataset->torch fix here
         elif isinstance(vp, torch.Tensor):
             vp2 = vp.cpu().detach().numpy()
         else:
@@ -1053,7 +1020,7 @@ def from_params(
         )
 
     @property
-    def vacuum_probe_intensity(self) -> np.ndarray | None:
+    def vacuum_probe_intensity(self) -> torch.Tensor | None:
         if self._vacuum_probe_intensity is None:
             return None
         return self._vacuum_probe_intensity
@@ -1066,7 +1033,7 @@ def vacuum_probe_intensity(self, vp: np.ndarray | Dataset4dstem | None):
         elif isinstance(vp, np.ndarray):
             vp2 = vp.astype(config.get("dtype_real"))
         elif isinstance(vp, (Dataset4dstem, Dataset2d)):
-            vp2 = cast(np.ndarray, vp.array) # TODO when finished Dataset->torch fix here
+            vp2 = cast(np.ndarray, vp.array)  # TODO when finished Dataset->torch fix here
         else:
             raise NotImplementedError(f"Unknown vacuum probe type: {type(vp)}")
 
@@ -1075,7 +1042,7 @@ def vacuum_probe_intensity(self, vp: np.ndarray | Dataset4dstem | None):
         elif vp2.ndim != 2:
             raise ValueError(f"Unexpected shape for vacuum probe: {vp2.shape}")
 
-        self._vacuum_probe_intensity = vp2
+        self._vacuum_probe_intensity = torch.tensor(vp2, dtype=torch.float32, device=self.device)
 
     @property
     def params(self) -> list[nn.Parameter]:
@@ -1265,22 +1232,6 @@ def model(self) -> "torch.nn.Module":
         """get the DIP model"""
         return self._model
 
-    @model.setter
-    def model(self, dip: "torch.nn.Module"):
-        """
-        This actually doesn't work -- can't have setters for torch sub modules
-        https://github.com/pytorch/pytorch/issues/52664
-        """
-        print("probe model setter hi")
-        if not isinstance(dip, torch.nn.Module):
-            raise TypeError(f"DIP must be a torch.nn.Module, got {type(dip)}")
-        if hasattr(dip, "dtype"):
-            dt = getattr(dip, "dtype")
-            if not dt.is_complex:
-                raise ValueError("DIP model must be a complex-valued model for probe objects")
-        self._model = dip.to(self.device)
-        self.set_pretrained_weights(self._model)
-
     @property
     def pretrained_weights(self) -> dict[str, torch.Tensor]:
         """get the pretrained weights of the DIP model"""
@@ -1360,8 +1311,8 @@ def probe(self) -> torch.Tensor:
     def _probe(self) -> torch.Tensor:
         return self.forward(None)  # type: ignore
 
-    def forward(self, fract_positions: torch.Tensor) -> torch.Tensor:
-        """Get shifted probes at fractional positions"""
+    def _noisy_model_input(self) -> torch.Tensor:
+        """model input with gaussian noise added when _input_noise_std > 0"""
         if self._input_noise_std > 0.0:
             noise = (
                 torch.randn(
@@ -1372,11 +1323,12 @@ def forward(self, fract_positions: torch.Tensor) -> torch.Tensor:
                 )
                 * self._input_noise_std
             )
-            model_input = self.model_input + noise
-        else:
-            model_input = self.model_input
+            return self.model_input + noise
+        return self.model_input
 
-        probe = self.model(model_input)[0]
+    def forward(self, fract_positions: torch.Tensor) -> torch.Tensor:
+        """Get shifted probes at fractional positions"""
+        probe = self.model(self._noisy_model_input())[0]
         shifted_probes = fourier_shift_expand(probe, fract_positions).swapaxes(0, 1)
         return shifted_probes
 
@@ -1392,6 +1344,7 @@ def set_initial_probe(
         super().set_initial_probe(
             roi_shape, reciprocal_sampling, mean_diffraction_intensity, device
         )
+        self._check_roi_shape()
 
         # could check if num_probes corresponds to out_channels of model
 
@@ -1463,7 +1416,7 @@ def pretrain(
                     f"Model target shape {pretrain_target.shape} does not match model input shape {self.model_input.shape}"
                 )
             self.pretrain_target = pretrain_target.clone().detach().to(self.device)
-        elif self.pretrain_target is None:
+        elif self._pretrain_target.numel() == 0:
             self.pretrain_target = self._initial_probe.clone().detach()
 
         loss_fn = get_loss_module(loss_fn, self.dtype)
@@ -1483,7 +1436,7 @@ def _pretrain(
         show: bool = False,
     ):
         """Pretrain the DIP model."""
-        if not hasattr(self, "pretrain_target"):
+        if self._pretrain_target.numel() == 0:
             raise ValueError("Pretrain target is not set. Use pretrain_target to set it.")
 
         self.model.train()
@@ -1496,19 +1449,7 @@ def _pretrain(
         output = self.probe
 
         for a0 in pbar:
-            if self._input_noise_std > 0.0:
-                noise = (
-                    torch.randn(
-                        self.model_input.shape,
-                        dtype=self.dtype,
-                        device=self.device,
-                        generator=self._rng_torch,
-                    )
-                    * self._input_noise_std
-                )
-                model_input = self.model_input + noise
-            else:
-                model_input = self.model_input
+            model_input = self._noisy_model_input()
 
             if apply_constraints:
                 output = self.apply_hard_constraints(self.model(model_input)[0])
@@ -1593,6 +1534,8 @@ def backward(self, propagated_gradient, obj_patches):
         )
 
     def _check_roi_shape(self):
+        if not hasattr(self, "_roi_shape"):
+            return
         num_layers = getattr(self.model, "num_layers", None)
         if num_layers is not None:
             if not np.all(np.array(self.roi_shape) % 2**num_layers == 0):
diff --git a/src/quantem/diffractive_imaging/ptycho_losses.py b/src/quantem/diffractive_imaging/ptycho_losses.py
index cd45970ae..9caa7cd26 100644
--- a/src/quantem/diffractive_imaging/ptycho_losses.py
+++ b/src/quantem/diffractive_imaging/ptycho_losses.py
@@ -122,6 +122,9 @@ class AmplitudeS3IM(DataCriterion):
     multi-GPU. It does, however, sit on a different absolute scale than the sum-based criteria
     (``L2``/``L1``/Poisson, which rescale to a full-dataset sum), so learning rates do **not**
     transfer between ``s3im_amplitude`` and those losses — retune the LR when switching.
+
+    Pass ``generator`` (a ``torch.Generator`` on the compute device) to make the random
+    permutations — and hence the loss values — reproducible run-to-run.
     """
 
     target_space: TargetSpace = "amplitude"
@@ -133,18 +136,20 @@ def __init__(
         patch_height: int = 32,
         window_size: int = 11,
         sigma: float = 1.5,
+        generator: torch.Generator | None = None,
     ):
         self.lambda_s3im = float(lambda_s3im)
         self.repeats = int(repeats)
         self.patch_height = int(patch_height)
         self.window_size = int(window_size)
         self.sigma = float(sigma)
+        self.generator = generator
 
     def _s3im(self, src: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
         num = src.numel()
         idx_list = [torch.arange(num, device=src.device)]
         for _ in range(self.repeats - 1):
-            idx_list.append(torch.randperm(num, device=src.device))
+            idx_list.append(torch.randperm(num, device=src.device, generator=self.generator))
         idx = torch.cat(idx_list)
         ph = self.patch_height
         usable = (idx.numel() // ph) * ph  # trim so it reshapes to (ph, -1)
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index ea277533b..dcf8833bd 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -1,4 +1,4 @@
-from math import ceil, floor
+from math import ceil
 from typing import Literal, Union, overload
 
 import numpy as np
@@ -37,42 +37,9 @@ def __init__(
             self.train_indices = np.asarray(train_indices, dtype=int)
             self.val_indices = np.asarray(val_indices, dtype=int)
         else:
-            # Validate ratio and split deterministically given rng
-            if val_ratio < 0 or val_ratio >= 1:
-                val_ratio = 0.0
-            n_val = int(round(len(self.indices) * val_ratio))
-            if n_val > 0:
-                if val_mode == "random":
-                    # Random unique selection for validation
-                    perm = self.rng.permutation(self.indices)
-                    self.val_indices = perm[:n_val]
-                    self.train_indices = np.setdiff1d(
-                        self.indices, self.val_indices, assume_unique=False
-                    )
-                else:  # grid/regular selection: every k-th index
-                    if val_ratio <= 0.5:
-                        k = max(1, int(round(1.0 / val_ratio)))
-                        invert = False
-                    else:
-                        k = max(1, int(round(1.0 / (1.0 - val_ratio))))
-                        invert = True
-
-                    grid_sel = self.indices[::k]
-                    if len(grid_sel) > n_val:
-                        grid_sel = grid_sel[:n_val]
-                    if invert:
-                        self.train_indices = grid_sel
-                        self.val_indices = np.setdiff1d(
-                            self.indices, grid_sel, assume_unique=False
-                        )
-                    else:
-                        self.val_indices = grid_sel
-                        self.train_indices = np.setdiff1d(
-                            self.indices, self.val_indices, assume_unique=False
-                        )
-            else:
-                self.val_indices = np.asarray([], dtype=int)
-                self.train_indices = self.indices
+            self.train_indices, self.val_indices = compute_train_val_split(
+                num, val_ratio, val_mode, self.rng
+            )
 
     @property
     def rng(self) -> np.random.Generator:
@@ -174,7 +141,8 @@ def fourier_shift_expand(
     array: ArrayLike, positions: ArrayLike, expand_dim: bool = True
 ) -> ArrayLike:
     """Fourier-shift array by flat array of positions."""
-    phase = fourier_translation_operator(positions, array.shape, expand_dim, dtype=array.dtype)
+    dtype = array.dtype if af.is_complex(array) else None
+    phase = fourier_translation_operator(positions, array.shape, expand_dim, dtype=dtype)
     fourier_array = af.fft2(array)
     shifted_fourier_array = fourier_array * phase
     shifted_array = af.ifft2(shifted_fourier_array)
@@ -221,42 +189,6 @@ def fourier_translation_operator(
     return ramp
 
 
-@overload
-def get_com_2d(ar: np.ndarray, corner_centered: bool = False) -> np.ndarray: ...
-@overload
-def get_com_2d(ar: "torch.Tensor", corner_centered: bool = False) -> "torch.Tensor": ...
-def get_com_2d(ar: ArrayLike, corner_centered: bool = False) -> ArrayLike:
-    """
-    Finds and returns the center of mass along last two dimensions.
-    If corner_centered is True, uses fftfreq for indices.
-    """
-    nr, nc = ar.shape[-2:]
-
-    if corner_centered:
-        c, r = np.meshgrid(np.fft.fftfreq(nc, 1 / nc), np.fft.fftfreq(nr, 1 / nr))
-    else:
-        c, r = np.meshgrid(np.arange(nc), np.arange(nr))
-
-    rc = af.match_device(np.stack([r, c]), ar)
-    com = (
-        af.sum(
-            rc * ar[..., None, :, :],
-            axis=(
-                -1,
-                -2,
-            ),
-        )
-        / af.sum(
-            ar,
-            axis=(
-                -1,
-                -2,
-            ),
-        )[:, None]
-    )
-    return com
-
-
 def sum_patches_base(
     patches: torch.Tensor, indices: torch.Tensor, obj_shape: tuple
 ) -> torch.Tensor:
@@ -541,7 +473,7 @@ def from_array(cls, T: np.ndarray):
             return cls()
         R[1] /= scale1
         shear1 /= scale1
-        angle = np.arccos(R[0, 0])
+        angle = np.arctan2(-R[0, 1], R[0, 0])
 
         if T.shape[0] > 2:
             t0, t1 = T[2]
@@ -663,9 +595,9 @@ def center_crop_arr(
                 raise ValueError(
                     f"Dimension {i} of shape ({s}) is larger than dimension {i} of arr ({a})."
                 )
-            pad[i] = [int(floor(s - a) / 2), int(ceil(s - a) / 2)]
+            pad[i] = [(s - a) // 2, -(-(s - a) // 2)]
 
-    if any(pad):
+    if any(p != [0, 0] for p in pad):
         arr = np.pad(arr, pad_width=pad, mode="constant")
 
     slices = []
diff --git a/src/quantem/diffractive_imaging/ptychography.py b/src/quantem/diffractive_imaging/ptychography.py
index 3f4206f7c..09276582a 100644
--- a/src/quantem/diffractive_imaging/ptychography.py
+++ b/src/quantem/diffractive_imaging/ptychography.py
@@ -212,8 +212,9 @@ def reset_recon(self) -> None:
         self.probe_model.reset_optimizer()
         self.dset.reset_optimizer()
 
-    def _record_iter(self, iter_loss: float) -> None:
+    def _record_iter(self, iter_loss: float, autograd: bool) -> None:
         self._iter_losses.append(iter_loss)
+        self._iter_recon_types.append("AD" if autograd else "GD")
         optimizers = self.optimizers
         all_keys = set(self._iter_lrs.keys()) | set(optimizers.keys())
         for key in all_keys:
@@ -318,10 +319,11 @@ def reconstruct(
             device if isinstance(device, list) else getattr(self, "_multi_gpu_devices", None)
         )
 
+        if (isinstance(devices_to_use, list) or is_distributed_launch()) and not autograd:
+            raise ValueError("Multi-GPU reconstruction requires autograd=True.")
+
         # Route to multi-GPU path
         if isinstance(devices_to_use, list) and not is_distributed_launch():
-            if not autograd:
-                raise ValueError("Multi-GPU reconstruction requires autograd=True.")
             return self._spawn_reconstruct(
                 devices=devices_to_use,
                 num_iters=num_iters,
@@ -420,6 +422,12 @@ def _reconstruct_inner(
             self.set_schedulers(self.scheduler_params, num_iter=num_iters)
 
         self.criterion = loss_type  # resolve name/instance -> DataCriterion
+        if not autograd and self._criterion.target_space != "amplitude":
+            raise ValueError(
+                "autograd=False uses the amplitude-projection update, which requires an "
+                f"amplitude-space loss; got loss_type with target_space="
+                f"{self._criterion.target_space!r}."
+            )
         self.dset._set_targets(self._criterion.target_space)
         self.compute_propagator_arrays()  # required to avoid issue if stopped learning probe tilt
 
@@ -463,7 +471,7 @@ def _reconstruct_inner(
                 )
                 pred_intensities = self.detector_model.forward(overlap)
 
-                batch_consistency_loss, targets = self.error_estimate(
+                batch_consistency_loss = self.error_estimate(
                     pred_intensities,
                     targets=targets,
                     global_n=global_n,
@@ -519,7 +527,7 @@ def _reconstruct_inner(
                             obj_patches, shifted_probes, descan_shifts
                         )
                         pred_intensities = self.detector_model.forward(overlap)
-                        batch_val_loss, _ = self.error_estimate(
+                        batch_val_loss = self.error_estimate(
                             pred_intensities,
                             targets=targets,
                             global_n=global_n,
@@ -539,7 +547,7 @@ def _reconstruct_inner(
                         self._iter_val_losses.append(val_loss)
 
             if _dist_rank == 0:
-                self._record_iter(total_loss)  # TODO record val loss as well
+                self._record_iter(total_loss, autograd)  # TODO record val loss as well
 
             # Step schedulers with current loss
             self.step_schedulers(total_loss)
@@ -666,6 +674,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
         # When reset=False the worker inherited existing history, so its lists are [old...new...].
         # n_before lets us take only the genuinely new tail in both cases.
         n_before = len(self._iter_losses)
+        v_before = len(self._iter_val_losses)
         is_reset = recon_kwargs.get("reset", False)
 
         if is_reset:
@@ -680,7 +689,7 @@ def _spawn_reconstruct(self, devices: list[int], **recon_kwargs) -> Self:
             self._iter_recon_types.extend(result.get("iter_recon_types", []))
         else:
             self._iter_losses.extend(result["iter_losses"][n_before:])
-            self._iter_val_losses.extend(result["iter_val_losses"][n_before:])
+            self._iter_val_losses.extend(result["iter_val_losses"][v_before:])
             for k, v in result.get("iter_lrs", {}).items():
                 if k not in self._iter_lrs:
                     self._iter_lrs[k] = []
@@ -701,7 +710,6 @@ def _get_current_lrs(self) -> dict[str, float]:
         return {
             param_name: optimizer.param_groups[0]["lr"]
             for param_name, optimizer in self.optimizers.items()
-            if optimizer is not None
         }
 
     def backward(
@@ -895,8 +903,6 @@ def from_file(
             Device to load the object on
         verbose : int | bool | None
             Verbosity level
-        rng : np.random.Generator | int | None
-            Random number generator
         auto_reload_dataset : bool
             Whether to automatically reload and preprocess the dataset from saved metadata
 
@@ -938,7 +944,7 @@ def from_file(
                             )
 
                     dset = PtychographyDatasetRaster.from_dataset4dstem(
-                        raw_dset, verbose=verbose or 1
+                        raw_dset, verbose=1 if verbose is None else verbose
                     )
                     # Apply preprocessing with saved parameters
                     preprocessing_params = metadata.get("preprocessing_params", {})
@@ -978,6 +984,9 @@ def from_file(
             #     "Please provide a dataset parameter or ensure the object was saved with dataset metadata."
             # )
 
+        if verbose is not None:
+            ptycho.verbose = verbose
+
         if device is not None:
             ptycho.to(device)
 
diff --git a/src/quantem/diffractive_imaging/ptychography_base.py b/src/quantem/diffractive_imaging/ptychography_base.py
index 22f7ae1d4..864f3cf63 100644
--- a/src/quantem/diffractive_imaging/ptychography_base.py
+++ b/src/quantem/diffractive_imaging/ptychography_base.py
@@ -125,7 +125,7 @@ def __init__(  # TODO prevent direct instantiation
         self._iter_lrs: dict[str, list[float]] = {}  # LRs/step_sizes across iterations
         self._snapshots: list[Snapshot] = []
         self._obj_padding_px = np.array([0, 0])
-        self.obj_fov_mask = torch.ones(self.dset._obj_shape_full_2d(self.obj_padding_px).shape)
+        self.obj_fov_mask = torch.ones(tuple(self.dset._obj_shape_full_2d(self.obj_padding_px)))
         self.batch_size = self.dset.num_gpts
         self._val_ratio = 0.0
         self._val_mode: Literal["grid", "random"] = "grid"
@@ -150,7 +150,7 @@ def __init__(  # TODO prevent direct instantiation
         self.to(self._single_device)
 
     # region --- preprocessing ---
-    ## hopefully will be able to remove some of thes preprocessing flags,
+    ## hopefully will be able to remove some of these preprocessing flags,
     ## convert plotting and vectorized to kwargs
     def preprocess(
         self,
@@ -195,8 +195,6 @@ def preprocess(
 
         # change obj_padding_px and whatever else needs to be changed
         self.obj_padding_px = obj_padding_px  # also initializes the object model
-        self.dset._set_initial_scan_positions_px(self.obj_padding_px)
-        self.dset._set_patch_indices(self.obj_padding_px)
 
         self.compute_propagator_arrays()
         self._set_obj_fov_mask(batch_size=batch_size)
@@ -212,7 +210,9 @@ def _set_obj_fov_mask(self, gaussian_sigma: float = 2.0, batch_size=None):
         overlap = self._get_probe_overlap(batch_size)
         ov = overlap > overlap.max() * 0.3
         ov = ndi.binary_closing(ov, iterations=5)
-        ov = ndi.binary_dilation(ov, iterations=min(32, np.min(self.obj_padding_px) // 4))
+        dilation_iters = min(32, np.min(self.obj_padding_px) // 4)
+        if dilation_iters > 0:
+            ov = ndi.binary_dilation(ov, iterations=dilation_iters)
         ov = ndi.gaussian_filter(ov.astype(config.get("dtype_real")), sigma=gaussian_sigma)
         self.obj_fov_mask = ov
         self.obj_model.mask = ov
@@ -457,7 +457,7 @@ def iter_lrs(self) -> dict[str, np.ndarray]:
 
     @property
     def probe(self) -> np.ndarray:
-        """Complex valued probe(s). Shape [num_probes, roi_reight, roi_width]"""
+        """Complex valued probe(s). Shape [num_probes, roi_height, roi_width]"""
         return self._to_numpy(self.probe_model.probe)
 
     @property
@@ -804,11 +804,6 @@ def _check_preprocessed(self):
                 "Preprocessing has not been completed. Please run Ptycho.preprocess()"
             )
 
-    def _check_rm_preprocessed(self, new_val: Any, name: str) -> None:
-        if hasattr(self, name):
-            if getattr(self, name) != new_val:
-                self._preprocessed = False
-
     def _to_numpy(self, array: "np.ndarray | torch.Tensor") -> np.ndarray:
         return to_numpy(array)
 
@@ -878,15 +873,6 @@ def _crop_rotate_obj_fov(
 
         return cropped
 
-    def _repeat_arr(
-        self, arr: "np.ndarray|torch.Tensor", repeats: int, axis: int
-    ) -> "np.ndarray|torch.Tensor":
-        """repeat the input array along the desired axis."""
-        if config.get("has_torch"):
-            if isinstance(arr, torch.Tensor):
-                return torch.repeat_interleave(arr, repeats, dim=axis)
-        return np.repeat(arr, repeats, axis=axis)
-
     def reset_recon(self) -> None:
         self._reset_rng()
         self.obj_model.reset()
@@ -1074,7 +1060,7 @@ def _build_dataloaders(
 
     # endregion
 
-    # region --- ptychography foRcard model ---
+    # region --- ptychography forward model ---
 
     def forward_operator(
         self,
@@ -1097,7 +1083,7 @@ def error_estimate(
         pred_intensities: torch.Tensor,
         targets: torch.Tensor,
         global_n: int | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
         """Data-fidelity loss for one batch via the active criterion (``self.criterion``).
 
         Maps predictions into the criterion's measurement space (amplitude or intensity),
@@ -1115,7 +1101,7 @@ def error_estimate(
         n = global_n if global_n is not None else self.dset.num_positions
         error = criterion(preds * mask, targets * mask, n)
         loss = error / self.dset.mean_diffraction_intensity
-        return loss, targets
+        return loss
 
     def overlap_projection(self, obj_patches, input_probe):
         """Multiplies `input_probes` with roi-shaped patches from `obj_array`.
@@ -1140,14 +1126,14 @@ def estimate_amplitudes(
         # incoherent sum of all probe components
         eps = 1e-9  # this is to avoid diverging gradients at sqrt(0)
         overlap_fft = torch.fft.fft2(overlap_array, norm="ortho")
-        amps = torch.sqrt(torch.sum(torch.abs(overlap_fft + eps) ** 2, dim=0))
+        amps = torch.sqrt(torch.sum(torch.abs(overlap_fft) ** 2, dim=0) + eps)
         if not corner_centered:  # default is shifted amplitudes matching exp data
             return torch.fft.fftshift(amps, dim=(-2, -1))
         else:
             return amps
 
     def estimate_intensities(self, overlap_array: "torch.Tensor") -> "torch.Tensor":
-        """Returns the estimated fourier amplitudes from real-valued `overlap_array`."""
+        """Returns the estimated fourier intensities from real-valued `overlap_array`."""
         # overlap shape: (nprobes, batch_size, roi_shape[0], roi_shape[1])
         # incoherent sum of all probe components
         overlap_fft = torch.fft.fft2(overlap_array, norm="ortho")
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index eea9fc71f..1eedacb05 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -4,7 +4,6 @@
 
 import numpy as np
 import torch
-import torch.nn as nn
 
 from quantem.core import config
 from quantem.core.datastructures import Dataset4dstem
@@ -71,13 +70,13 @@ def from_dataset(
             Object parameterization.
         num_probes : int
             Number of probe components (mixed state when >1).
-        energy, defocus, semiangle_cutoff, rolloff, polar_parameters
+        energy, defocus, semiangle_cutoff, polar_parameters
             Probe settings passed to ProbePixelated.
         vacuum_probe_intensity : np.ndarray | Dataset4dstem | None
             Optional corner-centered vacuum probe intensity for scaling/centering.
         initial_probe_weights : list[float] | np.ndarray | None
             Optional initial component weights (length=num_probes).
-        log_dir, log_prefix, log_suffix, log_images_every, log_probe_images, device, verbose, rng
+        log_dir, log_prefix, log_images_every, log_probe_images, device, verbose, rng
             Standard Ptychography configuration.
         """
 
@@ -205,6 +204,11 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         opt_params: dict[str, Any] | None
         scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
+            scheduler_dict: dict[str, Any] = {
+                "name": "exponential" if scheduler_type == "exp" else scheduler_type
+            }
+            if scheduler_type in ("exp", "plateau"):
+                scheduler_dict["factor"] = scheduler_factor
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -212,29 +216,20 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
                 },
             }
             scheduler_params = {
-                "object": {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                "object": dict(scheduler_dict),
             }
             if learn_probe:
                 opt_params["probe"] = {
                     "name": "adamw",
                     "lr": lr_probe,
                 }
-                scheduler_params["probe"] = {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                scheduler_params["probe"] = dict(scheduler_dict)
             if needs_dataset_optimizer:
                 opt_params["dataset"] = {
                     "name": "adamw",
                     "lr": lr_scan_positions,
                 }
-                scheduler_params["dataset"] = {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                scheduler_params["dataset"] = dict(scheduler_dict)
         else:
             opt_params = None
             scheduler_params = None
@@ -290,8 +285,8 @@ class PtychoLiteDIP(Ptychography):
     """
     High-level convenience wrapper around Ptychography.
 
-    Provides a from_dataset() constructor that builds pixelated object and probe
-    models from simple flags, then initializes a full Ptychography instance.
+    Provides a from_ptycholite() constructor that builds DIP object and probe
+    models from an existing PtychoLite, then initializes a full Ptychography instance.
     """
 
     @classmethod
@@ -307,7 +302,7 @@ def from_ptycholite(
         normalize_object_plotting: bool = True,
         # model settings
         cnn_num_layers: int = 3,
-        final_activation: "str | Callable[..., Any]" = nn.Identity(),
+        final_activation: "str | Callable[..., Any]" = "identity",
         # logging/device
         log_dir: os.PathLike[str] | str | None = None,
         log_prefix: str = "",
@@ -387,7 +382,7 @@ def from_ptycholite(
             logger = LoggerPtychography(
                 log_dir=log_dir,
                 run_prefix=log_prefix,
-                run_suffix="pix",
+                run_suffix="dip",
                 log_images_every=log_images_every,
                 log_probe_images=log_probe_images,
             )
@@ -404,7 +399,7 @@ def from_ptycholite(
             detector_model=ptycholite.detector_model,
             logger=logger if logger is not None else ptycholite.logger,
             device=device,
-            verbose=ptycholite.verbose,
+            verbose=verbose,
             rng=ptycholite.rng,
         )
 
@@ -452,6 +447,11 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         opt_params: dict[str, Any] | None
         scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
+            scheduler_dict: dict[str, Any] = {
+                "name": "exponential" if scheduler_type == "exp" else scheduler_type
+            }
+            if scheduler_type in ("exp", "plateau"):
+                scheduler_dict["factor"] = scheduler_factor
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -459,29 +459,20 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
                 },
             }
             scheduler_params = {
-                "object": {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                "object": dict(scheduler_dict),
             }
             if learn_probe:
                 opt_params["probe"] = {
                     "name": "adamw",
                     "lr": lr_probe,
                 }
-                scheduler_params["probe"] = {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                scheduler_params["probe"] = dict(scheduler_dict)
             if needs_dataset_optimizer:
                 opt_params["dataset"] = {
                     "name": "adamw",
                     "lr": lr_scan_positions,
                 }
-                scheduler_params["dataset"] = {
-                    "name": scheduler_type,
-                    "factor": scheduler_factor,
-                }
+                scheduler_params["dataset"] = dict(scheduler_dict)
         else:
             opt_params = None
             scheduler_params = None
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index dcb10ea3f..679295bc4 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -42,8 +42,6 @@ def show_obj(
         obj_iter = "Final"
         if obj is None:
             if snapshot_iter is not None:
-                if snapshot_iter < 0:
-                    snapshot_iter = len(self.snapshots) + snapshot_iter
                 snp = self.get_snapshot_by_iter(snapshot_iter, closest=True, cropped=True)
                 obj_np = snp["obj"]
                 obj_iter = snp["iteration"]
@@ -88,7 +86,7 @@ def show_obj(
             titles.extend([t + "Phase", t + "Amplitude"])
             cmaps.extend([ph_cmap, "gray"])
 
-        scalebar = [{"sampling": self.sampling[0], "units": "Å"}] + [None] * (len(ims) - 1)
+        scalebar = [{"sampling": self.sampling[1], "units": "Å"}] + [None] * (len(ims) - 1)
 
         show_2d(
             ims,
@@ -159,16 +157,12 @@ def show_obj_fft(
         elif self.obj_type == "pure_phase":
             windowed_obj = np.exp(1j * obj_np.sum(0)) * window_2d
         else:
-            windowed_obj = (
-                np.abs(obj_np).sum(0)
-                * window_2d
-                * np.exp(1j * np.angle(obj_np).sum(0) * window_2d)
-            )
+            windowed_obj = np.abs(obj_np).sum(0) * np.exp(1j * np.angle(obj_np).sum(0)) * window_2d
         obj_pad = np.pad(windowed_obj, pad, mode="constant", constant_values=0)
 
         obj_fft = np.fft.fftshift(np.fft.fft2(obj_pad))
 
-        fft_sampling = 1 / (self.sampling[0] * obj_pad.shape[0])
+        fft_sampling = 1 / (self.sampling[1] * obj_pad.shape[1])
         fft_scalebar = {"sampling": fft_sampling, "units": r"$\mathrm{A^{-1}}$"}
 
         t = kwargs.pop("title", "")
@@ -178,7 +172,7 @@ def show_obj_fft(
             t += f"Iter {obj_iter} "
 
         if show_obj:
-            obj_scalebar = {"sampling": self.sampling[0], "units": "Å"}
+            obj_scalebar = {"sampling": self.sampling[1], "units": "Å"}
             if self.obj_type == "potential":
                 obj_show = obj_pad
             else:  # complex or pure phase just show the phase
@@ -190,19 +184,17 @@ def show_obj_fft(
                 ],
                 title=[t + "Object", t + "Fourier Transform"],
                 scalebar=[obj_scalebar, fft_scalebar],
-                return_fig=True,
                 **kwargs,
             )
-            ax[1].set_aspect(obj_np.shape[-1] / obj_np.shape[-2])
+            ax[1].set_aspect(obj_pad.shape[-1] / obj_pad.shape[-2])
         else:
             fig, ax = show_2d(
                 np.abs(obj_fft),
                 scalebar=fft_scalebar,
                 title=t + "Fourier Transform",
-                return_fig=True,
                 **kwargs,
             )
-            ax.set_aspect(obj_np.shape[-1] / obj_np.shape[-2])
+            ax.set_aspect(obj_pad.shape[-1] / obj_pad.shape[-2])
         if return_fft:
             return obj_fft
         else:
@@ -236,8 +228,6 @@ def show_probe(
         probe_iter = "Final"
         if probe is None:
             if snapshot_iter is not None:
-                if snapshot_iter < 0:
-                    snapshot_iter = len(self.snapshots) + snapshot_iter
                 snp = self.get_snapshot_by_iter(snapshot_iter, closest=True, cropped=True)
                 probe = snp["probe"]
                 probe_iter = snp["iteration"]
@@ -256,7 +246,7 @@ def show_probe(
         if probe_iter != "Final":
             t += f"Iter {probe_iter} "
 
-        scalebar = [{"sampling": self.sampling[0], "units": "Å"}]
+        scalebar = [{"sampling": self.sampling[1], "units": "Å"}]
         if sum_probes:
             probes = [np.fft.fftshift(probe.sum(0))]
         else:
@@ -318,8 +308,6 @@ def show_probe_top_bottom(
         """
         if probe is None:
             if snapshot_iter is not None:
-                if snapshot_iter < 0:
-                    snapshot_iter = len(self.snapshots) + snapshot_iter
                 snp = self.get_snapshot_by_iter(snapshot_iter, closest=True, cropped=True)
                 probe = snp["probe"]
             else:
@@ -373,7 +361,7 @@ def show_probe_top_bottom(
             top_img = np.fft.fftshift(top_img)
             bottom_img = np.fft.fftshift(bottom_img)
 
-        scalebar = [{"sampling": self.sampling[0], "units": "Å"}, None]
+        scalebar = [{"sampling": self.sampling[1], "units": "Å"}, None]
         titles = [f"Top Surface {label}", f"Bottom Surface {label}"]
 
         fig, axs = show_2d(
@@ -410,7 +398,7 @@ def show_fourier_probe(self, probe: np.ndarray | None = None, **kwargs):
                 probe = probe[None, ...]
 
         probes = [np.fft.fftshift(np.fft.fft2(probe[i])) for i in range(len(probe))]
-        scalebar = [{"sampling": self.reciprocal_sampling[0], "units": r"$\mathrm{A^{-1}}$"}] + [
+        scalebar = [{"sampling": self.reciprocal_sampling[1], "units": r"$\mathrm{A^{-1}}$"}] + [
             None
         ] * (len(probes) - 1)
         if len(probes) > 1:
@@ -477,7 +465,7 @@ def show_obj_slices(
         self,
         obj: np.ndarray | None = None,
         cbar: bool = False,
-        interval_type: Literal["quantile", "manual"] = "quantile",
+        interval_type: Literal["quantile", "manual", "minmax", "abs"] = "quantile",
         interval_scaling: Literal["each", "all"] = "each",
         max_width: int = 4,
         return_fig: bool = False,
@@ -491,7 +479,7 @@ def show_obj_slices(
             The object to show. If None, the object from the last iteration is shown.
         cbar: bool, optional
             Whether to show a colorbar, by default False
-        interval_type: Literal["quantile", "manual"], optional
+        interval_type: Literal["quantile", "manual", "minmax", "abs"], optional
             The interval type to use for the colorbar, by default "quantile"
         interval_scaling: Literal["each", "all"], optional
             The interval scaling to use for the colorbar, by default "each"
@@ -514,9 +502,14 @@ def show_obj_slices(
             if obj.ndim == 2:
                 obj = obj[None, ...]
 
+        # slice_thicknesses has length num_slices - 1; entry i is the gap to the next slice
+        thicknesses = self.slice_thicknesses
         t_parts = []
         for i in range(len(obj)):
-            t_parts.append(f"{i + 1}/{len(obj)} | {self.slice_thicknesses[i - 1]:.1f} Å")
+            t_part = f"{i + 1}/{len(obj)}"
+            if i < len(thicknesses):
+                t_part += f" | {thicknesses[i]:.1f} Å"
+            t_parts.append(t_part)
 
         if self.obj_type == "potential":
             objs_flat = [np.abs(obj[i]) for i in range(len(obj))]
@@ -533,7 +526,7 @@ def show_obj_slices(
         titles = [titles_flat[i : i + max_width] for i in range(0, len(titles_flat), max_width)]
 
         scalebars: list = [[None for _ in row] for row in objs]
-        scalebars[0][0] = {"sampling": self.sampling[0], "units": "Å"}
+        scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
 
         if interval_type == "quantile":
             norm = {"interval_type": "quantile"}
@@ -548,8 +541,8 @@ def show_obj_slices(
                 norm["vmin"] = np.min(objs_flat)
                 norm["vmax"] = np.max(objs_flat)
             else:
-                norm["vmin"] = kwargs.get("vmin")
-                norm["vmax"] = kwargs.get("vmax")
+                norm["vmin"] = kwargs.pop("vmin", None)
+                norm["vmax"] = kwargs.pop("vmax", None)
         else:
             raise ValueError(f"Unknown interval type: {interval_type}")
 
@@ -594,7 +587,13 @@ def plot_losses(self, figax: tuple | None = None, plot_lrs: bool = True):
         if len(self.val_iter_losses) > 0:
             lines.extend(ax.semilogy(iters, self.iter_losses, c="k", label="train loss", lw=lw))
             lines.extend(
-                ax.semilogy(iters, self.val_iter_losses, c=colors[6], label="val loss", lw=lw)
+                ax.semilogy(
+                    np.arange(len(self.val_iter_losses)),
+                    self.val_iter_losses,
+                    c=colors[6],
+                    label="val loss",
+                    lw=lw,
+                )
             )
         else:
             lines.extend(ax.semilogy(iters, self.iter_losses, c="k", label="loss", lw=lw))
@@ -605,7 +604,7 @@ def plot_losses(self, figax: tuple | None = None, plot_lrs: bool = True):
         ax.set_xlabel("Iterations")
 
         # check if all lrs are constant and if so, don't plot lr
-        if all(np.all(lr == self.iter_lrs["object"][0]) for lr in self.iter_lrs.values()):
+        if all(np.all(lr == lr[0]) for lr in self.iter_lrs.values() if len(lr) > 0):
             plot_lrs = False
 
         if plot_lrs and len(self.iter_lrs) > 0:
@@ -655,7 +654,8 @@ def plot_losses(self, figax: tuple | None = None, plot_lrs: bool = True):
             # set title to each lr type
             title = ""
             for lr_type, lr_values in self.iter_lrs.items():
-                title += f"{lr_type} LR: {lr_values[0]:.1e} | "
+                if len(lr_values) > 0:
+                    title += f"{lr_type} LR: {lr_values[0]:.1e} | "
             ax.set_title(title[:-3], fontsize=10)
 
         labs = [lin.get_label() for lin in lines]
@@ -726,22 +726,28 @@ def show_iters(
         show_object : bool, optional
             Whether to show object reconstructions, by default True
         iters : list[int] | slice | None, optional
-            Specific iter iterations to display. If None, shows all available iters
+            Specific iter iterations to display. If None, shows all available iters.
+            Takes precedence over every_nth.
         every_nth : int | None, optional
-            Show every nth iter instead of all. Overrides iters parameter
+            Show every nth iter instead of all. Only used if iters is None
         max_n : int | None, optional
             Maximum number of iterations to display
         cbar : bool, optional
             Whether to show colorbars, by default False
         norm : str, optional
-            Normalization method for object display, by default "quantile",
+            Normalization method for object display, by default "quantile". Ignored if
+            show_probe and show_object are both True.
         interval_scaling : str, optional
             How to scale intervals: "each" for per-image scaling, "all" for global scaling across
-            all iterations, by default "each". Does not work if show_probe is True.
+            all iterations, by default "each". Ignored if show_probe and show_object are both
+            True.
         max_width : int, optional
             Maximum number of images per row, by default 4
         cropped : bool, optional
             Whether to show cropped objects (default True) or full objects
+        closest : bool, optional
+            Whether to fall back to the closest stored snapshot when an exact iteration is
+            not available, by default True
         **kwargs
             Additional arguments passed to show_2d
         """
@@ -780,6 +786,12 @@ def show_iters(
         ]
 
         if show_object and show_probe:
+            if norm != "quantile" or interval_scaling != "each":
+                warnings.warn(
+                    "norm and interval_scaling are ignored when both the object and probe "
+                    "are shown; pass show_probe=False to use them.",
+                    stacklevel=2,
+                )
             self._show_object_and_probe_iters(selected_snapshots, cbar, max_width, **kwargs)
         elif show_object:
             self._show_object_iters_only(
@@ -834,8 +846,8 @@ def _show_object_iters_only(
                     norm_dict["vmin"] = float(np.min(all_values_flat))
                     norm_dict["vmax"] = float(np.max(all_values_flat))
             else:
-                norm_dict["vmin"] = kwargs.get("vmin")
-                norm_dict["vmax"] = kwargs.get("vmax")
+                norm_dict["vmin"] = kwargs.pop("vmin", None)
+                norm_dict["vmax"] = kwargs.pop("vmax", None)
         else:
             raise ValueError(f"Unknown norm type: {norm}")
 
@@ -845,7 +857,7 @@ def _show_object_iters_only(
 
         scalebars: list = [[None for _ in row] for row in images_grid]
         if scalebars:
-            scalebars[0][0] = {"sampling": self.sampling[0], "units": "Å"}
+            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
 
         show_2d(
             images_grid,
@@ -884,10 +896,7 @@ def _show_probe_iters_only(
         # Set up scalebars
         scalebars: list = [[None for _ in row] for row in probes_grid]
         if scalebars:
-            scalebars[0][0] = {
-                "sampling": self.reciprocal_sampling[0],
-                "units": r"$\mathrm{A^{-1}}$",
-            }
+            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
 
         show_2d(
             probes_grid,
@@ -960,11 +969,8 @@ def _show_object_and_probe_iters(
 
         scalebars: list = [[None for _ in row] for row in all_images]
         if scalebars:
-            scalebars[0][0] = {"sampling": self.sampling[0], "units": "Å"}
-            scalebars[0][1] = {
-                "sampling": self.reciprocal_sampling[0],
-                "units": r"$\mathrm{A^{-1}}$",
-            }
+            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
+            scalebars[0][1] = {"sampling": self.sampling[1], "units": "Å"}
 
         show_2d(
             all_images,
@@ -1285,7 +1291,7 @@ def show_fourier_probe_and_amplitudes(
         else:
             amplitudes = self._to_numpy(amplitudes.sum(0))
 
-        scalebar = [{"sampling": self.reciprocal_sampling[0], "units": r"$\mathrm{A^{-1}}$"}]
+        scalebar = [{"sampling": self.reciprocal_sampling[1], "units": r"$\mathrm{A^{-1}}$"}]
 
         if fft_shift:
             probe_plot = np.fft.fftshift(probe_plot)

From 50b6ce322fbe5f9cc86b99274f9914a873f3148e Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Sun, 2 Aug 2026 12:23:32 -0700
Subject: [PATCH 58/59] some code simplifications and dedup

---
 .../diffractive_imaging/object_models.py      |  46 ++-----
 .../optimize_hyperparameters.py               | 124 ++++++------------
 .../diffractive_imaging/probe_models.py       |  21 +--
 .../diffractive_imaging/ptycho_utils.py       |  25 +++-
 .../diffractive_imaging/ptychography_lite.py  |  21 +--
 .../ptychography_visualizations.py            |  30 +++--
 6 files changed, 108 insertions(+), 159 deletions(-)

diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index 23783f1a3..0963dd9f9 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -32,7 +32,7 @@
 )
 from quantem.core.visualization import show_2d
 from quantem.core.visualization.custom_normalizations import CustomNormalization
-from quantem.diffractive_imaging.ptycho_utils import sum_patches
+from quantem.diffractive_imaging.ptycho_utils import add_input_noise, sum_patches
 
 object_type = Literal["potential", "pure_phase", "complex"]
 
@@ -919,8 +919,8 @@ def from_array(
                 initial = initial.angle()
             else:
                 raise ValueError(f"Complex initial_obj is not valid for obj_type '{obj_type}'")
-        obj_model._initial_obj = (
-            initial.clone().detach().to(dtype=obj_model.dtype, device=obj_model.device)
+        obj_model._initial_obj = initial.detach().to(
+            dtype=obj_model.dtype, device=obj_model.device, copy=True
         )
 
         return obj_model
@@ -1177,11 +1177,7 @@ def name(self) -> str:
     def dtype(self) -> "torch.dtype":
         if hasattr(self.model, "dtype"):
             return getattr(self.model, "dtype")
-        else:
-            if self.obj_type in ["complex"]:
-                return getattr(torch, config.get("dtype_complex"))
-            else:
-                return getattr(torch, config.get("dtype_real"))
+        return super().dtype
 
     @property
     def model(self) -> "torch.nn.Module":
@@ -1298,20 +1294,9 @@ def _obj(self):
 
     def forward(self, patch_indices: torch.Tensor):
         """Get object patches at given indices"""
-        if self._input_noise_std > 0.0:
-            noise = (
-                torch.randn(
-                    self.model_input.shape,
-                    dtype=self.dtype,
-                    device=self.device,
-                    generator=self._rng_torch,
-                )
-                * self._input_noise_std
-            )
-            model_input = self.model_input + noise
-        else:
-            model_input = self.model_input
-
+        model_input = add_input_noise(
+            self.model_input, self._input_noise_std, self.dtype, self.device, self._rng_torch
+        )
         obj_array = self.model(model_input)[0]
         if self.mask.numel() > 0:
             obj_array = obj_array * self._mask
@@ -1428,20 +1413,9 @@ def _pretrain(
         output = self.obj
 
         for a0 in pbar:
-            if self._input_noise_std > 0.0:
-                noise = (
-                    torch.randn(
-                        self.model_input.shape,
-                        dtype=self.dtype,
-                        device=self.device,
-                        generator=self._rng_torch,
-                    )
-                    * self._input_noise_std
-                )
-                model_input = self.model_input + noise
-            else:
-                model_input = self.model_input
-
+            model_input = add_input_noise(
+                self.model_input, self._input_noise_std, self.dtype, self.device, self._rng_torch
+            )
             if apply_constraints:
                 output = self.apply_hard_constraints(self.model(model_input)[0])
             else:
diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
index 99904e203..ea73bf02a 100644
--- a/src/quantem/diffractive_imaging/optimize_hyperparameters.py
+++ b/src/quantem/diffractive_imaging/optimize_hyperparameters.py
@@ -2,7 +2,6 @@
 
 import copy
 import gc
-import inspect
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Mapping, Optional
 
@@ -17,6 +16,7 @@
     PtychographyDatasetBase,
     PtychographyDatasetRaster,
 )
+from quantem.diffractive_imaging.ptychography_lite import PtychoLite, PtychoLiteDIP
 
 
 @dataclass
@@ -257,8 +257,8 @@ def _run_reconstruction_pipeline(recon_obj, resolved_kwargs):
     reconstruct_kwargs = resolved_kwargs.get("reconstruct")
     if reconstruct_kwargs:
         reconstruct_kwargs = dict(reconstruct_kwargs)
-        # only PtychoLite.reconstruct takes verbose, and it resets recon_obj.verbose
-        if "verbose" in inspect.signature(recon_obj.reconstruct).parameters:
+        # only PtychoLite/PtychoLiteDIP.reconstruct take verbose, and they reset recon_obj.verbose
+        if isinstance(recon_obj, (PtychoLite, PtychoLiteDIP)):
             reconstruct_kwargs.setdefault("verbose", False)
         recon_obj.reconstruct(**reconstruct_kwargs)
 
@@ -549,45 +549,7 @@ def visualize(self, figsize=None):
 
             # Second and third subplots: individual parameter plots
             for idx, param_name in enumerate(param_names):
-                ax = axes[idx + 1]
-
-                # Extract data
-                param_trials = [t for t in trials if param_name in t.params]
-                param_values = np.array([trial.params[param_name] for trial in param_trials])
-                losses = np.array([trial.value for trial in param_trials])
-
-                # Scatter plot
-                ax.scatter(
-                    param_values, losses, alpha=0.6, s=50, edgecolors="black", linewidth=0.5
-                )
-
-                # Highlight best trial
-                best_param_value = best_trial.params.get(param_name)
-                if best_param_value is not None:
-                    ax.scatter(
-                        [best_param_value],
-                        [best_value],
-                        color="red",
-                        s=200,
-                        marker="*",
-                        edgecolors="black",
-                        linewidth=1.5,
-                        zorder=5,
-                    )
-
-                    # Vertical line at optimal parameter value
-                    ax.axvline(
-                        best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7
-                    )
-
-                # Clean up parameter name for label
-                clean_name = param_name.split(".")[-1]
-
-                # Labels
-                ax.set_xlabel(clean_name, fontsize=11, fontweight="bold")
-                ax.set_ylabel("Loss", fontsize=11, fontweight="bold")
-                ax.set_title(f"{param_name}", fontsize=10)
-                ax.grid(True, alpha=0.3)
+                self._plot_param_panel(axes[idx + 1], trials, param_name, best_trial, best_value)
 
             plt.tight_layout()
             return fig, axes
@@ -602,41 +564,7 @@ def visualize(self, figsize=None):
 
         # Plot each parameter
         for idx, param_name in enumerate(param_names):
-            ax = axes[idx]
-
-            # Extract data
-            param_trials = [t for t in trials if param_name in t.params]
-            param_values = np.array([trial.params[param_name] for trial in param_trials])
-            losses = np.array([trial.value for trial in param_trials])
-
-            # Scatter plot
-            ax.scatter(param_values, losses, alpha=0.6, s=50, edgecolors="black", linewidth=0.5)
-
-            # Highlight best trial
-            best_param_value = best_trial.params.get(param_name)
-            if best_param_value is not None:
-                ax.scatter(
-                    [best_param_value],
-                    [best_value],
-                    color="red",
-                    s=200,
-                    marker="*",
-                    edgecolors="black",
-                    linewidth=1.5,
-                    zorder=5,
-                )
-
-                # Vertical line at optimal parameter value
-                ax.axvline(best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7)
-
-            # Clean up parameter name for label
-            clean_name = param_name.split(".")[-1]
-
-            # Labels
-            ax.set_xlabel(clean_name, fontsize=11, fontweight="bold")
-            ax.set_ylabel("Loss", fontsize=11, fontweight="bold")
-            ax.set_title(f"{param_name}", fontsize=10)
-            ax.grid(True, alpha=0.3)
+            self._plot_param_panel(axes[idx], trials, param_name, best_trial, best_value)
 
         # Hide unused subplots
         for idx in range(n_params, len(axes)):
@@ -645,6 +573,35 @@ def visualize(self, figsize=None):
         plt.tight_layout()
         return fig, axes
 
+    def _plot_param_panel(self, ax, trials, param_name, best_trial, best_value):
+        """Scatter one parameter's values vs loss, highlighting the best trial."""
+        param_trials = [t for t in trials if param_name in t.params]
+        param_values = np.array([trial.params[param_name] for trial in param_trials])
+        losses = np.array([trial.value for trial in param_trials])
+
+        ax.scatter(param_values, losses, alpha=0.6, s=50, edgecolors="black", linewidth=0.5)
+
+        best_param_value = best_trial.params.get(param_name)
+        if best_param_value is not None:
+            ax.scatter(
+                [best_param_value],
+                [best_value],
+                color="red",
+                s=200,
+                marker="*",
+                edgecolors="black",
+                linewidth=1.5,
+                zorder=5,
+            )
+            # Vertical line at optimal parameter value
+            ax.axvline(best_param_value, color="red", linestyle="--", linewidth=1.5, alpha=0.7)
+
+        clean_name = param_name.split(".")[-1]
+        ax.set_xlabel(clean_name, fontsize=11, fontweight="bold")
+        ax.set_ylabel("Loss", fontsize=11, fontweight="bold")
+        ax.set_title(f"{param_name}", fontsize=10)
+        ax.grid(True, alpha=0.3)
+
     def _extract_optimization_params(self):
         """Extract OptimizationParameter specs from stored config."""
         param_info = {}
@@ -685,8 +642,6 @@ def grid_search(self, plot_objects=True, figsize=None, return_results=False):
         """
         from itertools import product
 
-        import numpy as np
-
         if self.objective_func is None:
             raise RuntimeError("No objective function set. Use from_constructors() first.")
 
@@ -740,8 +695,7 @@ def grid_search(self, plot_objects=True, figsize=None, return_results=False):
                 gc.collect()
 
         # Find best
-        argfn = np.argmax if self.direction == "maximize" else np.argmin
-        best_idx = argfn([r["loss"] for r in results])
+        best_idx = self._best_index([r["loss"] for r in results])
         best_result = results[best_idx]
 
         # Plot objects
@@ -755,6 +709,11 @@ def grid_search(self, plot_objects=True, figsize=None, return_results=False):
                 "param_grids": param_grids,
             }
 
+    def _best_index(self, losses) -> int:
+        """Index of the best loss given the study direction."""
+        argfn = np.argmax if self.direction == "maximize" else np.argmin
+        return int(argfn(losses))
+
     def _run_reconstruction_with_params(self, params):
         """Run a single reconstruction with given parameters and return the object.
 
@@ -861,8 +820,7 @@ def _plot_grid_objects(self, results, param_names, figsize):
         axes = axes.flatten()
 
         # Find best result
-        losses = [r["loss"] for r in results]
-        best_idx = (np.argmax if self.direction == "maximize" else np.argmin)(losses)
+        best_idx = self._best_index([r["loss"] for r in results])
 
         for idx, result in enumerate(results):
             ax = axes[idx]
diff --git a/src/quantem/diffractive_imaging/probe_models.py b/src/quantem/diffractive_imaging/probe_models.py
index f9c4e4e2a..73a25877b 100644
--- a/src/quantem/diffractive_imaging/probe_models.py
+++ b/src/quantem/diffractive_imaging/probe_models.py
@@ -39,6 +39,7 @@
     real_space_probe,
 )
 from quantem.diffractive_imaging.ptycho_utils import (
+    add_input_noise,
     fourier_shift_expand,
     shift_array,
 )
@@ -1136,7 +1137,8 @@ def __init__(
         self.register_buffer("_pretrain_target", torch.tensor([]))
 
         self._model = model.to(self._device)
-        self._check_roi_shape()
+        if roi_shape is not None:
+            self._check_roi_shape()
         self.set_pretrained_weights(self._model)
 
         self._optimizer = None
@@ -1313,18 +1315,9 @@ def _probe(self) -> torch.Tensor:
 
     def _noisy_model_input(self) -> torch.Tensor:
         """model input with gaussian noise added when _input_noise_std > 0"""
-        if self._input_noise_std > 0.0:
-            noise = (
-                torch.randn(
-                    self.model_input.shape,
-                    dtype=self.dtype,
-                    device=self.device,
-                    generator=self._rng_torch,
-                )
-                * self._input_noise_std
-            )
-            return self.model_input + noise
-        return self.model_input
+        return add_input_noise(
+            self.model_input, self._input_noise_std, self.dtype, self.device, self._rng_torch
+        )
 
     def forward(self, fract_positions: torch.Tensor) -> torch.Tensor:
         """Get shifted probes at fractional positions"""
@@ -1534,8 +1527,6 @@ def backward(self, propagated_gradient, obj_patches):
         )
 
     def _check_roi_shape(self):
-        if not hasattr(self, "_roi_shape"):
-            return
         num_layers = getattr(self.model, "num_layers", None)
         if num_layers is not None:
             if not np.all(np.array(self.roi_shape) % 2**num_layers == 0):
diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py
index dcf8833bd..8a1027678 100644
--- a/src/quantem/diffractive_imaging/ptycho_utils.py
+++ b/src/quantem/diffractive_imaging/ptycho_utils.py
@@ -25,7 +25,6 @@ def __init__(
         train_indices: np.ndarray | None = None,
         val_indices: np.ndarray | None = None,
     ):
-        self.indices = np.arange(num)
         self.batch_size = batch_size if batch_size is not None else num
         self.shuffle = shuffle
         self.rng = rng
@@ -129,6 +128,28 @@ def compute_train_val_split(
     return np.asarray(train_indices, dtype=int), np.asarray(val_indices, dtype=int)
 
 
+def add_input_noise(
+    model_input: torch.Tensor,
+    noise_std: float,
+    dtype: torch.dtype,
+    device: "torch.device | str | int",
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Add gaussian noise to a DIP model input when noise_std > 0."""
+    if noise_std > 0.0:
+        noise = (
+            torch.randn(
+                model_input.shape,
+                dtype=dtype,
+                device=device,
+                generator=generator,
+            )
+            * noise_std
+        )
+        return model_input + noise
+    return model_input
+
+
 @overload
 def fourier_shift_expand(
     array: np.ndarray, positions: np.ndarray, expand_dim: bool = True
@@ -595,7 +616,7 @@ def center_crop_arr(
                 raise ValueError(
                     f"Dimension {i} of shape ({s}) is larger than dimension {i} of arr ({a})."
                 )
-            pad[i] = [(s - a) // 2, -(-(s - a) // 2)]
+            pad[i] = [(s - a) // 2, ceil((s - a) / 2)]
 
     if any(p != [0, 0] for p in pad):
         arr = np.pad(arr, pad_width=pad, mode="constant")
diff --git a/src/quantem/diffractive_imaging/ptychography_lite.py b/src/quantem/diffractive_imaging/ptychography_lite.py
index 1eedacb05..78d41596d 100644
--- a/src/quantem/diffractive_imaging/ptychography_lite.py
+++ b/src/quantem/diffractive_imaging/ptychography_lite.py
@@ -16,6 +16,15 @@
 from quantem.diffractive_imaging.ptychography import Ptychography
 
 
+def _scheduler_spec(scheduler_type: str, scheduler_factor: float) -> dict[str, Any]:
+    scheduler_dict: dict[str, Any] = {
+        "name": "exponential" if scheduler_type == "exp" else scheduler_type
+    }
+    if scheduler_type in ("exp", "plateau"):
+        scheduler_dict["factor"] = scheduler_factor
+    return scheduler_dict
+
+
 class PtychoLite(Ptychography):
     """
     High-level convenience wrapper around Ptychography.
@@ -204,11 +213,7 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         opt_params: dict[str, Any] | None
         scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
-            scheduler_dict: dict[str, Any] = {
-                "name": "exponential" if scheduler_type == "exp" else scheduler_type
-            }
-            if scheduler_type in ("exp", "plateau"):
-                scheduler_dict["factor"] = scheduler_factor
+            scheduler_dict = _scheduler_spec(scheduler_type, scheduler_factor)
             opt_params = {
                 "object": {
                     "name": "adamw",
@@ -447,11 +452,7 @@ def reconstruct(  # pyright: ignore[reportIncompatibleMethodOverride]
         opt_params: dict[str, Any] | None
         scheduler_params: dict[str, Any] | None
         if setup_new_optimizers or (needs_dataset_optimizer and "dataset" not in self.optimizers):
-            scheduler_dict: dict[str, Any] = {
-                "name": "exponential" if scheduler_type == "exp" else scheduler_type
-            }
-            if scheduler_type in ("exp", "plateau"):
-                scheduler_dict["factor"] = scheduler_factor
+            scheduler_dict = _scheduler_spec(scheduler_type, scheduler_factor)
             opt_params = {
                 "object": {
                     "name": "adamw",
diff --git a/src/quantem/diffractive_imaging/ptychography_visualizations.py b/src/quantem/diffractive_imaging/ptychography_visualizations.py
index 679295bc4..6ec82a385 100644
--- a/src/quantem/diffractive_imaging/ptychography_visualizations.py
+++ b/src/quantem/diffractive_imaging/ptychography_visualizations.py
@@ -15,6 +15,12 @@
 
 
 class PtychographyVisualizations(PtychographyBase):
+    def _scalebar_real(self) -> dict[str, Any]:
+        return {"sampling": self.sampling[1], "units": "Å"}
+
+    def _scalebar_recip(self) -> dict[str, Any]:
+        return {"sampling": self.reciprocal_sampling[1], "units": r"$\mathrm{A^{-1}}$"}
+
     def show_obj(
         self,
         obj: np.ndarray | None = None,
@@ -86,7 +92,7 @@ def show_obj(
             titles.extend([t + "Phase", t + "Amplitude"])
             cmaps.extend([ph_cmap, "gray"])
 
-        scalebar = [{"sampling": self.sampling[1], "units": "Å"}] + [None] * (len(ims) - 1)
+        scalebar = [self._scalebar_real()] + [None] * (len(ims) - 1)
 
         show_2d(
             ims,
@@ -172,7 +178,7 @@ def show_obj_fft(
             t += f"Iter {obj_iter} "
 
         if show_obj:
-            obj_scalebar = {"sampling": self.sampling[1], "units": "Å"}
+            obj_scalebar = self._scalebar_real()
             if self.obj_type == "potential":
                 obj_show = obj_pad
             else:  # complex or pure phase just show the phase
@@ -246,7 +252,7 @@ def show_probe(
         if probe_iter != "Final":
             t += f"Iter {probe_iter} "
 
-        scalebar = [{"sampling": self.sampling[1], "units": "Å"}]
+        scalebar = [self._scalebar_real()]
         if sum_probes:
             probes = [np.fft.fftshift(probe.sum(0))]
         else:
@@ -361,7 +367,7 @@ def show_probe_top_bottom(
             top_img = np.fft.fftshift(top_img)
             bottom_img = np.fft.fftshift(bottom_img)
 
-        scalebar = [{"sampling": self.sampling[1], "units": "Å"}, None]
+        scalebar = [self._scalebar_real(), None]
         titles = [f"Top Surface {label}", f"Bottom Surface {label}"]
 
         fig, axs = show_2d(
@@ -398,9 +404,7 @@ def show_fourier_probe(self, probe: np.ndarray | None = None, **kwargs):
                 probe = probe[None, ...]
 
         probes = [np.fft.fftshift(np.fft.fft2(probe[i])) for i in range(len(probe))]
-        scalebar = [{"sampling": self.reciprocal_sampling[1], "units": r"$\mathrm{A^{-1}}$"}] + [
-            None
-        ] * (len(probes) - 1)
+        scalebar = [self._scalebar_recip()] + [None] * (len(probes) - 1)
         if len(probes) > 1:
             titles = self.get_probe_intensities(probe)
             titles = [
@@ -526,7 +530,7 @@ def show_obj_slices(
         titles = [titles_flat[i : i + max_width] for i in range(0, len(titles_flat), max_width)]
 
         scalebars: list = [[None for _ in row] for row in objs]
-        scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
+        scalebars[0][0] = self._scalebar_real()
 
         if interval_type == "quantile":
             norm = {"interval_type": "quantile"}
@@ -857,7 +861,7 @@ def _show_object_iters_only(
 
         scalebars: list = [[None for _ in row] for row in images_grid]
         if scalebars:
-            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
+            scalebars[0][0] = self._scalebar_real()
 
         show_2d(
             images_grid,
@@ -896,7 +900,7 @@ def _show_probe_iters_only(
         # Set up scalebars
         scalebars: list = [[None for _ in row] for row in probes_grid]
         if scalebars:
-            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
+            scalebars[0][0] = self._scalebar_real()
 
         show_2d(
             probes_grid,
@@ -969,8 +973,8 @@ def _show_object_and_probe_iters(
 
         scalebars: list = [[None for _ in row] for row in all_images]
         if scalebars:
-            scalebars[0][0] = {"sampling": self.sampling[1], "units": "Å"}
-            scalebars[0][1] = {"sampling": self.sampling[1], "units": "Å"}
+            scalebars[0][0] = self._scalebar_real()
+            scalebars[0][1] = self._scalebar_real()
 
         show_2d(
             all_images,
@@ -1291,7 +1295,7 @@ def show_fourier_probe_and_amplitudes(
         else:
             amplitudes = self._to_numpy(amplitudes.sum(0))
 
-        scalebar = [{"sampling": self.reciprocal_sampling[1], "units": r"$\mathrm{A^{-1}}$"}]
+        scalebar = [self._scalebar_recip()]
 
         if fft_shift:
             probe_plot = np.fft.fftshift(probe_plot)

From 507acb0c3338f4acd507e187399ed606a26e579d Mon Sep 17 00:00:00 2001
From: arthurmccray <amccray@stanford.edu>
Date: Sun, 2 Aug 2026 12:46:12 -0700
Subject: [PATCH 59/59] fixing seed for INR tests

---
 tests/diffractive_imaging/test_object_inr.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/diffractive_imaging/test_object_inr.py b/tests/diffractive_imaging/test_object_inr.py
index b61f9cb3f..35c006dd5 100644
--- a/tests/diffractive_imaging/test_object_inr.py
+++ b/tests/diffractive_imaging/test_object_inr.py
@@ -291,13 +291,14 @@ def test_potential_obj_type_softplus_opt_in(self):
     def test_potential_identity_default_and_positivity_penalty(self):
         """Default ``potential`` activation is identity (vacuum is exactly 0); the soft
         ``positivity_weight`` penalty drives a forced-negative potential non-negative."""
+        torch.manual_seed(0)  # rng=0 only seeds coordinate sampling; HSiren init uses global RNG
         obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
         obj._initialize_obj((1, 24, 24))
         # identity + zeroed final layer -> vacuum is exactly 0 (not softplus(0) = ln 2)
         assert float(obj._materialize_obj().abs().max()) == pytest.approx(0.0, abs=1e-6)
         # force the whole potential negative via the (zero-weight) final-layer bias
         with torch.no_grad():
-            obj.model.net[-2].bias.fill_(-0.5) # type:ignore 
+            obj.model.net[-2].bias.fill_(-0.5)  # type:ignore
         assert float(obj._materialize_obj().min()) == pytest.approx(-0.5, abs=1e-3)
         obj.constraints = {"positivity_weight": 1.0}
         assert float(obj._sampled_positivity_loss(1.0)) == pytest.approx(0.5, abs=0.05)
@@ -314,7 +315,7 @@ def test_fix_potential_baseline_gauge(self):
         obj = ObjectINR.from_uniform(num_slices=1, obj_type="potential", hidden_features=32, rng=0)
         obj._initialize_obj((1, 16, 16))
         with torch.no_grad():
-            obj.model.net[-2].bias.fill_(1.0) # type:ignore  # constant +1 background
+            obj.model.net[-2].bias.fill_(1.0)  # type:ignore  # constant +1 background
         raw = obj._materialize_obj()
         assert float(raw.min()) == pytest.approx(1.0, abs=0.2)
         obj.constraints = {"fix_potential_baseline": True}