diff --git a/src/quantem/core/io/serialize.py b/src/quantem/core/io/serialize.py index 4dd03984..75f8784f 100644 --- a/src/quantem/core/io/serialize.py +++ b/src/quantem/core/io/serialize.py @@ -789,7 +789,9 @@ def _serialize_container( # Handle list/tuple containers if isinstance(value, (list, tuple)): - group.attrs["_container_type"] = type(value).__name__ + # normalize subclasses (torch.Size, namedtuples, ...) to their base container, + # since _deserialize_container only knows how to rebuild list/tuple + group.attrs["_container_type"] = "tuple" if isinstance(value, tuple) else "list" # Fast-path: homogeneous numeric scalars → single ndarray try: is_all_numeric = len(value) > 0 and all( diff --git a/src/quantem/core/utils/imaging_utils.py b/src/quantem/core/utils/imaging_utils.py index d352a051..9507bc53 100644 --- a/src/quantem/core/utils/imaging_utils.py +++ b/src/quantem/core/utils/imaging_utils.py @@ -23,6 +23,19 @@ def dft_upsample( Manuel Guizar-Sicairos, Samuel T. Thurman, and James R. Fienup, "Efficient subpixel image registration algorithms," Opt. Lett. 33, 156-158 (2008). http://www.sciencedirect.com/science/article/pii/S0045790612000778 + + Evaluates the inverse transform of ``F`` on a ``(2*du+1)`` square grid of spacing + ``1/up``, centered on ``shift``, where ``du = ceil(1.5 * up)``. The center tap is + therefore index ``du``, not ``up``. + + Parameters + ---------- + F : ndarray + Fourier-domain array, in FFT (unshifted) order. + up : int + Upsampling factor. The output samples ``shift + arange(-du, du+1) / up``. + shift : tuple of float + Position, in pixels of the *real-space* image, to center the fine grid on. """ if device == "gpu": import cupy as cp # type: ignore @@ -33,17 +46,17 @@ def dft_upsample( M, N = F.shape du = np.ceil(1.5 * up).astype(int) - row = np.arange(-du, du + 1) - col = np.arange(-du, du + 1) - r_shift = shift[0] - M // 2 - c_shift = shift[1] - N // 2 + span = np.arange(-du, du + 1) - kern_row = np.exp( - -2j * np.pi / (M * up) * np.outer(row, xp.fft.ifftshift(xp.arange(M)) - M // 2 + r_shift) - ) - kern_col = np.exp( - -2j * np.pi / (N * up) * np.outer(xp.fft.ifftshift(xp.arange(N)) - N // 2 + c_shift, col) - ) + # frequency of FFT bin m, i.e. `ifftshift` of the centered ramp + freq_row = xp.fft.ifftshift(xp.arange(M)) - M // 2 + freq_col = xp.fft.ifftshift(xp.arange(N)) - N // 2 + + # The offset re-centers the *output* sampling positions -- `shift + span/up` -- rather + # than translating the input frequencies, and the sign is the inverse transform's, so + # that this agrees with `ifft2` where the two grids coincide. + kern_row = xp.exp(2j * np.pi / (M * up) * np.outer(span + up * shift[0], freq_row)) + kern_col = xp.exp(2j * np.pi / (N * up) * np.outer(freq_col, span + up * shift[1])) return xp.real(kern_row @ F @ kern_col) @@ -142,7 +155,9 @@ def parabolic_peak(v): except (IndexError, ValueError): dxf = dyf = 0.0 - shifts = np.array([x0, y0]) + (np.array(peak) - upsample_factor) / upsample_factor + # the fine grid is centered on its middle tap, `ceil(1.5 * upsample_factor)` + center = np.ceil(1.5 * upsample_factor).astype(int) + shifts = np.array([x0, y0]) + (np.array(peak) - center) / upsample_factor shifts += np.array([dxf, dyf]) / upsample_factor shifts = (shifts + 0.5 * np.array(cc.shape)) % cc.shape - 0.5 * np.array(cc.shape) diff --git a/src/quantem/diffractive_imaging/__init__.py b/src/quantem/diffractive_imaging/__init__.py index 9c1520dd..0ecc2225 100644 --- a/src/quantem/diffractive_imaging/__init__.py +++ b/src/quantem/diffractive_imaging/__init__.py @@ -27,12 +27,25 @@ from quantem.diffractive_imaging.complex_probe import ( real_space_probe as real_space_probe, fourier_space_probe as fourier_space_probe, + FourierProbe as FourierProbe, ) from quantem.diffractive_imaging.direct_ptychography import ( DirectPtychography as DirectPtychography, ) +from quantem.diffractive_imaging.direct_ptychography_base import ( + OptimizationParameter as OptimizationParameter, +) + +from quantem.diffractive_imaging.direct_ptychography_montage import ( + DirectPtychographyMontage as DirectPtychographyMontage, +) + +from quantem.diffractive_imaging.direct_ptycho_utils import ( + estimate_frame_drift as estimate_frame_drift, +) + from quantem.diffractive_imaging.origin_models import ( CenterOfMassOriginModel as CenterOfMassOriginModel, ) diff --git a/src/quantem/diffractive_imaging/complex_probe.py b/src/quantem/diffractive_imaging/complex_probe.py index 8bf742c3..725f39cf 100644 --- a/src/quantem/diffractive_imaging/complex_probe.py +++ b/src/quantem/diffractive_imaging/complex_probe.py @@ -2,6 +2,7 @@ from collections import defaultdict from typing import Mapping, Tuple +import numpy as np import torch from numpy.typing import NDArray @@ -308,44 +309,292 @@ def aberration_surface_cartesian_gradients( return dchi_dx, dchi_dy +class FourierProbe: + """``psi(k)``, the probe in the detector plane, sampled wherever ``gamma_factor`` needs it. + + Two flavours behind one interface: + + - :meth:`from_aberrations` -- an aperture times ``exp(-i chi)``, evaluated in closed form + at any ``k``. What every electron entry point uses, and identical to calling + :func:`evaluate_probe` directly. + - :meth:`from_array` -- a measured or independently reconstructed complex array, sampled + on its own reciprocal grid. For a probe that no aperture and low-order aberrations + describe: a zone plate with a central stop, a speckled illumination, an X-ray optic. + + The difference that matters is that an analytic probe can be *evaluated* at the arbitrary + ``k -/+ q`` that the overlap function asks for, while an array can only be *sampled*. See + :meth:`at` for what that costs. + """ + + def __init__( + self, + wavelength: float, + *, + array: torch.Tensor | None = None, + reciprocal_sampling: Tuple[float, float] | None = None, + semiangle_cutoff: float | None = None, + aberration_coefs: Mapping[str, float | torch.Tensor] | None = None, + angular_sampling: Tuple[float, float] | None = None, + soft_edges: bool = True, + interpolation: str = "exact", + ): + self.wavelength = float(wavelength) + self.array = array + self.reciprocal_sampling = reciprocal_sampling + self.semiangle_cutoff = semiangle_cutoff + self.aberration_coefs = aberration_coefs or {} + self.angular_sampling = angular_sampling + self.soft_edges = soft_edges + self.interpolation = interpolation + + @classmethod + def from_aberrations( + cls, + wavelength: float, + semiangle_cutoff: float, + angular_sampling: Tuple[float, float], + aberration_coefs: Mapping[str, float | torch.Tensor] = {}, + soft_edges: bool = True, + ) -> "FourierProbe": + """The analytic probe: a soft or hard aperture times ``exp(-i chi(k))``.""" + return cls( + wavelength, + semiangle_cutoff=semiangle_cutoff, + angular_sampling=angular_sampling, + aberration_coefs=aberration_coefs, + soft_edges=soft_edges, + ) + + @classmethod + def from_array( + cls, + array: torch.Tensor | NDArray, + reciprocal_sampling: Tuple[float, float], + wavelength: float, + normalize: bool = True, + interpolation: str = "exact", + ) -> "FourierProbe": + """An empirical complex probe on the detector's own reciprocal grid. + + Parameters + ---------- + array : complex array + ``psi(k)``, corner-centered to match :func:`spatial_frequencies`. A probe + saved with its maximum at the array centre needs ``ifftshift`` first. + reciprocal_sampling : tuple of float + ``dq`` per axis, in inverse Angstrom. Its reciprocal, ``1 / (n * dq)`` per axis, + is the probe's real-space field of view, which :meth:`at` needs. + normalize : bool + Scale to unit total intensity, matching :func:`fourier_space_probe`. The + deconvolution kernels are scale invariant but the bright-field weight that + normalizes the finished object is not, so leaving this off changes the object's + absolute scale. + interpolation : {"exact", "bilinear"} + What to do when ``k -/+ q`` falls between grid points -- see :meth:`at`. + """ + tensor = array if isinstance(array, torch.Tensor) else torch.as_tensor(array) + if not torch.is_complex(tensor): + raise ValueError( + f"`array` must be a complex probe psi(k), got dtype {tensor.dtype}. Pass " + "`amplitude * exp(1j * phase)`, or a real array cast to complex if the " + "probe really has no phase structure." + ) + if tensor.ndim != 2: + raise ValueError(f"`array` must be 2D (Ny, Nx), got shape {tuple(tensor.shape)}") + if normalize: + tensor = tensor / tensor.abs().square().sum().sqrt() + + return cls( + wavelength, + array=tensor, + reciprocal_sampling=tuple(float(s) for s in reciprocal_sampling), + interpolation=interpolation, + ) + + @property + def is_analytic(self) -> bool: + """Whether the probe has aberration coefficients behind it. + + The parallax kernel, the ``sign(sin(chi))`` phase flip, the defocus gradient and any + hyperparameter search over aberrations are only defined when this is true. + """ + return self.array is None + + def to(self, device) -> "FourierProbe": + if self.array is None: + return self + moved = FourierProbe( + self.wavelength, + array=self.array.to(device), + reciprocal_sampling=self.reciprocal_sampling, + interpolation=self.interpolation, + ) + return moved + + def resampled_to(self, reciprocal_sampling: Tuple[float, float]) -> "FourierProbe": + """The same probe on a finer reciprocal grid, by zero-padding it in real space. + + This is not an approximation. ``psi`` is the transform of a probe confined to the + field of view its sampling implies, so it is band limited, and zero-padding then + transforming back is exactly the sinc interpolation that band limit licenses -- + unlike bilinear sampling, which is not. + + The requested step must divide the current one by a whole number per axis: refining + by a non-integer factor would need the real-space probe resampled rather than merely + extended, which *is* an approximation. + + This is what makes an empirical probe usable on a canvas larger than the probe's own + field of view: the canvas needs ``q`` spaced ``1 / canvas_fov``, and the probe + supplies ``1 / probe_fov``, so the padding factor is ``canvas_fov / probe_fov``. + """ + if self.array is None: + return self # analytic probes are evaluated, not sampled + + current = np.asarray(self.reciprocal_sampling, dtype=float) + target = np.asarray(reciprocal_sampling, dtype=float) + ratio = current / target + rounded = np.round(ratio) + if np.any(rounded < 1) or np.any(np.abs(ratio - rounded) > 1e-6): + raise ValueError( + f"An empirical probe can only be refined onto a reciprocal grid whose step " + f"divides its own by a whole number, but " + f"{tuple(self.reciprocal_sampling)} / {tuple(float(t) for t in target)} = " # ty:ignore[not-iterable] + f"{tuple(np.round(ratio, 4))}. Equivalently, the canvas field of view must " + f"be a whole multiple of the probe's." + ) + + factor = rounded.astype(int) + if np.all(factor == 1): + return self + + shape = tuple(int(n * f) for n, f in zip(self.array.shape, factor)) + real_space = torch.fft.ifft2(self.array) + padded = torch.zeros(shape, dtype=real_space.dtype, device=real_space.device) + # keep the corner-centered quadrants where they belong on the larger grid + n_rows, n_cols = self.array.shape + for row_slice, row_source in ( + (slice(0, (n_rows + 1) // 2), slice(0, (n_rows + 1) // 2)), + (slice(shape[0] - n_rows // 2, shape[0]), slice((n_rows + 1) // 2, n_rows)), + ): + for col_slice, col_source in ( + (slice(0, (n_cols + 1) // 2), slice(0, (n_cols + 1) // 2)), + (slice(shape[1] - n_cols // 2, shape[1]), slice((n_cols + 1) // 2, n_cols)), + ): + padded[row_slice, col_slice] = real_space[row_source, col_source] + + # no rescaling: the padded transform evaluated on the coarse sub-lattice is + # `sum_m real[m] exp(-2i.pi.k.m/n)`, which is the original psi exactly + refined = torch.fft.fft2(padded) + return FourierProbe( + self.wavelength, + array=refined, + reciprocal_sampling=tuple(float(t) for t in target), # ty:ignore[not-iterable] + interpolation=self.interpolation, + ) + + def at(self, kx: torch.Tensor, ky: torch.Tensor) -> torch.Tensor: + """``psi`` at the given spatial frequencies, in inverse Angstrom. + + For an analytic probe this is a closed-form evaluation and any ``k`` is fine. + + An array probe can only be read at its own grid points, and ``gamma_factor`` asks for + ``k -/+ q`` -- detector frequencies offset by canvas frequencies. Those coincide only + when the two grids are commensurate, that is when + + (probe field of view) / (canvas field of view) + + is an integer, the probe field of view being ``1 / reciprocal_sampling``. When it is, + every lookup is an exact gather. When it is not, ``interpolation="exact"`` raises + rather than quietly interpolating: a speckled probe varies over a few detector pixels + (0.83 amplitude spread pixel-to-pixel on the X-ray data this was written for), so + bilinear sampling is an approximation worth opting into explicitly. + + The usual fix is to zero-pad the probe in real space -- which refines its reciprocal + grid without changing the probe -- until the ratio is an integer. + """ + if self.array is None: + k, phi = polar_coordinates(kx, ky) + return evaluate_probe( + k * self.wavelength, + phi, + self.semiangle_cutoff, + self.angular_sampling, + self.wavelength, + self.soft_edges, + None, + self.aberration_coefs, + ) + + n_rows, n_cols = self.array.shape + dq_row, dq_col = self.reciprocal_sampling # ty:ignore[not-iterable] + row = kx / dq_row + col = ky / dq_col + + if self.interpolation == "bilinear": + return self._bilinear(row, col) + + rounded_row, rounded_col = torch.round(row), torch.round(col) + # 1e-3 of a detector pixel: loose enough for float32 k-grids, tight enough that a + # genuinely incommensurate canvas never slips through + offset = torch.maximum((row - rounded_row).abs().max(), (col - rounded_col).abs().max()) + if float(offset) > 1e-3: + fov = (1 / (n_rows * dq_row), 1 / (n_cols * dq_col)) + raise ValueError( + f"An empirical probe can only be sampled on its own reciprocal grid, but the " + f"requested frequencies miss it by up to {float(offset):.3f} of a detector " + f"pixel. The canvas field of view must divide the probe's, " + f"{fov[0]:.1f} x {fov[1]:.1f} Angstrom: either choose one that does, " + f"zero-pad the probe in real space to refine its grid, or pass " + f"`interpolation='bilinear'` to accept the sampling error." + ) + + return self._gather(rounded_row, rounded_col) + + def _gather(self, row: torch.Tensor, col: torch.Tensor) -> torch.Tensor: + """``psi`` at signed frequency indices, zero outside the detector. + + The indices are frequencies in units of ``dq``, so they run over + ``[-n//2, (n+1)//2)`` -- the same range ``fftfreq`` produces. Beyond it the detector + measured nothing, so the probe is zero there. Wrapping instead would alias the + opposite edge of the aperture into the answer, which is wrong wherever the grid is + cropped close to the probe: on the electron fixtures, where the bright-field mask is + cropped to the disk, that alone moved the reconstruction by 13%. + """ + array = self.array + n_rows, n_cols = array.shape # ty:ignore[possibly-unbound-attribute] + row, col = row.long(), col.long() + inside = ( + (row >= -(n_rows // 2)) + & (row < (n_rows + 1) // 2) + & (col >= -(n_cols // 2)) + & (col < (n_cols + 1) // 2) + ) + values = array[row % n_rows, col % n_cols] + return values * inside + + def _bilinear(self, row: torch.Tensor, col: torch.Tensor) -> torch.Tensor: + row0, col0 = torch.floor(row), torch.floor(col) + drow, dcol = row - row0, col - col0 + return ( + self._gather(row0, col0) * ((1 - drow) * (1 - dcol)) + + self._gather(row0 + 1, col0) * (drow * (1 - dcol)) + + self._gather(row0, col0 + 1) * ((1 - drow) * dcol) + + self._gather(row0 + 1, col0 + 1) * (drow * dcol) + ) + + def gamma_factor( qmks: tuple[torch.Tensor, torch.Tensor], qpks: tuple[torch.Tensor, torch.Tensor], cmplx_probe_at_k: torch.Tensor, - wavelength: float, - semiangle_cutoff: float, - soft_edges: bool, - aberration_coefs: Mapping[str, float | torch.Tensor], - angular_sampling: Tuple[float, float], + probe: FourierProbe, asymmetric_version: bool = True, normalize: bool = True, ): - """ """ - - q_m, phi_m = polar_coordinates(*qmks) - q_p, phi_p = polar_coordinates(*qpks) + """The overlap function ``Gamma(k, q)`` of a probe with itself, at ``k -/+ q``.""" - probe_m = evaluate_probe( - q_m * wavelength, - phi_m, - semiangle_cutoff, - angular_sampling, - wavelength, - soft_edges, - None, - aberration_coefs, - ) - - probe_p = evaluate_probe( - q_p * wavelength, - phi_p, - semiangle_cutoff, - angular_sampling, - wavelength, - soft_edges, - None, - aberration_coefs, - ) + probe_m = probe.at(*qmks) + probe_p = probe.at(*qpks) if asymmetric_version: gamma = probe_m * cmplx_probe_at_k.conj() - probe_p.conj() * cmplx_probe_at_k diff --git a/src/quantem/diffractive_imaging/direct_ptycho_utils.py b/src/quantem/diffractive_imaging/direct_ptycho_utils.py index ce08bfeb..9933dab5 100644 --- a/src/quantem/diffractive_imaging/direct_ptycho_utils.py +++ b/src/quantem/diffractive_imaging/direct_ptycho_utils.py @@ -1,5 +1,5 @@ from collections import defaultdict -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from quantem.core import config @@ -10,14 +10,20 @@ import torch import math +import warnings +import numpy as np from tqdm.auto import tqdm from quantem.core.utils.imaging_utils import cross_correlation_shift_torch, unwrap_phase_2d_torch +from quantem.core.utils.validators import validate_tensor from quantem.diffractive_imaging.complex_probe import ( spatial_frequencies, ) +# bilinear corner offsets: (row offset, col offset) +_BILINEAR_CORNERS = ((0, 0), (1, 0), (0, 1), (1, 1)) + # fmt: off ABERRATION_PRESETS = { "defocus": ["C10"], @@ -644,3 +650,889 @@ def _crop_corner_centered_mask(mask: torch.Tensor, bf_mask_padding_px: int): y0, y1 = ys.min() - px, ys.max() + px + 1 x0, x1 = xs.min() - px, xs.max() + px + 1 return torch.fft.ifftshift(mask_c[y0:y1, x0:x1]) + + +def preferred_float_dtype(device) -> torch.dtype: + """Widest float the device supports: float64 everywhere except MPS, which has none. + + Used for splat accumulators and for scan coordinates. On MPS the float32 fallback + resolves canvas positions to roughly 1e-3 pixels at a 10k-pixel canvas, well below the + sub-pixel detail any of this is trying to preserve. + """ + return torch.float32 if torch.device(device).type == "mps" else torch.float64 + + +def allocate_splat_buffers( + canvas_shape: tuple[int, int], + device, + dtype: torch.dtype | None = None, + accumulate_squares: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Flat, zeroed ``(sum_w, sum_wv, sum_wv2)`` accumulators for :func:`scatter_add_splat`.""" + if dtype is None: + dtype = preferred_float_dtype(device) + numel = canvas_shape[0] * canvas_shape[1] + + def _zeros(): + return torch.zeros(numel, device=device, dtype=dtype) + + return _zeros(), _zeros(), _zeros() if accumulate_squares else None + + +def _deposition_corners(coords: torch.Tensor, interpolation: str): + """``(base, frac, corners)`` for a sub-pixel deposition scheme. + + ``frac`` is ``None`` for nearest-neighbour, where every corner weight is one. + """ + if interpolation == "bilinear": + base = torch.floor(coords) + return base, coords - base, _BILINEAR_CORNERS + if interpolation == "nearest": + return torch.round(coords), None, ((0, 0),) + raise ValueError(f"`interpolation` must be 'bilinear' or 'nearest', got {interpolation!r}") + + +def _resolve_indices(row, col, weights, n_rows: int, n_cols: int, boundary: str): + """Apply the boundary rule and flatten to ``(flat_indices, weights)``.""" + if boundary == "wrap": + row = row % n_rows + col = col % n_cols + elif boundary == "pad": + # clamp the indices and zero their weights instead of masking, which would + # need a `nonzero()` and hence a device->host synchronization + valid = (row >= 0) & (row < n_rows) & (col >= 0) & (col < n_cols) + row = row.clamp(0, n_rows - 1) + col = col.clamp(0, n_cols - 1) + weights = weights * valid + else: + raise ValueError(f"`boundary` must be 'wrap' or 'pad', got {boundary!r}") + + return (row * n_cols + col).reshape(-1), weights + + +def scatter_add_splat( + values: torch.Tensor, + coords: torch.Tensor, + canvas_shape: tuple[int, int], + *, + boundary: Literal["wrap", "pad"] = "wrap", + interpolation: Literal["bilinear", "nearest"] = "bilinear", + out: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """ + Batched sub-pixel scatter-add of values onto a 2D canvas. + + This is the torch counterpart of :func:`quantem.core.utils.imaging_utils.bilinear_kde`'s + accumulation stage: it splits each point across the four surrounding pixels with bilinear + weights and accumulates ``w``, ``w * v`` and ``w * v**2`` via ``index_add_``. Unlike that + function it takes a leading batch axis, runs on any torch device, offers a drop + (``"pad"``) as well as a wrap boundary, and does no smoothing or normalization. + + Parameters + ---------- + values : torch.Tensor + ``(..., T)`` values to deposit. + coords : torch.Tensor + ``(..., T, 2)`` canvas coordinates in pixels, ordered ``(row, col)``. Broadcast + against ``values``. + canvas_shape : tuple of int + ``(n_rows, n_cols)`` of the output canvas. + boundary : {"wrap", "pad"} + ``"wrap"`` wraps coordinates periodically; ``"pad"`` drops out-of-bounds points. + interpolation : {"bilinear", "nearest"} + Sub-pixel deposition scheme. + out : tuple of torch.Tensor, optional + Pre-allocated flat buffers ``(sum_w, sum_wv, sum_wv2)`` to accumulate into, as + returned by :func:`allocate_splat_buffers`. ``sum_wv2`` may be ``None`` to skip + the sum-of-squares. If omitted, fresh buffers are allocated. + + Returns + ------- + sum_w, sum_wv, sum_wv2 : torch.Tensor + Flat ``(n_rows * n_cols,)`` accumulators. ``sum_wv2`` is ``None`` if not requested. + + Notes + ----- + ``index_add_`` uses atomics on CUDA and is therefore not bit-reproducible there; + compare results with a tolerance rather than for exact equality. + """ + n_rows, n_cols = int(canvas_shape[0]), int(canvas_shape[1]) + + if out is None: + out = allocate_splat_buffers(canvas_shape, coords.device) + sum_w, sum_wv, sum_wv2 = out + dtype = sum_w.dtype + + values = values.to(dtype) + coords = coords.to(dtype) + + base, frac, corners = _deposition_corners(coords, interpolation) + base_row = base[..., 0].to(torch.int64) + base_col = base[..., 1].to(torch.int64) + + for d_row, d_col in corners: + if frac is None: + weights = torch.ones_like(values) + else: + w_row = frac[..., 0] if d_row else 1 - frac[..., 0] + w_col = frac[..., 1] if d_col else 1 - frac[..., 1] + weights = w_row * w_col + weights = weights.expand_as(values) if weights.shape != values.shape else weights + + flat_indices, weights = _resolve_indices( + base_row + d_row, base_col + d_col, weights, n_rows, n_cols, boundary + ) + flat_weights = weights.reshape(-1) + flat_values = values.reshape(-1) + + sum_w.index_add_(0, flat_indices, flat_weights) + sum_wv.index_add_(0, flat_indices, flat_weights * flat_values) + if sum_wv2 is not None: + sum_wv2.index_add_(0, flat_indices, flat_weights * flat_values * flat_values) + + return sum_w, sum_wv, sum_wv2 + + +def splat_stack( + values: torch.Tensor, + coords: torch.Tensor, + canvas_shape: tuple[int, int], + *, + boundary: Literal["wrap", "pad"] = "wrap", + interpolation: Literal["bilinear", "nearest"] = "nearest", +) -> torch.Tensor: + """ + Splat each row of a batch onto its own canvas: ``(B, T)`` -> ``(B, n_rows, n_cols)``. + + :func:`scatter_add_splat` accumulates a whole batch into one shared canvas, which is what + the parallax kernel wants. A convolution kernel needs each bright-field image separately, + so that it can be convolved with that image's own kernel before the sum -- see + :func:`splat_and_convolve` and :func:`convolve_stack_fourier`. + + Deposition matches :func:`scatter_add_splat` exactly, so splatting and then convolving is + the same operator as :func:`scatter_add_convolve`, only reorganized. + """ + n_rows, n_cols = int(canvas_shape[0]), int(canvas_shape[1]) + batch = int(values.shape[0]) + dtype = values.dtype if values.is_floating_point() else torch.float32 + + flat = torch.zeros(batch * n_rows * n_cols, device=values.device, dtype=dtype) + values = values.to(dtype) + coords = coords.to(dtype) + + base, frac, corners = _deposition_corners(coords, interpolation) + base_row = base[..., 0].to(torch.int64) + base_col = base[..., 1].to(torch.int64) + # offset each batch element into its own slab of the flat buffer + slab = ( + torch.arange(batch, device=values.device, dtype=torch.int64).view(-1, 1) * n_rows * n_cols + ) + + for d_row, d_col in corners: + if frac is None: + weights = torch.ones_like(values) + else: + w_row = frac[..., 0] if d_row else 1 - frac[..., 0] + w_col = frac[..., 1] if d_col else 1 - frac[..., 1] + weights = w_row * w_col + weights = weights.expand_as(values) if weights.shape != values.shape else weights + + indices, weights = _resolve_indices( + base_row + d_row, base_col + d_col, weights, n_rows, n_cols, boundary + ) + indices = (indices.view(batch, -1) + slab).reshape(-1) + flat.index_add_(0, indices, (weights * values).reshape(-1)) + + return flat.view(batch, n_rows, n_cols) + + +def splat_and_convolve( + values: torch.Tensor, + coords: torch.Tensor, + canvas_shape: tuple[int, int], + stencil_weights: torch.Tensor, + radius: int, + *, + boundary: Literal["wrap", "pad"] = "wrap", + interpolation: Literal["bilinear", "nearest"] = "nearest", +) -> torch.Tensor: + """ + Splat each bright-field image, then convolve it with its own square kernel. + + The same operator as :func:`scatter_add_convolve`, reorganized so the convolution is a + grouped ``conv2d`` rather than a loop over taps. Measured on MPS with 167k bright-field + pixels and a 180x140 canvas, a radius-8 stencil takes 5.4 s here against 78 s there. + + ``stencil_weights`` is ``(B, (2 * radius + 1) ** 2)``, ordered as the ``"ij"`` meshgrid + :meth:`DirectPtychographyMontage._return_kernel_stencil` builds. + + Returns ``(B, n_rows, n_cols)``; sum over the batch to accumulate. + """ + n_rows, n_cols = int(canvas_shape[0]), int(canvas_shape[1]) + size = 2 * radius + 1 + + if boundary == "wrap": + stack = splat_stack( + values, coords, canvas_shape, boundary="wrap", interpolation=interpolation + ) + stack = torch.nn.functional.pad(stack.unsqueeze(0), (radius,) * 4, mode="circular") + else: + # `scatter_add_convolve` tests the boundary at the *deposit* position, so a point + # just outside still contributes inward through the kernel. Growing the canvas by + # the radius keeps those points; the unpadded conv2d below crops back. + grown = (n_rows + 2 * radius, n_cols + 2 * radius) + stack = splat_stack( + values, coords + radius, grown, boundary="pad", interpolation=interpolation + ).unsqueeze(0) + + batch = int(values.shape[0]) + # torch conv2d correlates rather than convolves, so flip the kernel + weight = torch.flip(stencil_weights.reshape(batch, 1, size, size), dims=(-2, -1)) + convolved = torch.nn.functional.conv2d(stack.to(weight.dtype), weight, groups=batch) + return convolved.view(batch, n_rows, n_cols) + + +def convolve_stack_fourier( + stack: torch.Tensor, + kernel_fourier: torch.Tensor, +) -> torch.Tensor: + """ + Multiply each canvas's transform by its own Fourier kernel, and return the sum in ``q``. + + Exact, where a stencil is truncated -- which matters because the SSB, OBF and + matched-filter kernels are never compact in real space (their transforms have ``r**-1.5`` + tails). It is also asymptotically cheaper for a kernel that spans the canvas: one FFT per + bright-field image against ``(2 * radius + 1) ** 2`` taps per point. + + The convolution is circular, as it is for any Fourier method -- ``DirectPtychography`` + included. Zero-pad the canvas beforehand to get a linear one. + """ + return (torch.fft.fft2(stack) * kernel_fourier).sum(0) + + +def scatter_add_convolve( + values: torch.Tensor, + coords: torch.Tensor, + canvas_shape: tuple[int, int], + stencil_offsets: torch.Tensor, + stencil_weights: torch.Tensor, + *, + boundary: Literal["wrap", "pad"] = "wrap", + interpolation: Literal["bilinear", "nearest"] = "bilinear", + out: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Scatter-add each value onto a canvas spread over a complex convolution stencil. + + Where :func:`scatter_add_splat` deposits a value at a point, this deposits + ``value * stencil_weights[b, s]`` at ``coords + stencil_offsets[s]`` for every tap ``s`` + -- the real-space form of multiplying a bright-field image by a Fourier kernel. The + parallax kernel is the special case of a single tap of weight one, which is why the + montage is cheap; SSB and OBF kernels need hundreds of taps. + + Kept separate from :func:`scatter_add_splat` rather than folded into it: the accumulator + here is complex, and the weight and sum-of-squares buffers that drive + ``variance_loss`` have no meaning when the taps are kernel weights rather than a + partition of unity. + + Parameters + ---------- + values : torch.Tensor + ``(B, T)`` values to deposit. + coords : torch.Tensor + ``(B, T, 2)`` canvas coordinates in pixels, ordered ``(row, col)``. + canvas_shape : tuple of int + ``(n_rows, n_cols)`` of the output canvas. + stencil_offsets : torch.Tensor + ``(S, 2)`` integer pixel offsets of the stencil taps. + stencil_weights : torch.Tensor + ``(B, S)`` complex weight of each tap, per batch element. + out : torch.Tensor, optional + Flat complex ``(n_rows * n_cols,)`` accumulator to add into. + + Returns + ------- + torch.Tensor + The flat complex accumulator. + + Notes + ----- + Taps are looped over rather than broadcast, so peak memory stays ``O(B * T)`` however + large the stencil is. ``index_add_`` uses atomics on CUDA and is not bit-reproducible + there. + """ + n_rows, n_cols = int(canvas_shape[0]), int(canvas_shape[1]) + + if out is None: + out = torch.zeros(n_rows * n_cols, device=values.device, dtype=torch.complex64) + values = values.to(out.dtype) + stencil_weights = stencil_weights.to(out.dtype) + + base, frac, corners = _deposition_corners(coords, interpolation) + base_row = base[..., 0].to(torch.int64) + base_col = base[..., 1].to(torch.int64) + + for d_row, d_col in corners: + if frac is None: + corner_weight = None + else: + w_row = frac[..., 0] if d_row else 1 - frac[..., 0] + w_col = frac[..., 1] if d_col else 1 - frac[..., 1] + corner_weight = (w_row * w_col).to(out.dtype) + + for tap, (s_row, s_col) in enumerate(stencil_offsets.tolist()): + contribution = values * stencil_weights[:, tap : tap + 1] + if corner_weight is not None: + contribution = contribution * corner_weight + + flat_indices, contribution = _resolve_indices( + base_row + d_row + int(s_row), + base_col + d_col + int(s_col), + contribution, + n_rows, + n_cols, + boundary, + ) + out.index_add_(0, flat_indices, contribution.reshape(-1)) + + return out + + +def validate_probe_positions(positions): + """``(N, 2)`` float64 probe positions in Angstrom, from an array, tensor or Dataset2d.""" + from quantem.core.datastructures import Dataset2d + + if isinstance(positions, Dataset2d): + if str(positions.units[0]) != "A": + raise ValueError(f"`positions` must be given in 'A', got {tuple(positions.units)!r}") + positions = positions.array + + positions = np.asarray( + positions.detach().cpu().numpy() if hasattr(positions, "detach") else positions, + dtype=np.float64, + ) + if positions.ndim != 2 or positions.shape[1] != 2: + raise ValueError(f"`positions` must have shape (N, 2), got {positions.shape}") + return positions + + +def infer_scan_sampling(positions_ang, max_points: int = 4096): + """Median nearest-neighbour spacing, isotropic, from a subsample of positions.""" + points = positions_ang + if len(points) > max_points: + points = points[np.linspace(0, len(points) - 1, max_points).astype(int)] + + distances = np.linalg.norm(points[:, None, :] - points[None, :, :], axis=-1) + np.fill_diagonal(distances, np.inf) + spacing = float(np.median(distances.min(axis=1))) + return (spacing, spacing) + + +def build_vbf_stack_from_dataset3d( + dataset, + positions, + scan_sampling, + device: str | int = "cpu", + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin=None, + force_fitted_origin=None, + rotation_angle: float | None = None, + intensity_threshold: float = 0.5, + normalization_order: int = 0, + bf_mask=None, +): + """ + Origin-correct and mask an ungridded diffraction stack into a flat vBF stack. + + The ungridded counterpart of :func:`build_vbf_stack_from_dataset4d`, shared by + ``DirectPtychographyMontage.from_dataset3d`` and ``DirectPtychography.from_dataset3d`` + so the two entry points cannot drift apart. + + Returns + ------- + vbf_stack : torch.Tensor + ``(N_bf, N)`` bright-field intensities, flattened over scan positions. + positions_px : ndarray + ``(N, 2)`` positions in scan pixels, anchored at the bounding-box corner. + bf_mask_dataset : Dataset2d + scan_gpts : tuple of int + Grid that just covers the positions. + scan_sampling : tuple of float + Resolved pixel size, with ``"auto"`` replaced by the inferred value. + rotation_angle : float + scan_origin : tuple of float + The bounding-box corner the positions were anchored at, in Angstrom. Keeping it + lets ``scan_origin + positions_px * scan_sampling`` recover the input positions, + so several acquisitions of the same region stay in one coordinate frame. + """ + from quantem.core.datastructures import Dataset2d + + positions_ang = validate_probe_positions(positions) + if positions_ang.shape[0] != dataset.shape[0]: + raise ValueError( + f"`positions` has {positions_ang.shape[0]} rows but `dataset` has " + f"{dataset.shape[0]} diffraction patterns." + ) + + if isinstance(scan_sampling, str): + if scan_sampling != "auto": + raise ValueError(f"`scan_sampling` must be a pair or 'auto', got {scan_sampling!r}") + scan_sampling = infer_scan_sampling(positions_ang) + warnings.warn( + f"Inferred scan_sampling={scan_sampling} Angstrom from the median " + "nearest-neighbour position spacing.", + stacklevel=3, + ) + scan_sampling = tuple(float(s) for s in scan_sampling) + + if normalization_order != 0: + raise ValueError( + "`normalization_order=1` fits a 2D linear background per bright-field image " + "and needs a scan grid, which an ungridded scan does not have; use " + "`normalization_order=0`." + ) + + shifted_tensor, rotation_angle = fit_and_shift_diffraction_origin( + dataset, + device=device, + max_batch_size=max_batch_size, + fit_method=fit_method, + mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + probe_positions=positions_ang, + ) + + if bf_mask is None: + bf_mask = bf_mask_from_mean_pattern(shifted_tensor, intensity_threshold) + else: + bf_mask = validate_tensor(bf_mask, "bf_mask", dtype=torch.bool).to(shifted_tensor.device) + if tuple(bf_mask.shape) != tuple(shifted_tensor.shape[-2:]): + raise ValueError( + f"`bf_mask` has shape {tuple(bf_mask.shape)} but the detector is " + f"{tuple(shifted_tensor.shape[-2:])}." + ) + bf_mask_dataset = Dataset2d.from_array( + bf_mask.cpu().numpy(), + name="BF mask", + units=dataset.units[-2:], + sampling=dataset.sampling[-2:], + ) + + vbf_stack = shifted_tensor[..., bf_mask].cpu() # (N, N_bf) + vbf_stack = normalize_vbf_stack(vbf_stack, normalization_order, vbf_stack.shape[:1]) + vbf_stack = vbf_stack.T.contiguous() # (N_bf, N) + + # anchor to the position bounding box, then convert to canvas pixels + scan_origin = positions_ang.min(axis=0) + positions_px = (positions_ang - scan_origin) / np.asarray(scan_sampling) + scan_gpts = tuple(int(math.ceil(v)) + 1 for v in positions_px.max(axis=0)) + + return ( + vbf_stack, + positions_px, + bf_mask_dataset, + scan_gpts, + scan_sampling, + rotation_angle, + tuple(float(v) for v in scan_origin), + ) + + +def regrid_vbf_stack( + vbf_stack, + positions_px, + scan_gpts, + interpolation: Literal["nearest", "bilinear"] = "nearest", + hole_fill: Literal["mean", "zero"] = "mean", +): + """ + Resample a flat ``(N_bf, N)`` vBF stack onto a regular scan grid. + + Splats each bright-field image at its probe positions with no aberration shift and + divides by the accumulated weight, which lets the scan-Fourier formulation accept an + ungridded acquisition. + + A grid finer than the scan is a supported case: the empty pixels between the probe + positions are then the sparse comb that ``upsampling_factor`` would build by Fourier + tiling, except that the teeth land on the real probe positions rather than a lattice, + and the deconvolution unfolds them from the bright-field shifts. + + Parameters + ---------- + interpolation : {"nearest", "bilinear"} + Deposition scheme, matching ``DirectPtychographyMontage.reconstruct``. ``"nearest"`` + by default: it keeps each measurement on one pixel, where bilinear smears it over + four and blurs away the sub-pixel detail a finer grid exists to capture (radial CTF + correlation 0.997 versus 0.986 on a scattered scan). + hole_fill : {"mean", "zero"} + What to put in grid pixels no probe position reached. + + ``"mean"`` (default) uses each bright-field image's mean over the visited pixels. + ``DirectPtychography._preprocess`` zeroes the DC bin, which subtracts the mean over + the whole grid, holes included -- so zero-filled holes sit at ``-mean``, a hard-edged + step the deconvolution then smears across the reconstruction. Filling with the + occupied mean puts them at the zero level and the step disappears. Measured on a + masked disk-shaped scan with 20% holes, correlation with ground truth goes from 0.25 + (zero-filled) to 0.69 (mean-filled), against 0.69 for the montage on the same data. + + The same choice serves a finer-than-scan grid: filling the comb's gaps with the + occupied mean and then zeroing the DC bin leaves them at zero, which is what the + unfolding needs. + + ``"zero"`` leaves them at zero without that centering, and inverts the + reconstruction. Kept for comparison. + + Returns + ------- + gridded : torch.Tensor + ``(N_bf, *scan_gpts)``. + hole_fraction : float + occupied : torch.Tensor + Boolean ``scan_gpts`` map of which pixels a probe position reached. + """ + if hole_fill not in ("mean", "zero"): + raise ValueError(f"`hole_fill` must be 'mean' or 'zero', got {hole_fill!r}") + + num_bf = int(vbf_stack.shape[0]) + device = vbf_stack.device + scan_gpts = (int(scan_gpts[0]), int(scan_gpts[1])) + coords = torch.as_tensor(positions_px, device=device, dtype=preferred_float_dtype(device))[ + None + ] + + # one image at a time: `scatter_add_splat` accumulates its whole batch into a single + # canvas, which is what the montage wants but would sum the stack away here + gridded = torch.empty((num_bf, *scan_gpts), device=device, dtype=torch.float32) + weights = None + for index in range(num_bf): + buffers = allocate_splat_buffers(scan_gpts, device, accumulate_squares=False) + sum_w, sum_wv, _ = scatter_add_splat( + vbf_stack[index : index + 1], + coords, + scan_gpts, + boundary="pad", + interpolation=interpolation, + out=buffers, + ) + if weights is None: + weights = sum_w + normalized = sum_wv / sum_w.clamp_min(torch.finfo(sum_w.dtype).tiny) + gridded[index] = normalized.reshape(scan_gpts).to(torch.float32) + + occupied = (weights > 0).reshape(scan_gpts) + hole_fraction = float((~occupied).sum()) / occupied.numel() + + if hole_fill == "mean" and bool(occupied.any()): + occupied_mean = gridded[:, occupied].mean(dim=1) + gridded[:, ~occupied] = occupied_mean[:, None] + + # Empty pixels only matter when there were enough positions to fill the grid and they + # still did not: on a deliberately finer grid the gaps are the point, and `nearest` + # deposition leaves ~30% empty even at the same size -- a configuration that measures + # *better* than a gapless bilinear one, so a low threshold would mislead. + positions_per_pixel = len(positions_px) / (scan_gpts[0] * scan_gpts[1]) + if positions_per_pixel >= 1.0 and hole_fraction > 0.5: + warnings.warn( + f"{hole_fraction:.1%} of the {scan_gpts[0]}x{scan_gpts[1]} scan grid received no " + f"probe position, despite {len(positions_px)} positions being available to cover " + f"{scan_gpts[0] * scan_gpts[1]} pixels, so the positions are clustered rather " + f"than merely sparse. Those pixels were filled with `hole_fill={hole_fill!r}`. " + "Use a coarser `scan_sampling`, or DirectPtychographyMontage, which needs no " + "grid at all.", + stacklevel=3, + ) + + return gridded, hole_fraction, occupied + + +def fit_and_shift_diffraction_origin( + dataset, + device: str | int = "cpu", + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin=None, + force_fitted_origin=None, + rotation_angle: float | None = None, + probe_positions=None, +): + """ + Measure, fit and remove the diffraction origin, returning a corner-centered stack. + + Works for both 4D ``(Rx, Ry, Qx, Qy)`` and 3D ``(N, Qx, Qy)`` datasets. For 3D input + ``probe_positions`` (``(N, 2)``) must be supplied for the background fit, and + ``rotation_angle`` must be given -- rotation estimation needs the 2D scan grid. + + Returns + ------- + shifted_tensor : torch.Tensor + Same shape as the input, with the diffraction origin moved to ``(0, 0)``. + rotation_angle : float + The supplied angle, or the estimated one when ``rotation_angle`` was ``None``. + """ + from quantem.diffractive_imaging.origin_models import CenterOfMassOriginModel + + origin = CenterOfMassOriginModel.from_dataset(dataset, device=device) + + # measure and fit origin + if force_fitted_origin is None: + if force_measured_origin is None: + origin.calculate_origin(max_batch_size) + else: + origin.origin_measured = force_measured_origin + if probe_positions is None: + origin.fit_origin_background(fit_method=fit_method) + else: + origin.fit_origin_background(probe_positions=probe_positions, fit_method=fit_method) + else: + origin.origin_fitted = force_fitted_origin + + if rotation_angle is None: + if dataset.ndim != 4: + raise ValueError( + "`rotation_angle` must be given for non-raster scans: detector rotation is " + "estimated from the curl of the center-of-mass over a 2D scan grid, which " + "requires 4D data." + ) + origin.estimate_detector_rotation() + rotation_angle = origin.detector_rotation_deg + + # shift to origin + origin.shift_origin_to( + max_batch_size=max_batch_size, + mode=mode, + ) + + return origin.shifted_tensor, rotation_angle + + +def estimate_frame_drift( + reconstructions, + upsample_factor: int = 16, + num_iterations: int = 3, + verbose: bool = True, +): + """ + Rigid drift of each frame of a multi-frame acquisition, in Angstrom. + + A long acquisition is often split into several interleaved frames -- successive passes + of a self-filling hexagonal grid, say -- so that specimen drift shows up as a shift + *between* frames rather than as a smear within one. Reconstruct each frame on its own, + pass the reconstructions here, and subtract the returned drift from the probe positions + before reconstructing them all together:: + + drift = estimate_frame_drift(montages) + combined_positions = np.concatenate([p - d for p, d in zip(positions, drift)]) + + Every reconstruction must cover the *same* window of the specimen, since the estimate + is a plain cross-correlation between them; pin it with ``reconstruct``'s ``obj_origin`` + and ``obj_fov``, which is checked here. + + Parameters + ---------- + reconstructions : sequence of DirectPtychographyBase + Reconstructed frames, in acquisition order. Each must have been reconstructed. + upsample_factor : int + Sub-pixel refinement of the correlation peak. + num_iterations : int + Leave-one-out refinement passes. Each frame is aligned against the mean of the + others, so no single frame is privileged as the reference; one pass is usually + enough, and the estimate converges within two or three. + verbose : bool + Report the drift per frame, and the largest update of the final pass. + + Returns + ------- + drift : ndarray + ``(n_frames, 2)`` drift in Angstrom, ordered ``(row, col)`` and referred to the mean + over frames, so it sums to zero rather than pinning frame 0. + + Notes + ----- + The drift is *rigid per frame*: it cannot represent drift accumulating within a frame, + which is what interleaving the frames is meant to avoid in the first place. For drift + that varies along the scan, see + :class:`~quantem.imaging.drift.DriftCorrection`, which warps individual scanlines. + """ + frames = list(reconstructions) + if len(frames) < 2: + raise ValueError("`estimate_frame_drift` needs at least two reconstructions.") + + images, samplings = [], [] + for i, frame in enumerate(frames): + obj = frame.obj + if obj is None: + raise ValueError(f"Frame {i} has not been reconstructed yet; call `.reconstruct()`.") + images.append(np.asarray(obj, dtype=np.float64)) + samplings.append(np.asarray(frame._obj_sampling, dtype=np.float64)) + + shapes = {img.shape for img in images} + if len(shapes) != 1: + raise ValueError( + f"Frames must share a canvas to be correlated, got shapes {sorted(shapes)}. " + "Reconstruct them with the same `obj_origin` and `obj_fov`." + ) + + # a canvas of the right shape in the wrong place is the subtler failure: the correlation + # would then measure the canvas offset rather than the drift, and silently succeed + origins = np.array([frame.obj_origin for frame in frames], dtype=np.float64) + sampling = samplings[0] + if not np.allclose(np.abs(origins - origins[0]).max(), 0.0, atol=1e-3 * sampling.min()): + raise ValueError( + "Frames share a canvas shape but not a canvas origin, so a correlation between " + f"them would measure that offset rather than the drift: {origins.tolist()}. " + "Reconstruct them with the same `obj_origin`." + ) + if not all(np.allclose(s, sampling) for s in samplings): + raise ValueError(f"Frames must share a sampling, got {[s.tolist() for s in samplings]}.") + + stack = torch.as_tensor(np.array(images), dtype=torch.float64) + spectra = torch.fft.fft2(stack) + kx = torch.fft.fftfreq(stack.shape[-2], dtype=torch.float64)[:, None] + ky = torch.fft.fftfreq(stack.shape[-1], dtype=torch.float64)[None, :] + + shifts = torch.zeros((len(frames), 2), dtype=torch.float64) + for _ in range(max(1, int(num_iterations))): + ramp = torch.exp( + -2j * np.pi * (kx * shifts[:, 0, None, None] + ky * shifts[:, 1, None, None]) + ) + aligned = torch.fft.ifft2(spectra * ramp).real + + total = aligned.sum(dim=0) + updates = torch.zeros_like(shifts) + for i in range(len(frames)): + # leave-one-out reference, so no frame is privileged as "the" reference + reference = (total - aligned[i]) / (len(frames) - 1) + updates[i] = cross_correlation_shift_torch( + reference, aligned[i], upsample_factor=upsample_factor + ) + shifts = shifts + updates + shifts = shifts - shifts.mean(dim=0, keepdim=True) + + # `cross_correlation_shift_torch` returns the shift that *undoes* the displacement, so + # the drift itself -- how far the frame moved -- is its negation + drift = -shifts.numpy() * sampling + + if verbose: + residual = float(updates.abs().max()) * float(sampling.max()) + print(f"Frame drift (Angstrom), final pass moved at most {residual:.2f} A:") + for i, d in enumerate(drift): + print(f" frame {i}: ({d[0]:+8.2f}, {d[1]:+8.2f})") + + return drift + + +def bf_mask_from_mean_pattern(shifted_tensor, intensity_threshold: float = 0.5): + """Bright-field mask from the mean diffraction pattern of a corner-centered stack.""" + scan_dims = tuple(range(shifted_tensor.ndim - 2)) + mean_dp = shifted_tensor.mean(dim=scan_dims) + return mean_dp > mean_dp.max() * intensity_threshold + + +def normalize_vbf_stack(vbf_stack, normalization_order: int, gpts: tuple[int, int]): + """ + Normalize a ``(*scan_gpts, N_bf)`` virtual bright-field stack. + + ``normalization_order=0`` scales each BF image to unity mean over the scan; + ``normalization_order=1`` divides out a least-squares linear background instead. + """ + if normalization_order == 0: + scan_dims = tuple(range(vbf_stack.ndim - 1)) + vbf_stack = vbf_stack / vbf_stack.mean(scan_dims) # unity mean, important + + elif normalization_order == 1: + # Fit linear background to each BF image + x = torch.linspace(-0.5, 0.5, gpts[0]) + y = torch.linspace(-0.5, 0.5, gpts[1]) + ya, xa = torch.meshgrid(y, x, indexing="ij") + + # Basis for linear fit: [1, x, y] + basis = torch.stack( + [torch.ones_like(xa.ravel()), xa.ravel(), ya.ravel()], dim=1 + ) # shape: [N_pixels, 3] + + # Fit each BF image + for k in range(vbf_stack.shape[-1]): + intensities = vbf_stack[..., k].ravel() + + # Least squares + coefs = torch.linalg.lstsq(basis, intensities).solution + + # Normalize + background = (basis @ coefs).reshape(gpts) + vbf_stack[..., k] /= background + else: + raise ValueError(f"`normalization_order` must be 0 or 1, got {normalization_order!r}") + + return vbf_stack + + +def build_vbf_stack_from_dataset4d( + dataset, + device: str | int = "cpu", + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin=None, + force_fitted_origin=None, + rotation_angle: float | None = None, + intensity_threshold: float = 0.5, + normalization_order: int = 0, + edge_blend_pixels: int = 0, +): + """ + Turn a 4D-STEM dataset into the virtual bright-field stack the direct-ptychography + classes consume. + + Returns + ------- + vbf_dataset : Dataset3d + ``(N_bf, Rx, Ry)`` stack of virtual bright-field images. + bf_mask_dataset : Dataset2d + Corner-centered bright-field mask on the detector grid. + rotation_angle : float + The supplied angle, or the estimated one when ``rotation_angle`` was ``None``. + """ + from quantem.core.datastructures import Dataset2d, Dataset3d + + shifted_tensor, rotation_angle = fit_and_shift_diffraction_origin( + dataset, + device=device, + max_batch_size=max_batch_size, + fit_method=fit_method, + mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + ) + + bf_mask = bf_mask_from_mean_pattern(shifted_tensor, intensity_threshold) + bf_mask_dataset = Dataset2d.from_array( + bf_mask.cpu().numpy(), + name="BF mask", + units=dataset.units[-2:], + sampling=dataset.sampling[-2:], + ) + + # vbf_stack + vbf_stack = shifted_tensor[..., bf_mask].cpu() + gpts = vbf_stack.shape[:2] + vbf_stack = normalize_vbf_stack(vbf_stack, normalization_order, gpts) + + # smooth window + window_edge = create_edge_window(shape=gpts, edge_blend_pixels=edge_blend_pixels, device="cpu") + vbf_stack = (1 - window_edge[..., None]) + window_edge[..., None] * vbf_stack + + vbf_stack = torch.moveaxis(vbf_stack, (0, 1, 2), (1, 2, 0)) + vbf_dataset = Dataset3d.from_array( + vbf_stack.numpy(), + name="vBF stack", + units=("index",) + tuple(dataset.units[:2]), + sampling=(1,) + tuple(dataset.sampling[:2]), + ) + + return vbf_dataset, bf_mask_dataset, rotation_angle diff --git a/src/quantem/diffractive_imaging/direct_ptychography.py b/src/quantem/diffractive_imaging/direct_ptychography.py index 912b4bd2..c5fe2f5f 100644 --- a/src/quantem/diffractive_imaging/direct_ptychography.py +++ b/src/quantem/diffractive_imaging/direct_ptychography.py @@ -1,37 +1,29 @@ import gc import math -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Dict, Literal, Tuple +import warnings +from typing import TYPE_CHECKING, Literal, Tuple import numpy as np -import optuna from numpy.typing import NDArray from tqdm.auto import tqdm from quantem.core import config from quantem.core.datastructures import Dataset2d, Dataset3d, Dataset4d -from quantem.core.io.serialize import AutoSerialize -from quantem.core.utils.rng import RNGMixin -from quantem.core.utils.utils import electron_wavelength_angstrom, to_numpy from quantem.core.utils.validators import ( validate_aberration_coefficients, - validate_gt, - validate_int, validate_tensor, ) -from quantem.core.visualization import show_2d from quantem.diffractive_imaging.complex_probe import ( + FourierProbe, aberration_surface, aberration_surface_cartesian_basis, aberration_surface_cartesian_gradients, aperture, - evaluate_probe, gamma_factor, merge_aberration_coefficients, polar_coordinates, spatial_frequencies, ) -from quantem.diffractive_imaging.origin_models import CenterOfMassOriginModel from quantem.diffractive_imaging.ptycho_utils import SimpleBatcher if TYPE_CHECKING: @@ -40,175 +32,53 @@ if config.get("has_torch"): import torch -from itertools import product - 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, + build_vbf_stack_from_dataset3d, + build_vbf_stack_from_dataset4d, fit_aberrations_from_shifts, group_basis_by_method, + regrid_vbf_stack, unwrap_bf_overlap_phase_torch, ) +from quantem.diffractive_imaging.direct_ptychography_base import ( + DirectPtychographyBase, + HyperparameterState, +) -optuna.logging.set_verbosity(optuna.logging.WARNING) - - -@dataclass -class OptimizationParameter: - low: float - high: float - log: bool = False - n_points: int | None = None - - def grid_values(self): - """Return an array of grid values for this parameter.""" - if self.n_points is None: - raise ValueError("n_points must be specified for grid search parameters.") - if self.log: - return np.geomspace(self.low, self.high, self.n_points) - else: - return np.linspace(self.low, self.high, self.n_points) - - -@dataclass -class HyperparameterState: - initial_aberrations: dict[str, float] = field(default_factory=dict) - initial_rotation_angle: float | None = None - optimized_aberrations: dict[str, float] = field(default_factory=dict) - optimized_rotation_angle: float | None = None - optimized_keys: set[str] = field(default_factory=set) - study: optuna.Study | None = None - - def __post_init__(self): - self.initial_aberrations = validate_aberration_coefficients(dict(self.initial_aberrations)) - self.optimized_aberrations = validate_aberration_coefficients(self.optimized_aberrations) - - if self.optimized_keys: - canonical = validate_aberration_coefficients( - {k: 0.0 for k in self.optimized_keys if k != "rotation_angle"} - ) - self.optimized_keys = set(canonical.keys()) | { - k for k in self.optimized_keys if k == "rotation_angle" - } - - def __repr__(self) -> str: - return self.summarize(which="all") - - def current_aberrations( - self, override_fixed: dict[str, float] | None = None - ) -> Dict[str, float]: - """Return full aberration dictionary (fixed ⊕ optimized).""" - out = dict(self.initial_aberrations) - out.update(self.optimized_aberrations) - if override_fixed is not None: - out.update(validate_aberration_coefficients(override_fixed)) - return out - - def current_rotation_angle(self, override_fixed: float | None = None) -> float: - """Return rotation angle (optimized takes precedence).""" - if override_fixed is not None: - return override_fixed - if self.optimized_rotation_angle is not None: - return self.optimized_rotation_angle - if self.initial_rotation_angle is not None: - return self.initial_rotation_angle - return 0.0 - - def clear_optimized(self): - """Clear all optimized aberrations and rotation angle.""" - self.optimized_aberrations.clear() - self.optimized_rotation_angle = None - self.optimized_keys.clear() - self.study = None - - def clear_all(self): - """Clear everything: initial and optimized hyperparameters.""" - self.initial_aberrations.clear() - self.initial_rotation_angle = None - self.clear_optimized() - - def copy(self): - """ """ - return HyperparameterState( - initial_aberrations=self.initial_aberrations, - optimized_aberrations=self.optimized_aberrations, - initial_rotation_angle=self.initial_rotation_angle, - optimized_rotation_angle=self.optimized_rotation_angle, - optimized_keys=self.optimized_keys, - study=self.study, - ) - - def summarize( - self, - *, - which: str = "current", - override_aberration_coefs: dict[str, float] | None = None, - override_rotation_angle: float | None = None, - ) -> str: - cls = self.__class__.__name__ - lines: list[str] = [] - - def add(name: str, value): - lines.append(f" {name}={value!r},") - - if which == "initial": - if self.initial_aberrations: - add("initial_aberrations", self.initial_aberrations) - if self.initial_rotation_angle is not None: - 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_deg", self.optimized_rotation_angle) - - elif which == "current": - current_abers = self.current_aberrations(override_aberration_coefs) - current_rot = self.current_rotation_angle(override_rotation_angle) - - if current_abers: - add("current_aberrations", current_abers) - if current_rot is not None: - 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_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_deg", self.optimized_rotation_angle) - - else: - raise ValueError( - f"`which` must be one of " - f'{{"initial", "optimized", "current", "all"}}, got {which!r}' - ) +# re-exported because this was `OptimizationParameter`'s public path in v0.1.9, and +# published notebooks still import it from this module +from quantem.diffractive_imaging.direct_ptychography_base import ( + OptimizationParameter as OptimizationParameter, +) - if not lines: - return f"{cls}()" - body = "\n".join(lines) - return f"{cls}(\n{body}\n)" +class DirectPtychography(DirectPtychographyBase): + """ + Direct ptychography in the scan-space Fourier domain. + Every kernel here -- SSB, OBF, matched-filter, iCoM and parallax -- is a multiplier on + ``G(k, q)``, the virtual bright-field stack Fourier transformed over the scan. The + object is recovered by summing those multipliers over bright-field pixels and + transforming back once. -@dataclass(frozen=True) -class BrightFieldContext: - bf_mask: torch.Tensor - bf_inds_i: torch.Tensor - bf_inds_j: torch.Tensor - num_bf: int - vbf_index_mapping: torch.Tensor + Because the multiplication happens in ``q``, the scan has to lie on a regular grid. + :meth:`from_dataset4d` and :meth:`from_virtual_bfs` are given one; :meth:`from_dataset3d` + resamples ungridded positions onto one, which is where its caveats come from. + The alternative is + :class:`~quantem.diffractive_imaging.direct_ptychography_montage.DirectPtychographyMontage`, + which accumulates the scan onto a real-space canvas and needs no grid. See its class + docstring for the full comparison; in short, this class is exact and cheapest on a + gridded scan, but a bright-field mask beyond a few tens of thousands of pixels will not + fit in memory here. -class DirectPtychography(RNGMixin, AutoSerialize): - """ """ + Instantiate with :meth:`from_dataset4d`, :meth:`from_virtual_bfs` or + :meth:`from_dataset3d`. + """ _token = object() @@ -216,16 +86,18 @@ def __init__( self, vbf_dataset: Dataset3d, bf_mask_dataset: Dataset2d, - energy: float, + energy: float | None, rotation_angle: float, aberration_coefs: dict, - semiangle_cutoff: float | None, + semiangle_cutoff: float, soft_edges: bool, crop_bf_mask: bool, bf_mask_padding_px: int, rng: np.random.Generator | int | None, device: str | int, verbose: int | bool, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, _token: object | None = None, ): """ """ @@ -241,19 +113,20 @@ def __init__( if crop_bf_mask: self.bf_mask = _crop_corner_centered_mask(self.bf_mask, bf_mask_padding_px) - self.wavelength = electron_wavelength_angstrom(energy) + self.wavelength = self._resolve_wavelength(energy, wavelength) self.scan_units = vbf_dataset.units[-2:] self.detector_units = bf_mask_dataset.units - self.scan_gpts = vbf_dataset.shape[-2:] + self.scan_gpts = tuple(int(n) for n in vbf_dataset.shape[-2:]) self.scan_sampling = vbf_dataset.sampling[-2:] self.reciprocal_sampling = bf_mask_dataset.sampling self.angular_sampling = tuple(d * 1e3 * self.wavelength for d in self.reciprocal_sampling) self.num_bf = vbf_dataset.shape[0] - self.gpts = self.bf_mask.shape[:2] + self.gpts = tuple(int(n) for n in self.bf_mask.shape[:2]) self.sampling = tuple(1 / s / n for n, s in zip(self.reciprocal_sampling, self.gpts)) + self.fourier_probe = fourier_probe self.semiangle_cutoff = semiangle_cutoff # ty:ignore[invalid-assignment] self.soft_edges = soft_edges self.rng = rng @@ -269,10 +142,12 @@ def from_virtual_bfs( cls, vbf_dataset: Dataset3d, bf_mask_dataset: Dataset2d, - energy: float, - rotation_angle: float, - aberration_coefs: dict = {}, + energy: float | None = None, + rotation_angle: float | None = None, semiangle_cutoff: float | None = None, + aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, soft_edges: bool = True, crop_bf_mask: bool = True, bf_mask_padding_px: int = 1, @@ -286,9 +161,11 @@ def from_virtual_bfs( vbf_dataset=vbf_dataset, bf_mask_dataset=bf_mask_dataset, energy=energy, + wavelength=wavelength, rotation_angle=rotation_angle, aberration_coefs=aberration_coefs, semiangle_cutoff=semiangle_cutoff, + fourier_probe=fourier_probe, soft_edges=soft_edges, crop_bf_mask=crop_bf_mask, bf_mask_padding_px=bf_mask_padding_px, @@ -302,9 +179,11 @@ def from_virtual_bfs( def from_dataset4d( cls, dataset: Dataset4d, - energy: float, - semiangle_cutoff: float, + energy: float | None = None, + semiangle_cutoff: float | None = None, aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, rotation_angle: float | None = None, max_batch_size: int | None = None, fit_method: str = "plane", @@ -323,99 +202,271 @@ def from_dataset4d( ): """ """ - origin = CenterOfMassOriginModel.from_dataset(dataset, device=device) - - # measure and fit origin - if force_fitted_origin is None: - if force_measured_origin is None: - origin.calculate_origin(max_batch_size) - else: - origin.origin_measured = force_measured_origin # ty:ignore[invalid-assignment] - origin.fit_origin_background(fit_method=fit_method) - else: - origin.origin_fitted = force_fitted_origin # ty:ignore[invalid-assignment] - - if rotation_angle is None: - origin.estimate_detector_rotation() - rotation_angle = origin.detector_rotation_deg - - # shift to origin - origin.shift_origin_to( + vbf_dataset, bf_mask_dataset, rotation_angle = build_vbf_stack_from_dataset4d( + dataset, + device=device, max_batch_size=max_batch_size, + fit_method=fit_method, mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + intensity_threshold=intensity_threshold, + normalization_order=normalization_order, + edge_blend_pixels=edge_blend_pixels, ) - shifted_tensor = origin.shifted_tensor - # bf_mask - mean_dp = shifted_tensor.mean(dim=(0, 1)) - bf_mask = mean_dp > mean_dp.max() * intensity_threshold - - bf_mask_dataset = Dataset2d.from_array( - bf_mask.cpu().numpy(), - name="BF mask", - units=dataset.units[-2:], - sampling=dataset.sampling[-2:], + return cls( + vbf_dataset=vbf_dataset, + bf_mask_dataset=bf_mask_dataset, + energy=energy, + wavelength=wavelength, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + semiangle_cutoff=semiangle_cutoff, + fourier_probe=fourier_probe, + soft_edges=soft_edges, + crop_bf_mask=crop_bf_mask, + bf_mask_padding_px=bf_mask_padding_px, + rng=rng, + device=device, + verbose=verbose, + _token=cls._token, ) - # vbf_stack - vbf_stack = shifted_tensor[..., bf_mask].cpu() - gpts = vbf_stack.shape[:2] - - if normalization_order == 0: - vbf_stack = vbf_stack / vbf_stack.mean((0, 1)) # unity mean, important - - elif normalization_order == 1: - # Fit linear background to each BF image - x = torch.linspace(-0.5, 0.5, gpts[0]) - y = torch.linspace(-0.5, 0.5, gpts[1]) - ya, xa = torch.meshgrid(y, x, indexing="ij") - - # Basis for linear fit: [1, x, y] - basis = torch.stack( - [torch.ones_like(xa.ravel()), xa.ravel(), ya.ravel()], dim=1 - ) # shape: [N_pixels, 3] + @classmethod + def from_dataset3d( + cls, + dataset: Dataset3d, + positions: Dataset2d | torch.Tensor | NDArray, + energy: float | None = None, + semiangle_cutoff: float | None = None, + rotation_angle: float | None = None, + scan_sampling: Tuple[float, float] | Literal["auto"] = "auto", + aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, + bf_mask: torch.Tensor | NDArray | None = None, + scan_gpts: Tuple[int, int] | None = None, + interpolation: Literal["nearest", "bilinear"] = "nearest", + hole_fill: Literal["mean", "zero"] = "mean", + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + force_fitted_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + intensity_threshold: float = 0.5, + soft_edges: bool = True, + crop_bf_mask: bool = True, + bf_mask_padding_px: int = 1, + rng: np.random.Generator | int | None = None, + device: str | int = "cpu", + verbose: int | bool = True, + normalization_order: int = 0, + ): + """ + Build from an ungridded stack of diffraction patterns and their probe positions. - # Fit each BF image - for k in range(vbf_stack.shape[-1]): - intensities = vbf_stack[..., k].ravel() + Every kernel here is a multiplier on the scan-space Fourier transform, which needs a + regular grid. So the bright-field images are resampled onto one first, by splatting + them at their probe positions with no aberration shift and dividing by the + accumulated weight. Once gridded, SSB, OBF, matched-filter and iCoM all run exactly + as they do for a raster scan. - # Least squares - coefs = torch.linalg.lstsq(basis, intensities).solution + Parameters + ---------- + dataset : Dataset3d + ``(N, Qx, Qy)`` diffraction patterns, reciprocal units ``"A^-1"`` or ``"mrad"``. + positions : Dataset2d, torch.Tensor or ndarray + ``(N, 2)`` probe positions in Angstrom, ordered ``(row, col)`` to match the + diffraction axes. A ``Dataset2d`` must carry units ``"A"``. + rotation_angle : float + Detector rotation in degrees. Required: rotation is otherwise estimated from the + curl of the center of mass over a 2D scan grid, which an ungridded scan lacks. + scan_sampling : tuple of float or "auto" + Grid pixel size in Angstrom. ``"auto"`` uses the median nearest-neighbour + position spacing and warns with the inferred value. + + This, rather than ``reconstruct(upsampling_factor=...)``, is how to sample more + finely than the scan. Pass a fraction of the position spacing and the empty + pixels in between become the sparse comb the deconvolution unfolds, with its + teeth on the real probe positions. Measured against the analytical parallax CTF + on a scan scattered by a full pixel, correlation over the band above the scan + Nyquist: 0.720 for binning then ``upsampling_factor=2``, against 0.997 for + halving ``scan_sampling`` instead, and 0.982 for the montage. + scan_gpts : tuple of int, optional + Grid size. Defaults to whatever just covers the positions at ``scan_sampling``, + and may only be larger than that -- ``scan_sampling`` stays authoritative, so + this pads the canvas rather than rescaling the positions into it. + interpolation : {"nearest", "bilinear"} + Deposition scheme, matching ``DirectPtychographyMontage.reconstruct``. + ``"nearest"`` by default; see :func:`regrid_vbf_stack`. + hole_fill : {"mean", "zero"} + What to put in grid pixels no probe position reached. Defaults to each + bright-field image's mean over the visited pixels, which matters on a masked or + irregular scan and is also what centers a comb's gaps -- see + :func:`regrid_vbf_stack`. + + Notes + ----- + Regridding has two failure modes, both of which raise a warning. + + Avoid combining this with ``reconstruct(upsampling_factor > 1)`` on an irregular + scan; pass a finer ``scan_sampling`` instead. Upsampling at reconstruct time recovers + detail above the scan Nyquist from where each probe sat, which binning onto a grid + discards, so the extra band returns as a replica of the contrast-transfer function + rather than an extension of it. With no holes, correlation over the extension band + falls 0.997 -> 0.822 -> 0.711 as the sub-pixel scatter grows from 0 to 0.5 to 1.0 + grid pixels. + + The second is holes: a grid pixel no probe reached has to be invented. ``hole_fill`` + keeps that from becoming a hard-edged step, but a large filled region still biases + the result, and a coarser ``scan_gpts`` is the usual fix. Contiguous holes are worse + than scattered ones at equal fraction -- scattered holes behave like noise, while an + excluded region is a low-frequency mask the deconvolution spreads everywhere. Note + that a masked scan and an upsampled one want opposite treatments, filled holes versus + deliberate gaps, and ``hole_fill`` cannot do both at once. + + Both are properties of the regridding rather than the deconvolution, so neither + applies to + :class:`~quantem.diffractive_imaging.direct_ptychography_montage.DirectPtychographyMontage`, + which uses the positions as measured. Prefer it for a scan that is sparse, strongly + irregular, or both masked and upsampled; see its class docstring for the comparison. + """ + ( + vbf_stack, + positions_px, + bf_mask_dataset, + fitted_gpts, + scan_sampling, + rotation_angle, + scan_origin, + ) = build_vbf_stack_from_dataset3d( + dataset, + positions, + scan_sampling, + device=device, + max_batch_size=max_batch_size, + fit_method=fit_method, + mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + intensity_threshold=intensity_threshold, + normalization_order=normalization_order, + bf_mask=bf_mask, + ) - # Normalize - background = (basis @ coefs).reshape(gpts) - vbf_stack[..., k] /= background + if scan_gpts is None: + scan_gpts = fitted_gpts else: - raise ValueError() + # `scan_sampling` stays authoritative, so the grid is padded, never rescaled: + # resizing the pixel to fill a requested shape would contradict the caller + scan_gpts = tuple(int(n) for n in scan_gpts) + if any(n < f for n, f in zip(scan_gpts, fitted_gpts)): + raise ValueError( + f"`scan_gpts={scan_gpts}` is smaller than the {fitted_gpts} needed to " + f"cover the positions at scan_sampling={scan_sampling}; positions would " + "be dropped. Pass a coarser `scan_sampling` instead." + ) - # smooth window - window_edge = create_edge_window( - shape=gpts, edge_blend_pixels=edge_blend_pixels, device="cpu" + gridded, hole_fraction, occupied = regrid_vbf_stack( + vbf_stack, + positions_px, + scan_gpts, + interpolation=interpolation, + hole_fill=hole_fill, ) - vbf_stack = (1 - window_edge[..., None]) + window_edge[..., None] * vbf_stack + if verbose: + occupancy = vbf_stack.shape[1] / (scan_gpts[0] * scan_gpts[1]) + detail = ( + f"grid is finer than the scan ({occupancy:.2f} positions per pixel), so " + f"{hole_fraction:.1%} of it is comb gaps for the deconvolution to unfold" + if occupancy < 1 + else f"{hole_fraction:.1%} of grid pixels unvisited" + ) + print( + f"Regridded {vbf_stack.shape[1]} positions onto " + f"{scan_gpts[0]}x{scan_gpts[1]}; {detail}." + ) - vbf_stack = torch.moveaxis(vbf_stack, (0, 1, 2), (1, 2, 0)) vbf_dataset = Dataset3d.from_array( - vbf_stack.numpy(), - name="vBF stack", - units=("index",) + tuple(dataset.units[:2]), - sampling=(1,) + tuple(dataset.sampling[:2]), + gridded.cpu().numpy(), + name="regridded virtual BF stack", + sampling=(1.0, *scan_sampling), + units=("index", "A", "A"), ) - return cls( + reconstruction = cls.from_virtual_bfs( vbf_dataset=vbf_dataset, bf_mask_dataset=bf_mask_dataset, energy=energy, + wavelength=wavelength, rotation_angle=rotation_angle, - aberration_coefs=aberration_coefs, semiangle_cutoff=semiangle_cutoff, + aberration_coefs=aberration_coefs, + fourier_probe=fourier_probe, soft_edges=soft_edges, crop_bf_mask=crop_bf_mask, bf_mask_padding_px=bf_mask_padding_px, rng=rng, device=device, verbose=verbose, - _token=cls._token, + ) + # the grid was anchored at the position bounding box, so record where that was -- + # this is what lets `obj_origin` be read in the caller's own coordinates + reconstruction.scan_origin = scan_origin + + # how far the positions sit from the grid they were binned onto. Upsampling unfolds + # aliased detail from where each probe actually was, which regridding discards, so + # this predicts whether `upsampling_factor` can work at all + subpixel = positions_px - np.round(positions_px) + lattice_rms_px = float(np.sqrt((subpixel**2).sum(axis=1).mean())) + + reconstruction._regrid_info = { + "hole_fraction": hole_fraction, + "occupied": occupied, + "positions_px": positions_px, + "scan_gpts": scan_gpts, + "scan_sampling": scan_sampling, + "lattice_rms_px": lattice_rms_px, + } + return reconstruction + + #: sub-pixel scatter, in grid pixels, above which regridding has destroyed enough of the + #: probe positions that upsampling replicates the band instead of extending it + _UNFOLDING_RMS_LIMIT = 0.1 + + def _warn_if_upsampling_cannot_unfold(self, upsampling_factor): + """Warn when upsampling an irregular regridded scan, which cannot unfold. + + ``upsampling_factor`` recovers detail above the scan Nyquist from the aliased content + of the bright-field images, which depends on where each probe actually sat. Binning + those measurements onto a grid throws that away, so on an irregular scan the extra + band comes back as a replica of the CTF rather than an extension of it. + + Measured on a white-noise object against the analytical parallax CTF, with no holes + at all: correlation over the band above the scan Nyquist falls 0.997 -> 0.965 -> + 0.822 -> 0.711 as the sub-pixel scatter grows 0 -> 0.25 -> 0.5 -> 1.0 grid pixels, + while the montage holds 0.995 -> 0.983 on the same data. + """ + info = getattr(self, "_regrid_info", None) + if upsampling_factor <= 1 or info is None: + return + if info["lattice_rms_px"] <= self._UNFOLDING_RMS_LIMIT: + return + + warnings.warn( + f"This reconstruction was regridded from ungridded positions that sit " + f"{info['lattice_rms_px']:.2f} grid pixels from the grid on average, and " + f"`upsampling_factor={upsampling_factor}` cannot unfold that. Upsampling recovers " + "detail above the scan Nyquist from where each probe actually was, which the " + "regridding has already discarded, so the extra band will come back as a replica " + "of the contrast-transfer function rather than an extension of it. Rebuild with " + f"a `scan_sampling` {upsampling_factor}x finer instead, which keeps the probe " + "positions on the grid rather than binning them away first, or use " + "DirectPtychographyMontage, which needs no grid at all.", + stacklevel=2, ) def _preprocess( @@ -431,24 +482,6 @@ def _preprocess( return self - def _return_bf_context(self, bf_mask): - """ - Given a BF mask, compute all BF-dependent geometry and indexing. - """ - bf_mask = torch.as_tensor(bf_mask, dtype=torch.bool, device=self.device) - - bf_inds_i, bf_inds_j = torch.nonzero(bf_mask, as_tuple=True) - vbf_index_mapping = torch.where(bf_mask[self.bf_mask])[0] - num_bf = bf_inds_i.numel() - - return BrightFieldContext( - bf_mask=bf_mask, - bf_inds_i=bf_inds_i, - bf_inds_j=bf_inds_j, - num_bf=num_bf, - vbf_index_mapping=vbf_index_mapping, - ) - def _return_upsampled_qgrid( self, upsampling_factor=None, @@ -468,14 +501,6 @@ def _return_upsampled_qgrid( return qxa, qya - @property - def verbose(self) -> int: - return self._verbose - - @verbose.setter - def verbose(self, v: bool | int | float) -> None: - self._verbose = validate_int(validate_gt(v, -1, "verbose"), "verbose") - @property def vbf_stack(self) -> torch.Tensor: return self._vbf_stack @@ -486,78 +511,6 @@ def vbf_stack(self, value: torch.Tensor): device=self.device ) - @property - def bf_mask(self) -> torch.Tensor: - return self._bf_mask - - @bf_mask.setter - def bf_mask(self, value: torch.Tensor): - self._bf_mask = validate_tensor(value, "bf_mask", dtype=torch.bool).to(device=self.device) - - @property - def rotation_angle(self) -> float: - """Current detector rotation angle in degrees.""" - return self.hyperparameter_state.current_rotation_angle() - - @property - def aberration_coefs(self) -> dict: - return self.hyperparameter_state.current_aberrations() - - @property - def semiangle_cutoff(self) -> float: - return self._semiangle_cutoff - - @semiangle_cutoff.setter - def semiangle_cutoff(self, value: float): - validate_gt(value, 0.0, "semiangle_cutoff") - self._semiangle_cutoff = value - - @property - def device(self) -> str | torch.device: - """This should be of form 'cuda:X' or 'cpu', as defined by quantem.config""" - if hasattr(self, "_device"): - return self._device # ty:ignore[invalid-return-type] - else: - return config.get("device") - - @device.setter - def device(self, device: str | int | None): - if device is not None: - dev, _id = config.validate_device(device) - self._device = dev - - @property - def scan_sampling(self) -> NDArray: - return self._scan_sampling # ty:ignore[invalid-return-type] - - @scan_sampling.setter - def scan_sampling(self, value: NDArray | tuple | list) -> None: - """ - Units A or raises error - """ - units = self.scan_units - if units[0] == "A": - self._scan_sampling = value - else: - raise ValueError("real-space needs to be given in 'A'") - - @property - def reciprocal_sampling(self) -> NDArray: - return self._reciprocal_sampling # ty:ignore[invalid-return-type] - - @reciprocal_sampling.setter - def reciprocal_sampling(self, value: NDArray | tuple | list) -> None: - """ - Units A or raises error - """ - units = self.detector_units - if units[0] == "A^-1": - self._reciprocal_sampling = value - elif units[0] == "mrad": - self._reciprocal_sampling = tuple(val / self.wavelength / 1e3 for val in value) - else: - raise ValueError("reciprocal-space needs to be given in 'A^-1' or 'mrad'") - def _return_kernel_contributions( self, bf, @@ -594,11 +547,7 @@ def _return_kernel_contributions( (qmkxa, qmkya), (qpkxa, qpkya), cmplx_probe_at_k, - self.wavelength, - self.semiangle_cutoff, - self.soft_edges, - angular_sampling=self.angular_sampling, - aberration_coefs=aberration_coefs, + self._return_probe(aberration_coefs), normalize=False, ) @@ -628,31 +577,6 @@ def _return_kernel_contributions( return fourier_factor, power - def _normalize_kernel_name(self, kernel): - kernel = kernel.lower() - - aliases = { - "ssb": "ssb", - "single-sideband": "ssb", - "acbf": "ssb", - "aberration-corrected-bright-field": "ssb", - "obf": "obf", - "optimum-bright-field": "obf", - "mf": "mf", - "matched-filter": "mf", - "prlx": "prlx", - "parallax": "prlx", - "tcbf": "prlx", - "tilt-corrected-bright-field": "prlx", - "icom": "icom", - "center-of-mass": "icom", - } - - if kernel not in aliases: - raise ValueError(f"Unknown deconvolution kernel '{kernel}'") - - return aliases[kernel] - def reconstruct( self, bf_mask=None, @@ -725,6 +649,7 @@ def reconstruct( if upsampling_factor is None: upsampling_factor = 1 upsampling_factor = math.ceil(upsampling_factor) + self._warn_if_upsampling_cannot_unfold(upsampling_factor) if bf_mask is None: bf_mask = self.bf_mask @@ -738,6 +663,9 @@ def reconstruct( max_batch_size = num_bf deconvolution_kernel = self._normalize_kernel_name(deconvolution_kernel) + if deconvolution_kernel == "prlx": + # iCoM is exempt: `k . q / |q|**2` never reads the probe + self._require_analytic_probe("The prlx kernel") # Get upsampled q-space grid qxa, qya = self._return_upsampled_qgrid(upsampling_factor) @@ -776,14 +704,7 @@ def reconstruct( sign_sin_chi_q = None # compute global / cheap functions for all - cmplx_probe_k = evaluate_probe( - k * self.wavelength, - phi, - self.semiangle_cutoff, - self.angular_sampling, - self.wavelength, - aberration_coefs=aberration_coefs, - ) + cmplx_probe_k = self._return_probe_on_grid(k, phi, aberration_coefs) BF_weights = cmplx_probe_k[bf_mask].abs().square().sum() butterworth_env = torch.ones_like(q) @@ -901,290 +822,6 @@ def variance_loss(self): ) return variance_loss - @property - def obj(self) -> np.ndarray: - obj = to_numpy(self.corrected_bf) - return obj - - def visualize( - self, - return_fig: bool = False, - show_obj_fft: bool = True, - apply_hanning_window: 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 - 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 - return None - - def optimize_hyperparameters( - self, - aberration_coefs: dict[str, float | OptimizationParameter] | None = None, - rotation_angle: float | OptimizationParameter | None = None, - n_trials=50, - sampler=None, - verbose=None, - **reconstruct_kwargs, - ): - """ - Optimize hyperparameters (aberrations and/or rotation) using Optuna. - - Parameters - ---------- - 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, in degrees. - n_trials : int - Number of Optuna trials. - sampler : optuna.samplers.BaseSampler, optional - Custom Optuna sampler. - direction : str - "minimize" or "maximize" (default: "minimize"). - show_progress_bar : bool - Show progress bar during optimization. - add_fixed_to_hyperparameter_state: bool - fixed aberrations included will be passed on to hyperparameter state - **reconstruct_kwargs : - Extra arguments passed to reconstruct(). - """ - - if verbose is None: - verbose = self.verbose - - sampler = sampler or optuna.samplers.TPESampler() - - state = self.hyperparameter_state - aberration_coefs = aberration_coefs or {} - - # Reset optimized bookkeeping - state.clear_optimized() - - # Partition inputs - fixed_override_aberrations = {} - optimizable_aberrations = {} - - for name, val in aberration_coefs.items(): - if isinstance(val, OptimizationParameter): - optimizable_aberrations[name] = val - state.optimized_keys.add(name) - else: - fixed_override_aberrations[name] = val - - if isinstance(rotation_angle, OptimizationParameter): - state.optimized_keys.add("rotation_angle") - - def objective(trial): - trial_aberrations = {} - for name, val in optimizable_aberrations.items(): - trial_aberrations[name] = trial.suggest_float(name, val.low, val.high, log=val.log) - - trial_aberrations |= fixed_override_aberrations - - if isinstance(rotation_angle, OptimizationParameter): - rot = trial.suggest_float( - "rotation_angle", - rotation_angle.low, - rotation_angle.high, - log=rotation_angle.log, - ) - else: - rot = rotation_angle - - self.reconstruct( - override_aberration_coefs=trial_aberrations, - override_rotation_angle=rot, - verbose=False, - **reconstruct_kwargs, - ) - return float(self.variance_loss()) - - study = optuna.create_study(direction="minimize", sampler=sampler) - study.optimize(objective, n_trials=n_trials, show_progress_bar=verbose) - - # Write back optimized results - best = study.best_params.copy() - state.optimized_rotation_angle = best.pop("rotation_angle", None) - - if state.optimized_rotation_angle is None and rotation_angle is not None: - state.optimized_rotation_angle = rotation_angle # ty:ignore[invalid-assignment] - - state.optimized_aberrations = best - state.optimized_aberrations = state.current_aberrations(fixed_override_aberrations) - state.study = study - - if verbose: - print("Optimized state:\n\n", self.hyperparameter_state) - - self.reconstruct(verbose=False, **reconstruct_kwargs) - return self - - def grid_search_hyperparameters( - self, - aberration_coefs: dict[str, float | OptimizationParameter] | None = None, - rotation_angle: float | OptimizationParameter | None = None, - verbose=None, - **reconstruct_kwargs, - ): - if verbose is None: - verbose = self.verbose - - aberration_coefs = aberration_coefs or {} - state = self.hyperparameter_state - - # Reset optimized bookkeeping - state.clear_optimized() - - # Partition inputs - fixed_override_aberrations: dict[str, float] = {} - optimizable_aberrations: dict[str, OptimizationParameter] = {} - - for name, val in aberration_coefs.items(): - if isinstance(val, OptimizationParameter): - optimizable_aberrations[name] = val - state.optimized_keys.add(name) - else: - fixed_override_aberrations[name] = val - - optimize_rotation = isinstance(rotation_angle, OptimizationParameter) - if optimize_rotation: - state.optimized_keys.add("rotation_angle") - - # Build parameter grid (only over optimizable parameters) - param_grid: dict[str, list[float]] = {} - - for name, param in optimizable_aberrations.items(): - param_grid[name] = param.grid_values() - - if optimize_rotation: - param_grid["rotation_angle"] = ( - rotation_angle.grid_values() # ty:ignore[possibly-missing-attribute] - ) - - # Cartesian product - keys = list(param_grid.keys()) - grid = list(product(*(param_grid[k] for k in keys))) - - best_loss = float("inf") - best_params: dict[str, float] | None = None - results = [] - - for combo in tqdm(grid, disable=not verbose): - trial_params = dict(zip(keys, combo)) - - trial_aberrations = dict(fixed_override_aberrations) - trial_aberrations.update( - {k: v for k, v in trial_params.items() if k != "rotation_angle"} - ) - - if optimize_rotation: - rot = trial_params["rotation_angle"] - else: - rot = rotation_angle - - self.reconstruct( - override_aberration_coefs=trial_aberrations, - override_rotation_angle=rot, - verbose=False, - **reconstruct_kwargs, - ) - - loss = float(self.variance_loss()) - results.append((trial_params, loss)) - - if loss < best_loss: - best_loss = loss - best_params = trial_params - - self._grid_search_results = results - - # Write back best optimized values - if best_params is not None: - best_params = best_params.copy() - state.optimized_rotation_angle = best_params.pop("rotation_angle", None) - state.optimized_aberrations = best_params - - if state.optimized_rotation_angle is None and rotation_angle is not None: - state.optimized_rotation_angle = rotation_angle # ty:ignore[invalid-assignment] - - state.optimized_aberrations = state.current_aberrations(fixed_override_aberrations) - - if verbose: - print("Optimized state:\n\n", self.hyperparameter_state) - - # Final reconstruction using merged state - self.reconstruct(verbose=False, **reconstruct_kwargs) - return self - - def _return_lateral_shifts( - self, - rotation_angle, - aberration_coefs, - bf_mask, - ): - # Get initial shifts - kxa, kya = spatial_frequencies( - self.gpts, - self.sampling, - rotation_angle=_rotation_degrees_to_radians(rotation_angle), - device=self.device, - ) - k, phi = polar_coordinates(kxa, kya) - - dx, dy = aberration_surface_cartesian_gradients( - k * self.wavelength, - phi, - aberration_coefs=aberration_coefs, - ) - grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1) - lateral_shifts = grad_k / 2 / np.pi - return lateral_shifts - def fit_hyperparameters_cross_correlation( self, bf_mask: torch.Tensor | None = None, @@ -1550,7 +1187,8 @@ def ap_to_mask(ap, eps=1e-6): # --------------------------------------------------------- # Solve LS # --------------------------------------------------------- - sol = torch.linalg.lstsq(A, b).solution + # Fall back to CPU (to support MPS), as fit_linear_plane does for eigh + sol = torch.linalg.lstsq(A.cpu(), b.cpu()).solution.to(A.device) delta_cartesian = {name: sol[i] for i, name in enumerate(cartesian_basis)} @@ -1681,18 +1319,6 @@ def _reconstruct_all_permutations(self, verbose=None, **reconstruct_kwargs): return recons - def _make_checkerboard_bf_masks(self, gpts, bf_mask): - """ """ - i_coords = torch.arange(gpts[0], device=self.device) - j_coords = torch.arange(gpts[1], device=self.device) - i_grid, j_grid = torch.meshgrid(i_coords, j_coords, indexing="ij") - checkerboard = torch.fft.ifftshift(((i_grid + j_grid) % 2).bool()) - - bf1 = bf_mask & checkerboard - bf2 = bf_mask & (~checkerboard) - - return [bf1, bf2] - def _reconstruct_with_halfsets(self, verbose=None, **reconstruct_kwargs): """ Compute two half-set reconstructions using alternating BF pixels (checkerboard pattern). diff --git a/src/quantem/diffractive_imaging/direct_ptychography_base.py b/src/quantem/diffractive_imaging/direct_ptychography_base.py new file mode 100644 index 00000000..4c9c42f1 --- /dev/null +++ b/src/quantem/diffractive_imaging/direct_ptychography_base.py @@ -0,0 +1,921 @@ +from dataclasses import dataclass, field +from itertools import product +from typing import TYPE_CHECKING, Dict, Tuple + +import numpy as np +import optuna +from numpy.typing import NDArray +from tqdm.auto import tqdm + +from quantem.core import config +from quantem.core.io.serialize import AutoSerialize +from quantem.core.utils.rng import RNGMixin +from quantem.core.utils.utils import electron_wavelength_angstrom, to_numpy +from quantem.core.utils.validators import ( + validate_aberration_coefficients, + validate_gt, + validate_int, + validate_tensor, +) +from quantem.core.visualization import show_2d +from quantem.diffractive_imaging.complex_probe import ( + FourierProbe, + aberration_surface_cartesian_gradients, + evaluate_probe, + polar_coordinates, + spatial_frequencies, +) +from quantem.diffractive_imaging.direct_ptycho_utils import _rotation_degrees_to_radians +from quantem.diffractive_imaging.ptycho_utils import ( + OptimizationParameter as OptimizationParameter, # re-export: the documented path +) + +if TYPE_CHECKING: + import torch +else: + if config.get("has_torch"): + import torch + +optuna.logging.set_verbosity(optuna.logging.WARNING) + + +@dataclass +class HyperparameterState: + initial_aberrations: dict[str, float] = field(default_factory=dict) + initial_rotation_angle: float | None = None + optimized_aberrations: dict[str, float] = field(default_factory=dict) + optimized_rotation_angle: float | None = None + optimized_keys: set[str] = field(default_factory=set) + study: optuna.Study | None = None + + def __post_init__(self): + self.initial_aberrations = validate_aberration_coefficients(dict(self.initial_aberrations)) + self.optimized_aberrations = validate_aberration_coefficients(self.optimized_aberrations) + + if self.optimized_keys: + canonical = validate_aberration_coefficients( + {k: 0.0 for k in self.optimized_keys if k != "rotation_angle"} + ) + self.optimized_keys = set(canonical.keys()) | { + k for k in self.optimized_keys if k == "rotation_angle" + } + + def __repr__(self) -> str: + return self.summarize(which="all") + + def current_aberrations( + self, override_fixed: dict[str, float] | None = None + ) -> Dict[str, float]: + """Return full aberration dictionary (fixed ⊕ optimized).""" + out = dict(self.initial_aberrations) + out.update(self.optimized_aberrations) + if override_fixed is not None: + out.update(validate_aberration_coefficients(override_fixed)) + return out + + def current_rotation_angle(self, override_fixed: float | None = None) -> float: + """Return rotation angle (optimized takes precedence).""" + if override_fixed is not None: + return override_fixed + if self.optimized_rotation_angle is not None: + return self.optimized_rotation_angle + if self.initial_rotation_angle is not None: + return self.initial_rotation_angle + return 0.0 + + def clear_optimized(self): + """Clear all optimized aberrations and rotation angle.""" + self.optimized_aberrations.clear() + self.optimized_rotation_angle = None + self.optimized_keys.clear() + self.study = None + + def clear_all(self): + """Clear everything: initial and optimized hyperparameters.""" + self.initial_aberrations.clear() + self.initial_rotation_angle = None + self.clear_optimized() + + def copy(self): + """ """ + return HyperparameterState( + initial_aberrations=self.initial_aberrations, + optimized_aberrations=self.optimized_aberrations, + initial_rotation_angle=self.initial_rotation_angle, + optimized_rotation_angle=self.optimized_rotation_angle, + optimized_keys=self.optimized_keys, + study=self.study, + ) + + def summarize( + self, + *, + which: str = "current", + override_aberration_coefs: dict[str, float] | None = None, + override_rotation_angle: float | None = None, + ) -> str: + cls = self.__class__.__name__ + lines: list[str] = [] + + def add(name: str, value): + lines.append(f" {name}={value!r},") + + if which == "initial": + if self.initial_aberrations: + add("initial_aberrations", self.initial_aberrations) + if self.initial_rotation_angle is not None: + 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_deg", self.optimized_rotation_angle) + + elif which == "current": + current_abers = self.current_aberrations(override_aberration_coefs) + current_rot = self.current_rotation_angle(override_rotation_angle) + + if current_abers: + add("current_aberrations", current_abers) + if current_rot is not None: + 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_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_deg", self.optimized_rotation_angle) + + else: + raise ValueError( + f"`which` must be one of " + f'{{"initial", "optimized", "current", "all"}}, got {which!r}' + ) + + if not lines: + return f"{cls}()" + + body = "\n".join(lines) + return f"{cls}(\n{body}\n)" + + +@dataclass(frozen=True) +class BrightFieldContext: + bf_mask: torch.Tensor + bf_inds_i: torch.Tensor + bf_inds_j: torch.Tensor + num_bf: int + vbf_index_mapping: torch.Tensor + + +class DirectPtychographyBase(RNGMixin, AutoSerialize): + """ + Shared state and hyperparameter machinery for direct-ptychography reconstructions. + + This class holds everything that does not depend on *how* the deconvolution is + performed: geometry/sampling bookkeeping, the aberration and rotation + :class:`HyperparameterState`, bright-field mask indexing, visualization, and the + optuna / grid-search drivers. Subclasses supply the reconstruction itself. + + Subclass contract + ----------------- + Attributes a subclass must set in ``__init__``, in this order (the ``scan_sampling`` + and ``reciprocal_sampling`` setters read the units and wavelength): + + 1. ``device``, ``verbose``, ``rng`` + 2. ``wavelength``, ``scan_units``, ``detector_units`` + 3. ``scan_sampling``, ``reciprocal_sampling``, ``angular_sampling`` + 4. ``gpts``, ``sampling`` (detector grid), ``bf_mask``, ``semiangle_cutoff`` + 5. ``hyperparameter_state`` + + Methods a subclass must implement: + + - ``reconstruct(*, override_aberration_coefs, override_rotation_angle, verbose, ...)`` + returning ``self`` + - ``variance_loss()`` returning a scalar tensor (a *method*, not a property -- + the optimizers call ``float(self.variance_loss())``) + - ``corrected_bf`` property returning the reconstructed image, or ``None`` + + A subclass whose reconstruction spans more than the scan field of view (e.g. a padded + canvas) must override ``_obj_fov``; upsampling needs no bookkeeping, since + ``_obj_sampling`` reads the object's own shape. + """ + + # --- state the subclass __init__ must provide (annotation only, no default) --- + hyperparameter_state: HyperparameterState + scan_units: Tuple[str, str] + detector_units: Tuple[str, str] + scan_gpts: Tuple[int, int] + gpts: Tuple[int, int] + sampling: Tuple[float, float] + angular_sampling: Tuple[float, float] + + @property + def wavelength(self) -> float: + """Probe wavelength in Angstrom.""" + return self._wavelength + + @wavelength.setter + def wavelength(self, value: float) -> None: + self._wavelength = float(validate_gt(value, 0.0, "wavelength")) + + @property + def fourier_probe(self) -> "FourierProbe | None": + """An empirical complex probe, or ``None`` when the probe is analytic. + + Set it to reconstruct with a measured ``psi(k)`` instead of an aperture plus + aberrations. Everything that reads the aberration surface -- the parallax kernel, + the ``sign(sin(chi))`` phase flip, the defocus gradient, any hyperparameter search + over aberrations -- has no meaning then, and raises. + """ + return getattr(self, "_fourier_probe", None) + + @fourier_probe.setter + def fourier_probe(self, value) -> None: + if value is None: + self._fourier_probe = None + return + if not isinstance(value, FourierProbe): + raise TypeError( + "`fourier_probe` must be a FourierProbe; build one with " + "`FourierProbe.from_array(psi, reciprocal_sampling, wavelength)`." + ) + if value.array is None: + raise ValueError( + "`fourier_probe` is for an empirical probe; an analytic one is already " + "described by `semiangle_cutoff` and the aberration coefficients." + ) + if tuple(value.array.shape) != tuple(self.gpts): + raise ValueError( + f"`fourier_probe` has shape {tuple(value.array.shape)} but the detector grid " + f"is {tuple(self.gpts)}. They must match -- crop the probe the same way the " + "diffraction patterns were cropped." + ) + self._fourier_probe = value.to(self.device) + + def _require_analytic_probe(self, what: str) -> None: + """Guard for everything that only means something for an aperture plus aberrations.""" + if self.fourier_probe is not None: + raise NotImplementedError( + f"{what} is defined by the aberration surface chi(k), which an empirical " + "`fourier_probe` does not have. Use a deconvolution kernel that does not " + "need it -- 'ssb', 'obf' and 'mf' read only the probe itself, and 'icom' " + "does not read the probe at all." + ) + + def _return_probe(self, aberration_coefs) -> "FourierProbe": + """The probe object the overlap function samples, empirical or analytic.""" + probe = self.fourier_probe + if probe is not None: + return probe + return FourierProbe.from_aberrations( + self.wavelength, + self.semiangle_cutoff, + self.angular_sampling, + aberration_coefs, + self.soft_edges, + ) + + def _return_probe_on_grid(self, k, phi, aberration_coefs): + """``psi(k)`` on the detector grid, whose squared sum is the bright-field weight. + + Leaves ``soft_edges`` at ``evaluate_probe``'s default rather than taking + ``self.soft_edges``, so the normalization matches between the two classes. This is + an inconsistency, kept for compatibility. + """ + probe = self.fourier_probe + if probe is not None: + return probe.array + return evaluate_probe( + k * self.wavelength, + phi, + self.semiangle_cutoff, + self.angular_sampling, + self.wavelength, + aberration_coefs=aberration_coefs, + ) + + @staticmethod + def _resolve_wavelength(energy: float | None, wavelength: float | None) -> float: + """Wavelength in Angstrom, from an electron accelerating voltage or given directly. + + ``energy`` goes through the relativistic electron de Broglie formula, which is what + every electron entry point wants. ``wavelength`` skips it, which is what anything + that is not an electron needs: for a 7.9 keV photon the electron formula returns + 0.137 Angstrom against the correct ``hc/E`` = 1.569. + """ + if (energy is None) == (wavelength is None): + raise ValueError( + "Pass exactly one of `energy` (electron accelerating voltage, in volts) or " + "`wavelength` (in Angstrom, for photons or anything else non-electron), " + f"got energy={energy!r} and wavelength={wavelength!r}." + ) + if wavelength is not None: + return float(validate_gt(wavelength, 0.0, "wavelength")) + return electron_wavelength_angstrom(validate_gt(energy, 0.0, "energy")) + + # ------------------------------------------------------------------ + # subclass hooks + # ------------------------------------------------------------------ + + def reconstruct(self, *args, **kwargs): + raise NotImplementedError(f"{type(self).__name__} does not implement reconstruct().") + + def variance_loss(self): + raise NotImplementedError(f"{type(self).__name__} does not implement variance_loss().") + + def rms_gradient_loss(self): + """ + Negated RMS gradient of the reconstruction, per Angstrom -- a sharpness objective. + + The classic autofocus metric: a correctly deconvolved image has sharp edges and a + large gradient, a mis-set aberration blurs them. Negated so that, like + :meth:`variance_loss`, it is minimized. + + It is better conditioned than the variance loss -- 28% dynamic range against 0.08% + over a defocus series, agreeing on the optimum -- because the variance loss compares + bright-field images with each other and saturates once they agree, while this + measures the reconstruction. It is also insensitive to how the canvas is sized, so it + needs no pinning the way a patch fit does. + + Two caveats: + + - It rewards *amplitude*, not only sharpness, since it is not normalized by the + image's own spread. Aberrations and rotation barely change the overall scale, so + this is safe for the searches here, but a hyperparameter that could inflate the + object would game it. + - It is defined for every deconvolution kernel, where + ``DirectPtychographyMontage.variance_loss`` is defined only for the parallax one. + That makes it the way to drive a search over ``"ssb"`` or ``"obf"``. + + Returns + ------- + float or None + ``None`` before :meth:`reconstruct`, mirroring :meth:`variance_loss`. + """ + obj = self.corrected_bf + if obj is None: + return None + if min(obj.shape[-2:]) < 2: + raise ValueError( + f"An object of shape {tuple(obj.shape)} has no gradient to measure; " + "reconstruct onto a canvas at least 2x2." + ) + + # per Angstrom rather than per pixel, so the value is comparable across upsampling + # factors and samplings rather than only within one search + spacing = tuple(float(s) for s in self._obj_sampling) + grad_rows, grad_cols = torch.gradient(obj.to(torch.float32), spacing=spacing, dim=(-2, -1)) + return -float(torch.sqrt((grad_rows.square() + grad_cols.square()).mean())) + + #: objectives the hyperparameter searches accept by name, all minimized + _LOSS_FUNCTIONS = { + "variance": "variance_loss", + "rms_gradient": "rms_gradient_loss", + } + + def _return_loss_value(self, loss) -> float: + """Evaluate a search objective on the current reconstruction.""" + if callable(loss): + return float(loss(self)) + try: + method = self._LOSS_FUNCTIONS[loss] + except (KeyError, TypeError): + raise ValueError( + f"`loss` must be a callable or one of {sorted(self._LOSS_FUNCTIONS)}, got {loss!r}" + ) from None + return float(getattr(self, method)()) + + @property + def corrected_bf(self): + raise NotImplementedError(f"{type(self).__name__} does not implement corrected_bf.") + + @property + def fov(self) -> tuple[float, float]: + """Field of view of the scan, in Angstrom. Fixed by the acquisition.""" + return tuple(n * s for n, s in zip(self.scan_gpts, self.scan_sampling)) + + @property + def scan_origin(self) -> tuple[float, float]: + """Position of scan pixel ``(0, 0)``, in Angstrom, in the caller's coordinates. + + Zero for a raster acquisition, whose grid *defines* the coordinates. The ungridded + constructors anchor the scan grid at the corner of the probe-position bounding box, + and record that corner here so pixels can be mapped back to the positions that were + passed in -- which is what makes two acquisitions of the same region comparable. + """ + return getattr(self, "_scan_origin", (0.0, 0.0)) + + @scan_origin.setter + def scan_origin(self, value) -> None: + if value is None: + self._scan_origin = (0.0, 0.0) + return + origin = tuple(float(v) for v in np.asarray(value, dtype=np.float64).reshape(-1)) + if len(origin) != 2: + raise ValueError(f"`scan_origin` must be a (row, col) pair, got {value!r}") + self._scan_origin = origin + + @property + def _obj_fov(self) -> tuple[float, float]: + """Field of view the reconstructed object spans, in Angstrom. + + Defaults to the scan field of view, which is what a reconstruction sampled on the + scan grid covers at any upsampling factor. Override where the object spans more. + """ + return self.fov + + @property + def obj_origin(self) -> tuple[float, float]: + """Position of object pixel ``(0, 0)``, in Angstrom, in the caller's coordinates. + + Together with :attr:`_obj_sampling` this is the full map from object pixels back to + the probe positions that were passed in: ``origin + pixel * sampling``. Defaults to + :attr:`scan_origin`, for a reconstruction sampled on the scan grid; a class whose + canvas starts elsewhere must override it. + """ + return self.scan_origin + + @property + def _obj_sampling(self) -> tuple[float, float]: + """Real-space sampling of the reconstructed object, in Angstrom. + + Derived from the object's own shape rather than tracked across reconstructions, so + upsampling needs no bookkeeping and the scalebar cannot fall out of sync. + """ + obj = self.corrected_bf + if obj is None: + return tuple(self.scan_sampling) + return tuple(f / n for f, n in zip(self._obj_fov, obj.shape[-2:])) + + # ------------------------------------------------------------------ + # properties + # ------------------------------------------------------------------ + + @property + def verbose(self) -> int: + return self._verbose + + @verbose.setter + def verbose(self, v: bool | int | float) -> None: + self._verbose = validate_int(validate_gt(v, -1, "verbose"), "verbose") + + @property + def bf_mask(self) -> torch.Tensor: + return self._bf_mask + + @bf_mask.setter + def bf_mask(self, value: torch.Tensor): + self._bf_mask = validate_tensor(value, "bf_mask", dtype=torch.bool).to(device=self.device) + + @property + def rotation_angle(self) -> float: + """Current detector rotation angle in degrees.""" + return self.hyperparameter_state.current_rotation_angle() + + @property + def aberration_coefs(self) -> dict: + return self.hyperparameter_state.current_aberrations() + + @property + def semiangle_cutoff(self) -> float: + return self._semiangle_cutoff + + @semiangle_cutoff.setter + def semiangle_cutoff(self, value: float): + if value is None: + # an empirical probe already carries its own aperture, whatever shape it is + if self.fourier_probe is not None: + self._semiangle_cutoff = None + return + raise ValueError( + "`semiangle_cutoff` is required, in mrad: it sets the aperture used to build " + "the probe and the deconvolution kernels. Pass a `fourier_probe` instead if " + "the probe is measured rather than described by an aperture." + ) + validate_gt(value, 0.0, "semiangle_cutoff") + self._semiangle_cutoff = value + + @property + def device(self) -> str | torch.device: + """This should be of form 'cuda:X' or 'cpu', as defined by quantem.config""" + if hasattr(self, "_device"): + return self._device # ty:ignore[invalid-return-type] + else: + return config.get("device") + + @device.setter + def device(self, device: str | int | None): + if device is not None: + dev, _id = config.validate_device(device) + self._device = dev + + @property + def scan_sampling(self) -> NDArray: + return self._scan_sampling # ty:ignore[invalid-return-type] + + @scan_sampling.setter + def scan_sampling(self, value: NDArray | tuple | list) -> None: + """ + Units A or raises error + """ + units = self.scan_units + if units[0] == "A": + self._scan_sampling = value + else: + raise ValueError("real-space needs to be given in 'A'") + + @property + def reciprocal_sampling(self) -> NDArray: + return self._reciprocal_sampling # ty:ignore[invalid-return-type] + + @reciprocal_sampling.setter + def reciprocal_sampling(self, value: NDArray | tuple | list) -> None: + """ + Units A or raises error + """ + units = self.detector_units + if units[0] == "A^-1": + self._reciprocal_sampling = value + elif units[0] == "mrad": + self._reciprocal_sampling = tuple(val / self.wavelength / 1e3 for val in value) + else: + raise ValueError("reciprocal-space needs to be given in 'A^-1' or 'mrad'") + + @property + def obj(self) -> np.ndarray | None: + """Reconstructed object as a numpy array, or ``None`` before :meth:`reconstruct`. + + Mirrors ``corrected_bf`` rather than raising: ``AutoSerialize._recursive_load`` + walks ``dir(obj)`` and evaluates every property, so a ``to_numpy(None)`` here made + save/load fail for a instance that had not been reconstructed yet. + """ + corrected_bf = self.corrected_bf + if corrected_bf is None: + return None + return to_numpy(corrected_bf) + + # ------------------------------------------------------------------ + # bright-field geometry + # ------------------------------------------------------------------ + + def _return_bf_context(self, bf_mask): + """ + Given a BF mask, compute all BF-dependent geometry and indexing. + """ + bf_mask = torch.as_tensor(bf_mask, dtype=torch.bool, device=self.device) + + bf_inds_i, bf_inds_j = torch.nonzero(bf_mask, as_tuple=True) + vbf_index_mapping = torch.where(bf_mask[self.bf_mask])[0] + num_bf = bf_inds_i.numel() + + return BrightFieldContext( + bf_mask=bf_mask, + bf_inds_i=bf_inds_i, + bf_inds_j=bf_inds_j, + num_bf=num_bf, + vbf_index_mapping=vbf_index_mapping, + ) + + def _make_checkerboard_bf_masks(self, gpts, bf_mask): + """ """ + i_coords = torch.arange(gpts[0], device=self.device) + j_coords = torch.arange(gpts[1], device=self.device) + i_grid, j_grid = torch.meshgrid(i_coords, j_coords, indexing="ij") + checkerboard = torch.fft.ifftshift(((i_grid + j_grid) % 2).bool()) + + bf1 = bf_mask & checkerboard + bf2 = bf_mask & (~checkerboard) + + return [bf1, bf2] + + def _normalize_kernel_name(self, kernel): + kernel = kernel.lower() + + aliases = { + "ssb": "ssb", + "single-sideband": "ssb", + "acbf": "ssb", + "aberration-corrected-bright-field": "ssb", + "obf": "obf", + "optimum-bright-field": "obf", + "mf": "mf", + "matched-filter": "mf", + "prlx": "prlx", + "parallax": "prlx", + "tcbf": "prlx", + "tilt-corrected-bright-field": "prlx", + "icom": "icom", + "center-of-mass": "icom", + } + + if kernel not in aliases: + raise ValueError(f"Unknown deconvolution kernel '{kernel}'") + + return aliases[kernel] + + def _return_lateral_shifts( + self, + rotation_angle, + aberration_coefs, + bf_mask, + ): + """Aberration-induced lateral shift of each BF pixel, in Angstrom.""" + # Get initial shifts + kxa, kya = spatial_frequencies( + self.gpts, + self.sampling, + rotation_angle=_rotation_degrees_to_radians(rotation_angle), + device=self.device, + ) + k, phi = polar_coordinates(kxa, kya) + + dx, dy = aberration_surface_cartesian_gradients( + k * self.wavelength, + phi, + aberration_coefs=aberration_coefs, + ) + grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1) + lateral_shifts = grad_k / 2 / np.pi + return lateral_shifts + + # ------------------------------------------------------------------ + # visualization + # ------------------------------------------------------------------ + + def visualize( + self, + return_fig: bool = False, + show_obj_fft: bool = True, + apply_hanning_window: 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 + obj_sampling = self._obj_sampling + obj_scalebar = {"sampling": obj_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 / (obj_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 + return None + + # ------------------------------------------------------------------ + # hyperparameter optimization + # ------------------------------------------------------------------ + + def optimize_hyperparameters( + self, + aberration_coefs: dict[str, float | OptimizationParameter] | None = None, + rotation_angle: float | OptimizationParameter | None = None, + n_trials=50, + sampler=None, + loss="variance", + verbose=None, + **reconstruct_kwargs, + ): + """ + Optimize hyperparameters (aberrations and/or rotation) using Optuna. + + Parameters + ---------- + 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, in degrees. + n_trials : int + Number of Optuna trials. + sampler : optuna.samplers.BaseSampler, optional + Custom Optuna sampler. + loss : {"variance", "rms_gradient"} or callable + Objective to minimize. ``"variance"`` is :meth:`variance_loss`, the spread + between bright-field images. ``"rms_gradient"`` is + :meth:`rms_gradient_loss`, an image-sharpness objective that is better + conditioned -- 28% dynamic range against 0.08% over a defocus series -- and is + defined for every deconvolution kernel. A callable is passed the reconstruction + and must return a float to minimize. + verbose : bool, optional + Report the search and show Optuna's progress bar. Defaults to :attr:`verbose`. + **reconstruct_kwargs : + Extra arguments passed to reconstruct(). + """ + + if verbose is None: + verbose = self.verbose + + sampler = sampler or optuna.samplers.TPESampler() + + state = self.hyperparameter_state + aberration_coefs = aberration_coefs or {} + + # Reset optimized bookkeeping + state.clear_optimized() + + # Partition inputs + fixed_override_aberrations = {} + optimizable_aberrations = {} + + for name, val in aberration_coefs.items(): + if isinstance(val, OptimizationParameter): + optimizable_aberrations[name] = val + state.optimized_keys.add(name) + else: + fixed_override_aberrations[name] = val + + if isinstance(rotation_angle, OptimizationParameter): + state.optimized_keys.add("rotation_angle") + + def objective(trial): + trial_aberrations = {} + for name, val in optimizable_aberrations.items(): + trial_aberrations[name] = trial.suggest_float(name, val.low, val.high, log=val.log) + + trial_aberrations |= fixed_override_aberrations + + if isinstance(rotation_angle, OptimizationParameter): + rot = trial.suggest_float( + "rotation_angle", + rotation_angle.low, + rotation_angle.high, + log=rotation_angle.log, + ) + else: + rot = rotation_angle + + self.reconstruct( + override_aberration_coefs=trial_aberrations, + override_rotation_angle=rot, + verbose=False, + **reconstruct_kwargs, + ) + return self._return_loss_value(loss) + + study = optuna.create_study(direction="minimize", sampler=sampler) + study.optimize(objective, n_trials=n_trials, show_progress_bar=bool(verbose)) + + # Write back optimized results + best = study.best_params.copy() + state.optimized_rotation_angle = best.pop("rotation_angle", None) + + if state.optimized_rotation_angle is None and rotation_angle is not None: + state.optimized_rotation_angle = rotation_angle # ty:ignore[invalid-assignment] + + state.optimized_aberrations = best + state.optimized_aberrations = state.current_aberrations(fixed_override_aberrations) + state.study = study + + if verbose: + print("Optimized state:\n\n", self.hyperparameter_state) + + self.reconstruct(verbose=False, **reconstruct_kwargs) + return self + + def grid_search_hyperparameters( + self, + aberration_coefs: dict[str, float | OptimizationParameter] | None = None, + rotation_angle: float | OptimizationParameter | None = None, + loss="variance", + verbose=None, + **reconstruct_kwargs, + ): + """ + Exhaustive search over a grid of hyperparameter values. + + Parameters + ---------- + loss : {"variance", "rms_gradient"} or callable + Objective to minimize; see :meth:`optimize_hyperparameters`. + """ + if verbose is None: + verbose = self.verbose + + aberration_coefs = aberration_coefs or {} + state = self.hyperparameter_state + + # Reset optimized bookkeeping + state.clear_optimized() + + # Partition inputs + fixed_override_aberrations: dict[str, float] = {} + optimizable_aberrations: dict[str, OptimizationParameter] = {} + + for name, val in aberration_coefs.items(): + if isinstance(val, OptimizationParameter): + optimizable_aberrations[name] = val + state.optimized_keys.add(name) + else: + fixed_override_aberrations[name] = val + + optimize_rotation = isinstance(rotation_angle, OptimizationParameter) + if optimize_rotation: + state.optimized_keys.add("rotation_angle") + + # Build parameter grid (only over optimizable parameters) + param_grid: dict[str, list[float]] = {} + + for name, param in optimizable_aberrations.items(): + param_grid[name] = param.grid_values() + + # isinstance inline rather than reusing `optimize_rotation`, so the type narrows + if isinstance(rotation_angle, OptimizationParameter): + param_grid["rotation_angle"] = rotation_angle.grid_values() + + # Cartesian product + keys = list(param_grid.keys()) + grid = list(product(*(param_grid[k] for k in keys))) + + best_loss = float("inf") + best_params: dict[str, float] | None = None + results = [] + + for combo in tqdm(grid, disable=not verbose): + trial_params = dict(zip(keys, combo)) + + trial_aberrations = dict(fixed_override_aberrations) + trial_aberrations.update( + {k: v for k, v in trial_params.items() if k != "rotation_angle"} + ) + + if optimize_rotation: + rot = trial_params["rotation_angle"] + else: + rot = rotation_angle + + self.reconstruct( + override_aberration_coefs=trial_aberrations, + override_rotation_angle=rot, + verbose=False, + **reconstruct_kwargs, + ) + + loss_value = self._return_loss_value(loss) + results.append((trial_params, loss_value)) + + if loss_value < best_loss: + best_loss = loss_value + best_params = trial_params + + self._grid_search_results = results + + # Write back best optimized values + if best_params is not None: + best_params = best_params.copy() + state.optimized_rotation_angle = best_params.pop("rotation_angle", None) + state.optimized_aberrations = best_params + + if state.optimized_rotation_angle is None and rotation_angle is not None: + state.optimized_rotation_angle = rotation_angle # ty:ignore[invalid-assignment] + + state.optimized_aberrations = state.current_aberrations(fixed_override_aberrations) + + if verbose: + print("Optimized state:\n\n", self.hyperparameter_state) + + # Final reconstruction using merged state + self.reconstruct(verbose=False, **reconstruct_kwargs) + return self diff --git a/src/quantem/diffractive_imaging/direct_ptychography_montage.py b/src/quantem/diffractive_imaging/direct_ptychography_montage.py new file mode 100644 index 00000000..3a54dca0 --- /dev/null +++ b/src/quantem/diffractive_imaging/direct_ptychography_montage.py @@ -0,0 +1,2114 @@ +import gc +import math +import warnings +from typing import TYPE_CHECKING, Literal, Tuple + +import numpy as np +from numpy.typing import NDArray +from tqdm.auto import tqdm + +from quantem.core import config +from quantem.core.datastructures import Dataset2d, Dataset3d, Dataset4d +from quantem.core.utils.utils import to_numpy +from quantem.core.utils.validators import ( + validate_aberration_coefficients, + validate_tensor, +) +from quantem.diffractive_imaging.complex_probe import ( + FourierProbe, + aberration_surface, + aberration_surface_cartesian_gradients, + gamma_factor, + polar_coordinates, + spatial_frequencies, +) +from quantem.diffractive_imaging.ptycho_utils import SimpleBatcher + +if TYPE_CHECKING: + import torch +else: + if config.get("has_torch"): + import torch + +from quantem.diffractive_imaging.direct_ptycho_utils import ( + _crop_corner_centered_mask, + _rotation_degrees_to_radians, + allocate_splat_buffers, + build_vbf_stack_from_dataset3d, + build_vbf_stack_from_dataset4d, + convolve_stack_fourier, + preferred_float_dtype, + scatter_add_splat, + splat_and_convolve, + splat_stack, +) +from quantem.diffractive_imaging.direct_ptychography_base import DirectPtychographyBase + +# target number of (BF pixel, scan position) points per splat batch +_DEFAULT_POINTS_PER_BATCH = 4_194_304 + + +def _snap_to_integer(values: torch.Tensor, tolerance: float = 1e-4) -> torch.Tensor: + """Round values that are integers to within ``tolerance``, leave the rest alone. + + Canvas bounds go through ``floor``/``ceil``, where a shift of exactly 4 arriving as + 4.0000001 costs a whole pixel. The k-grid is float32 and positions are float32 on MPS + (which has no float64), so that noise is unavoidable -- and without snapping the same + data yields a canvas one pixel larger on CPU and a different one again on MPS. + """ + rounded = torch.round(values) + return torch.where((values - rounded).abs() < tolerance, rounded, values) + + +class DirectPtychographyMontage(DirectPtychographyBase): + """ + Direct ptychography that montages the scan onto a shared canvas. + + Every kernel of + :class:`~quantem.diffractive_imaging.direct_ptychography.DirectPtychography` is a + multiplier on the scan-space Fourier transform. Here that transform is not taken: + each virtual bright-field image is deposited onto one canvas at its own probe position. + The detector axis is handled the same way in both classes, summed over bright-field + pixels each carrying its own kernel. + + Kernels + ------- + ``prlx`` + A pure translation by ``grad_chi / (2 * pi)`` Angstrom, exact, one deposit per + point. The shadow-montage (tilt-corrected bright field) construction [1]_. + ``ssb``, ``obf``, ``mf`` + Convolutions rather than translations. Exact with ``convolution_mode="fft"`` (the + default); truncated to a box stencil with ``"stencil"``. + ``icom`` + Exact by FFT. Truncated it is riCOM [2]_, where the radius is a high-pass cutoff + rather than an error. Being linear in ``k``, it collapses to two convolutions of + the centre-of-mass shift regardless of the number of bright-field pixels. + + The parallax equivalence holds both ways: ``exp(-1j * grad_chi . q)`` is a translation, + and Fourier-space tiling by ``U`` is real-space zero-insertion at every ``U``-th pixel. + Working in real space instead gives: + + - no scan-space FFT, so the scan positions need not lie on a grid -- + see :meth:`from_dataset3d`; + - the phase-flip and Butterworth filters, which do not depend on the bright-field index, + collapse into a single post-hoc filter on the finished image; + - a per-position defocus, which models a tilted sample -- see :attr:`defocus_gradient` + and :meth:`fit_defocus_gradient`. A Fourier multiplier is global over the scan and + cannot express this. + + Choosing between the two classes + -------------------------------- + ================================== ========================================== + gridded scan ``DirectPtychography`` -- exact, and + cheapest at one scan FFT + ungridded scan either: ``from_dataset3d`` regrids onto a + lattice first, this class never grids + sub-pixel positions matter here -- regridding discards them (0.997 + against 0.982 CTF correlation, quoted in + ``DirectPtychography.from_dataset3d``) + scan both masked and upsampled here -- ``hole_fill`` cannot serve filled + holes and deliberate gaps at once + position-dependent defocus here only + large bright-field mask here -- see below + ================================== ========================================== + + The last is a memory limit rather than a preference. ``DirectPtychography._preprocess`` + materializes the scan transform as ``(N_bf, Ry, Rx)`` complex64 -- 34 GB at 167k + bright-field pixels on a 128x170 canvas -- where this class streams over detector pixels + into one canvas for roughly 800 MB. + + Instantiate with :meth:`from_dataset4d`, :meth:`from_virtual_bfs` or + :meth:`from_dataset3d`. + + References + ---------- + .. [1] *Microscopy and Microanalysis* 32(1), ozaf126 (2026). + https://doi.org/10.1093/mam/ozaf126 + .. [2] Yu et al., *Microscopy and Microanalysis* 28, 1526 (2022). riCOM. + + Related, though neither is the kernel implemented here: the first convolves a WDD + kernel where this convolves SSB/OBF/MF kernels, and the second uses a segmented + detector rather than a pixelated one. + + .. [3] Convolution WDD: *Ultramicroscopy* 285, 114411 (2026). + https://doi.org/10.1016/j.ultramic.2026.114411 + .. [4] Segmented-detector OBF: *Ultramicroscopy* 220, 113133 (2021). + https://doi.org/10.1016/j.ultramic.2020.113133 + """ + + _token = object() + + def __init__( + self, + vbf_stack: torch.Tensor | NDArray, + positions_px: torch.Tensor | NDArray, + bf_mask_dataset: Dataset2d, + energy: float | None, + rotation_angle: float, + aberration_coefs: dict, + semiangle_cutoff: float, + scan_sampling: Tuple[float, float], + scan_units: Tuple[str, str], + scan_gpts: Tuple[int, int], + boundary: Literal["wrap", "pad"], + gridded_scan: bool, + subtract_frame_mean: bool, + soft_edges: bool, + crop_bf_mask: bool, + bf_mask_padding_px: int, + rng: np.random.Generator | int | None, + device: str | int, + verbose: int | bool, + defocus_gradient: Tuple[float, float] | None = None, + scan_origin: Tuple[float, float] | None = None, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, + _token: object | None = None, + ): + """ """ + if _token is not self._token: + raise RuntimeError( + "Use DirectPtychographyMontage.from_dataset4d(), .from_virtual_bfs() or " + ".from_dataset3d() to instantiate this class." + ) + + self.device = device + self.verbose = verbose + self.vbf_stack = vbf_stack + self.positions_px = positions_px + self.bf_mask = bf_mask_dataset.array # ty:ignore[invalid-assignment] + if crop_bf_mask: + self.bf_mask = _crop_corner_centered_mask(self.bf_mask, bf_mask_padding_px) + + if rotation_angle is None: + raise ValueError( + "`rotation_angle` is required, in degrees: it sets the detector rotation " + "relative to the scan. Pass 0.0 if the two frames already agree." + ) + + self.wavelength = self._resolve_wavelength(energy, wavelength) + self.scan_units = scan_units + self.detector_units = bf_mask_dataset.units + + self.scan_gpts = tuple(int(n) for n in scan_gpts) + self.scan_sampling = scan_sampling + self.scan_origin = scan_origin + self.reciprocal_sampling = bf_mask_dataset.sampling + self.angular_sampling = tuple(d * 1e3 * self.wavelength for d in self.reciprocal_sampling) + + self.num_bf = int(self.vbf_stack.shape[0]) + self.num_positions = int(self.vbf_stack.shape[1]) + self.gpts = tuple(int(n) for n in self.bf_mask.shape[:2]) + self.sampling = tuple(1 / s / n for n, s in zip(self.reciprocal_sampling, self.gpts)) + + self.fourier_probe = fourier_probe + self.semiangle_cutoff = semiangle_cutoff + self.soft_edges = soft_edges + self.boundary = boundary + #: whether the positions lie on a regular lattice; drives the sampling-density + #: correction that `weight_normalize` applies + self.gridded_scan = gridded_scan + self.subtract_frame_mean = subtract_frame_mean + self.defocus_gradient = defocus_gradient + self.rng = rng + + if self.positions_px.shape[0] != self.num_positions: + raise ValueError( + f"`positions_px` has {self.positions_px.shape[0]} rows but `vbf_stack` has " + f"{self.num_positions} scan positions." + ) + + self.hyperparameter_state = self._make_hyperparameter_state( + aberration_coefs, rotation_angle + ) + + self._preprocess() + + @staticmethod + def _make_hyperparameter_state(aberration_coefs, rotation_angle): + from quantem.diffractive_imaging.direct_ptychography_base import HyperparameterState + + return HyperparameterState( + initial_aberrations=aberration_coefs, initial_rotation_angle=rotation_angle + ) + + # ------------------------------------------------------------------ + # constructors + # ------------------------------------------------------------------ + + @classmethod + def from_virtual_bfs( + cls, + vbf_dataset: Dataset3d, + bf_mask_dataset: Dataset2d, + energy: float | None = None, + rotation_angle: float | None = None, + semiangle_cutoff: float | None = None, + aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, + boundary: Literal["wrap", "pad"] = "wrap", + defocus_gradient: Tuple[float, float] | None = None, + subtract_frame_mean: bool = False, + soft_edges: bool = True, + crop_bf_mask: bool = True, + bf_mask_padding_px: int = 1, + rng: np.random.Generator | int | None = None, + device: str | int = "cpu", + verbose: int | bool = True, + ): + """ + Build from a gridded virtual bright-field stack. + + Accepts exactly the ``(N_bf, Rx, Ry)`` ``Dataset3d`` that + :meth:`DirectPtychography.from_virtual_bfs` takes, so the same stack can be fed to + both classes; the trailing scan axes are flattened internally. + """ + scan_gpts = tuple(int(n) for n in vbf_dataset.shape[-2:]) + vbf_stack = np.asarray(vbf_dataset.array).reshape(vbf_dataset.shape[0], -1) + + return cls( + vbf_stack=vbf_stack, + positions_px=cls._raster_positions_px(scan_gpts), + bf_mask_dataset=bf_mask_dataset, + energy=energy, + wavelength=wavelength, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + semiangle_cutoff=semiangle_cutoff, + scan_sampling=tuple(vbf_dataset.sampling[-2:]), + scan_units=tuple(vbf_dataset.units[-2:]), + scan_gpts=scan_gpts, + fourier_probe=fourier_probe, + boundary=boundary, + gridded_scan=True, + defocus_gradient=defocus_gradient, + subtract_frame_mean=subtract_frame_mean, + soft_edges=soft_edges, + crop_bf_mask=crop_bf_mask, + bf_mask_padding_px=bf_mask_padding_px, + rng=rng, + device=device, + verbose=verbose, + _token=cls._token, + ) + + @classmethod + def from_dataset4d( + cls, + dataset: Dataset4d, + energy: float | None = None, + semiangle_cutoff: float | None = None, + aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, + rotation_angle: float | None = None, + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + force_fitted_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + intensity_threshold: float = 0.5, + boundary: Literal["wrap", "pad"] = "wrap", + defocus_gradient: Tuple[float, float] | None = None, + subtract_frame_mean: bool = False, + soft_edges: bool = True, + crop_bf_mask: bool = True, + bf_mask_padding_px: int = 1, + rng: np.random.Generator | int | None = None, + device: str | int = "cpu", + verbose: int | bool = True, + normalization_order: int = 0, + edge_blend_pixels: int = 0, + ): + """ + Build from a raster-scanned 4D-STEM dataset. + + Runs the same origin-correction, bright-field masking and normalization pipeline as + :meth:`DirectPtychography.from_dataset4d` (they share + :func:`~quantem.diffractive_imaging.direct_ptycho_utils.build_vbf_stack_from_dataset4d`), + then flattens the scan axes onto an integer position grid. + """ + vbf_dataset, bf_mask_dataset, rotation_angle = build_vbf_stack_from_dataset4d( + dataset, + device=device, + max_batch_size=max_batch_size, + fit_method=fit_method, + mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + intensity_threshold=intensity_threshold, + normalization_order=normalization_order, + edge_blend_pixels=edge_blend_pixels, + ) + + return cls.from_virtual_bfs( + vbf_dataset=vbf_dataset, + bf_mask_dataset=bf_mask_dataset, + energy=energy, + wavelength=wavelength, + rotation_angle=rotation_angle, + semiangle_cutoff=semiangle_cutoff, + aberration_coefs=aberration_coefs, + fourier_probe=fourier_probe, + boundary=boundary, + defocus_gradient=defocus_gradient, + subtract_frame_mean=subtract_frame_mean, + soft_edges=soft_edges, + crop_bf_mask=crop_bf_mask, + bf_mask_padding_px=bf_mask_padding_px, + rng=rng, + device=device, + verbose=verbose, + ) + + @classmethod + def from_dataset3d( + cls, + dataset: Dataset3d, + positions: Dataset2d | torch.Tensor | NDArray, + energy: float | None = None, + semiangle_cutoff: float | None = None, + rotation_angle: float | None = None, + scan_sampling: Tuple[float, float] | Literal["auto"] = "auto", + aberration_coefs: dict = {}, + wavelength: float | None = None, + fourier_probe: "FourierProbe | None" = None, + bf_mask: torch.Tensor | NDArray | None = None, + max_batch_size: int | None = None, + fit_method: str = "plane", + mode: str = "bilinear", + force_measured_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + force_fitted_origin: Tuple[float, float] | torch.Tensor | NDArray | None = None, + intensity_threshold: float = 0.5, + boundary: Literal["wrap", "pad"] = "pad", + defocus_gradient: Tuple[float, float] | None = None, + subtract_frame_mean: bool = False, + soft_edges: bool = True, + crop_bf_mask: bool = True, + bf_mask_padding_px: int = 1, + rng: np.random.Generator | int | None = None, + device: str | int = "cpu", + verbose: int | bool = True, + normalization_order: int = 0, + ): + """ + Build from an ungridded stack of diffraction patterns and their probe positions. + + Parameters + ---------- + dataset : Dataset3d + ``(N, Qx, Qy)`` diffraction patterns, reciprocal units ``"A^-1"`` or ``"mrad"``. + positions : Dataset2d, torch.Tensor or ndarray + ``(N, 2)`` probe positions in Angstrom, ordered ``(row, col)`` to match the + diffraction axes. A ``Dataset2d`` must carry units ``"A"``. + rotation_angle : float + Detector rotation in degrees. Required: rotation is otherwise estimated from the + curl of the center of mass over a 2D scan grid, which an ungridded scan lacks. + scan_sampling : tuple of float or "auto" + Canvas pixel size in Angstrom. ``"auto"`` uses the median nearest-neighbour + position spacing and warns with the inferred value. + + Notes + ----- + Positions are *not* rotated: the detector rotation already enters through the + bright-field k-grid, and rotating the positions as well would double-count it. + """ + ( + vbf_stack, + positions_px, + bf_mask_dataset, + scan_gpts, + scan_sampling, + rotation_angle, + scan_origin, + ) = build_vbf_stack_from_dataset3d( + dataset, + positions, + scan_sampling, + device=device, + max_batch_size=max_batch_size, + fit_method=fit_method, + mode=mode, + force_measured_origin=force_measured_origin, + force_fitted_origin=force_fitted_origin, + rotation_angle=rotation_angle, + intensity_threshold=intensity_threshold, + normalization_order=normalization_order, + bf_mask=bf_mask, + ) + + return cls( + vbf_stack=vbf_stack, + positions_px=positions_px, + bf_mask_dataset=bf_mask_dataset, + energy=energy, + wavelength=wavelength, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + semiangle_cutoff=semiangle_cutoff, + scan_sampling=scan_sampling, + scan_units=("A", "A"), + scan_gpts=scan_gpts, + scan_origin=scan_origin, + fourier_probe=fourier_probe, + boundary=boundary, + gridded_scan=False, + defocus_gradient=defocus_gradient, + subtract_frame_mean=subtract_frame_mean, + soft_edges=soft_edges, + crop_bf_mask=crop_bf_mask, + bf_mask_padding_px=bf_mask_padding_px, + rng=rng, + device=device, + verbose=verbose, + _token=cls._token, + ) + + @staticmethod + def _raster_positions_px(scan_gpts: Tuple[int, int]) -> NDArray: + """Integer ``(Rx*Ry, 2)`` raster positions in scan pixels, "ij" ordered.""" + ii, jj = np.meshgrid(np.arange(scan_gpts[0]), np.arange(scan_gpts[1]), indexing="ij") + return np.stack((ii.ravel(), jj.ravel()), axis=-1).astype(np.float64) + + # ------------------------------------------------------------------ + # properties + # ------------------------------------------------------------------ + + @property + def vbf_stack(self) -> torch.Tensor: + """``(N_bf, N_pos)`` virtual bright-field stack, flattened over scan positions.""" + return self._vbf_stack + + @vbf_stack.setter + def vbf_stack(self, value): + stack = validate_tensor(value, "vbf_stack", dtype=torch.float).to(device=self.device) + if stack.ndim != 2: + raise ValueError( + f"`vbf_stack` must have shape (N_bf, N_pos), got {tuple(stack.shape)}" + ) + self._vbf_stack = stack + + @property + def _float_dtype(self) -> torch.dtype: + """Widest float this device supports. MPS has no float64, so everything positional + -- coordinates, shifts, canvas origins -- has to follow the device rather than + hardcode float64.""" + return preferred_float_dtype(self.device) + + @property + def positions_px(self) -> torch.Tensor: + """``(N_pos, 2)`` scan positions in canvas pixels at ``upsampling_factor=1``.""" + return self._positions_px + + @positions_px.setter + def positions_px(self, value): + positions = validate_tensor(value, "positions_px", dtype=self._float_dtype).to( + device=self.device + ) + if positions.ndim != 2 or positions.shape[1] != 2: + raise ValueError( + f"`positions_px` must have shape (N_pos, 2), got {tuple(positions.shape)}" + ) + self._positions_px = positions + + @property + def defocus_gradient(self) -> Tuple[float, float] | None: + """``(d C10 / d row, d C10 / d col)`` in Angstrom per Angstrom, or ``None``. + + Models a tilted sample, whose defocus varies linearly across the field of view as + ``C10(r) = C10_global + g . (r - r_centroid)``. The magnitude is the tangent of the + sample tilt, so a 5 degree tilt is ``|g| = 0.087``. + + Measuring from the centroid of the scan positions makes ``mean(delta C10) = 0`` + exactly, so the gradient is orthogonal to the global ``C10``: a hyperparameter search + over ``C10`` stays well posed with a gradient set. + + This only matters when the defocus swing across the field of view is comparable to + the depth of field ``wavelength / semiangle**2`` -- about 1200 Angstrom at 4 mrad, + where it is irrelevant, but only 22 Angstrom at 30 mrad, where it dominates. + """ + return self._defocus_gradient + + @defocus_gradient.setter + def defocus_gradient(self, value): + if value is None: + self._defocus_gradient = None + return + value = tuple(float(v) for v in np.asarray(value, dtype=np.float64).reshape(-1)) + if len(value) != 2: + raise ValueError( + f"`defocus_gradient` must be a (row, col) pair or None, got {value!r}" + ) + self._defocus_gradient = value + + @property + def defocus_map_results(self) -> dict | None: + """What the last :meth:`defocus_map` measured, or ``None`` if it has not run. + + :meth:`fit_defocus_gradient` leaves its map here, so the per-patch loss curves can be + plotted without paying for a second pass -- and a fit should be looked at before it + is trusted, since a patch whose minimum sits on an endpoint of ``c10_values`` is + dropped from the plane silently apart from the reported count. + """ + return getattr(self, "_defocus_map_results", None) + + @property + def positions_centroid_px(self) -> torch.Tensor: + """Centroid of the scan positions, in canvas pixels. Where ``delta C10`` vanishes.""" + return self._positions_px.mean(dim=0) + + @property + def corrected_bf(self) -> torch.Tensor | None: + """Reconstructed phase image, or ``None`` before :meth:`reconstruct`.""" + return self._corrected_bf + + @property + def weights(self) -> torch.Tensor | None: + """Accumulated splat weight per canvas pixel -- the montage's local support.""" + if self._sum_w is None: + return None + return self._sum_w.reshape(self._canvas_shape) + + @property + def variance_map(self) -> torch.Tensor | None: + """Per-pixel variance across bright-field images (see :meth:`variance_loss`).""" + if self._sum_wv2 is None: + return None + _, var, _ = self._weighted_moments() + return var.reshape(self._canvas_shape) + + @property + def _obj_fov(self) -> tuple[float, float]: + """Field of view of the canvas, in Angstrom. + + With ``boundary="pad"`` the canvas grows past the scan to cover the shifted + positions, so it spans more than :attr:`fov`. Computed by :meth:`_return_canvas` + alongside the canvas shape, so the two always agree. + """ + if self._canvas_fov is None: + return self.fov + return self._canvas_fov + + @property + def obj_origin(self) -> tuple[float, float]: + """Position of object pixel ``(0, 0)``, in Angstrom, in the caller's coordinates. + + The canvas corner, which ``"pad"`` places below the scan origin to make room for the + aberration shifts. With :attr:`_obj_sampling` this maps the reconstruction back onto + the probe positions that were passed in, and hence onto any other reconstruction of + the same region -- see :meth:`reconstruct`'s ``obj_origin`` and ``obj_fov``. + """ + if self._canvas_origin_px is None: + return self.scan_origin + origin_px = to_numpy(self._canvas_origin_px) + return tuple( + float(o + p * s) for o, p, s in zip(self.scan_origin, origin_px, self._obj_sampling) + ) + + # ------------------------------------------------------------------ + # preprocessing + # ------------------------------------------------------------------ + + def _preprocess(self): + """ + Remove the scan mean of each bright-field image. + + This is the real-space equivalent of zeroing the DC bin of each image's scan-space + Fourier transform, which is what ``DirectPtychography._preprocess`` does. It also + centers the accumulated values on zero, which keeps the ``E[v^2] - E[v]^2`` variance + accumulation well conditioned. + """ + self._dc_per_image = self._vbf_stack.mean(dim=1) + self._vbf_stack = self._vbf_stack - self._dc_per_image[:, None] + + if self.subtract_frame_mean: + self._vbf_stack = self._vbf_stack - self._vbf_stack.mean(dim=0, keepdim=True) + + self._reset_reconstruction() + return self + + def _reset_reconstruction(self): + self._sum_w = None + self._sum_wv = None + self._sum_wv2 = None + self._corrected_bf = None + self._canvas_shape = None + self._canvas_origin_px = None + self._canvas_fov = None + self._bf_weights = None + self._kernel = "prlx" + self._stencil_info = None + + # ------------------------------------------------------------------ + # reconstruction + # ------------------------------------------------------------------ + + def _return_k_grid(self, rotation_angle): + """``(kxa, kya, k, phi)`` on the rotated detector k-grid.""" + kxa, kya = spatial_frequencies( + self.gpts, + self.sampling, + rotation_angle=_rotation_degrees_to_radians(rotation_angle), + device=self.device, + ) + k, phi = polar_coordinates(kxa, kya) + return kxa, kya, k, phi + + def _upsampled_sampling(self, upsampling_factor) -> torch.Tensor: + return torch.as_tensor( + [s / upsampling_factor for s in self.scan_sampling], + device=self.device, + dtype=self._float_dtype, + ) + + def _return_defocus_rate_px(self, rotation_angle, bf_mask, upsampling_factor): + """``(num_bf, 2)`` change in lateral shift per Angstrom of defocus, in canvas pixels. + + ``chi`` is linear in every aberration magnitude, so the rate is exactly the shift of + a unit ``C10`` and is *independent of all other aberrations* -- hence the absent + ``aberration_coefs`` argument. Evaluating it through + ``aberration_surface_cartesian_gradients`` rather than hand-coding the analytic + ``wavelength * k`` keeps it tied to the same expression the Fourier class uses. + """ + self._require_analytic_probe("A defocus gradient") + _, _, k, phi = self._return_k_grid(rotation_angle) + dx, dy = aberration_surface_cartesian_gradients( + k * self.wavelength, phi, aberration_coefs={"C10": 1.0} + ) + rate = torch.stack((dx[bf_mask], dy[bf_mask]), -1) + return ( + rate.to(self._float_dtype) + / (2 * math.pi) + / self._upsampled_sampling(upsampling_factor) + ) + + def _return_delta_c10(self, defocus_gradient) -> torch.Tensor | None: + """``(N_pos,)`` local defocus offset in Angstrom, or ``None`` for no gradient. + + The positions are taken in the unrotated scan frame: defocus varies with physical + position on the specimen, whereas the detector rotation belongs to the k-grid in + :meth:`_return_defocus_rate_px`. Rotating here as well would double-count it. + """ + if defocus_gradient is None or (defocus_gradient[0] == 0.0 and defocus_gradient[1] == 0.0): + return None + self._require_analytic_probe("A defocus gradient") + + scan_sampling = torch.as_tensor( + tuple(self.scan_sampling), device=self.device, dtype=self._float_dtype + ) + offsets_ang = (self._positions_px - self.positions_centroid_px) * scan_sampling + gradient = torch.as_tensor(defocus_gradient, device=self.device, dtype=self._float_dtype) + return offsets_ang @ gradient + + def _return_shifts_px(self, rotation_angle, aberration_coefs, bf_mask, upsampling_factor): + """``(num_bf, 2)`` parallax shifts in upsampled canvas pixels, plus the BF weight.""" + _, _, k, phi = self._return_k_grid(rotation_angle) + + dx, dy = aberration_surface_cartesian_gradients( + k * self.wavelength, + phi, + aberration_coefs=aberration_coefs, + ) + grad_k = torch.stack((dx[bf_mask], dy[bf_mask]), -1) + + upsampled_sampling = self._upsampled_sampling(upsampling_factor) + shifts_px = grad_k.to(self._float_dtype) / (2 * math.pi) / upsampled_sampling + + # matches DirectPtychography.reconstruct: soft_edges is left at evaluate_probe's + # default rather than taking self.soft_edges, so the two normalizations agree + cmplx_probe_k = self._return_probe_on_grid(k, phi, aberration_coefs) + bf_weights = cmplx_probe_k[bf_mask].abs().square().sum() + + return shifts_px, bf_weights + + @staticmethod + def _return_shift_extrema(shifts_px, defocus_rate_px, delta_c10): + """``(lo, hi)`` over every ``(bright-field pixel, scan position)`` shift. + + The total shift is ``shifts[m] + rate[m] * delta_c10[n]``, which is monotone in + ``delta_c10``, so its extrema over ``n`` are attained at the extremes of + ``delta_c10`` and there is no need to materialize the ``(num_bf, N_pos, 2)`` array. + """ + if delta_c10 is None: + return shifts_px.amin(0), shifts_px.amax(0) + + at_min = shifts_px + defocus_rate_px * delta_c10.min() + at_max = shifts_px + defocus_rate_px * delta_c10.max() + return ( + torch.minimum(at_min, at_max).amin(0), + torch.maximum(at_min, at_max).amax(0), + ) + + def _return_kernel_fourier( + self, batch_idx, bf, kernel, qxa, qya, kxa, kya, cmplx_probe_k, probe, norm + ): + """``(B, Ny, Nx)`` Fourier deconvolution kernel for a batch of bright-field pixels. + + Mirrors ``DirectPtychography._return_kernel_contributions`` term for term, minus the + data, so the two classes cannot drift apart. + """ + ind_i = bf.bf_inds_i[batch_idx] + ind_j = bf.bf_inds_j[batch_idx] + kx = kxa[ind_i, ind_j].view(-1, 1, 1) + ky = kya[ind_i, ind_j].view(-1, 1, 1) + + if kernel == "icom": + # `k . q / |q|**2`, which never reads the probe -- so it needs no overlap + # function, and an empirical probe raises no sampling question here + q_square = qxa.square() + qya.square() + qx_op = -1.0j * qxa / q_square + qy_op = -1.0j * qya / q_square + qx_op[0, 0] = 0.0 + qy_op[0, 0] = 0.0 + return kx * qx_op.unsqueeze(0) + ky * qy_op.unsqueeze(0), None + + gamma = gamma_factor( + (qxa.unsqueeze(0) - kx, qya.unsqueeze(0) - ky), + (qxa.unsqueeze(0) + kx, qya.unsqueeze(0) + ky), + cmplx_probe_k[ind_i, ind_j].view(-1, 1, 1), + probe, + normalize=False, + ) + + if kernel == "ssb": + return -1.0j * gamma.conj() / gamma.abs().clip(1e-8), gamma + return -1.0j * gamma.conj() / (1.0 if norm is None else norm), gamma + + @staticmethod + def _resolve_convolution_mode(convolution_mode, kernel, stencil_radius): + """Which convolution route to take for a non-parallax kernel. + + ``"auto"`` reads ``stencil_radius``: naming one is a request to truncate, so it takes + the stencil; leaving it at ``"auto"`` takes the exact FFT. That is cheap to decide, + where actually measuring which is faster would cost a full pass over the kernels. + """ + if kernel == "prlx": + return "splat" + if convolution_mode not in ("auto", "fft", "stencil"): + raise ValueError( + f"`convolution_mode` must be 'auto', 'fft' or 'stencil', got {convolution_mode!r}" + ) + if convolution_mode != "auto": + return convolution_mode + return "stencil" if stencil_radius != "auto" else "fft" + + def _probe_rotation_is_exact(self, rotation_angle) -> bool: + """Whether a rotation leaves `k -/+ q` on the probe's own reciprocal lattice. + + `_return_k_grid` rotates the k-grid into the scan frame, which is what the analytic + aperture wants -- it is *evaluated* there. An array can only be *read* on its lattice, + and while `k` itself stays put, the offset `q` arrives rotated. Only rotations that + map the lattice onto itself keep it there; everything else has to be interpolated, + which `_resample_probe` handles by refining the grid first. + """ + turns = float(rotation_angle) / 90.0 + square = abs(self.reciprocal_sampling[0] - self.reciprocal_sampling[1]) < 1e-9 + return abs(turns - round(turns)) < 1e-6 and (round(turns) % 2 == 0 or square) + + def _resample_probe(self, probe, q_step, oversample=1): + """`probe` on the canvas's reciprocal grid, cached across reconstructions. + + Two transforms over a large detector are not free, and the canvas rarely changes + between calls -- a defocus sweep or a hyperparameter search repeats the same one. + + ``oversample`` refines further, which is what makes an off-lattice sampling accurate: + bilinear error falls as the square of the refinement. Measured on a speckled X-ray + probe at a generic sub-pixel offset: 21% unrefined, 7.5% at 2x, 1.7% at 4x, 0.52% at + 8x, 0.064% at 16x -- against a memory cost that grows as the square. + """ + key = (q_step, int(oversample)) + cached = getattr(self, "_resampled_probe", None) + if cached is not None and cached[0] == key and cached[1] is probe: + return cached[2] + refined = probe.resampled_to(tuple(q / oversample for q in q_step)) + self._resampled_probe = (key, probe, refined) + return refined + + def _return_com_shift(self, bf, rotation_angle, upsampling_factor): + """``(2, N_pos)`` centre-of-mass shift, in upsampled canvas pixels. + + The k-weighted first moment of each diffraction pattern, summed over the detector. + This is the collapse that makes riCOM cheap: the iCoM kernel is linear in ``k``, so + + sum_m FFT(V_m) K_m(q) = A(q) FFT(splat(com_x)) + B(q) FFT(splat(com_y)) + + with ``A``, ``B`` the two components of ``-i q / |q|**2``. Summing over the detector + *before* convolving turns ``num_bf`` transforms into two -- 167k into two on the + X-ray data this was written against -- and is what the riCOM paper does. + """ + kxa, kya, _, _ = self._return_k_grid(rotation_angle) + k_vectors = torch.stack((kxa[bf.bf_mask], kya[bf.bf_mask]), dim=-1).to(self._float_dtype) + values = self._vbf_stack[bf.vbf_index_mapping].to(self._float_dtype) + return torch.einsum("mn,md->dn", values, k_vectors) + + def _return_icom_operators(self, qxa, qya): + """``(2, Ny, Nx)`` complex ``-i q / |q|**2``, the two halves of the iCoM kernel.""" + q_square = qxa.square() + qya.square() + operators = torch.stack((-1.0j * qxa / q_square, -1.0j * qya / q_square), dim=0) + operators[:, 0, 0] = 0.0 + return operators + + def _truncate_icom_operators(self, operators, canvas_shape, stencil_radius, max_radius): + """``((2, S*S) weights, info)`` -- the riCOM box stencil, from the iCoM operators. + + The real-space kernel is ``r / (2 * pi * |r|**2)``, centred at the origin, so the box + is taken about the origin rather than about a parallax shift: iCoM carries no shift + to divide out. The radius is riCOM's ``(n - 1) / 2``, and it is a high-pass cutoff by + intent, not a truncation error -- so nothing is reported as one. + """ + n_rows, n_cols = canvas_shape + limit = max(1, min(n_rows, n_cols) // 2 - 1) + radius = limit if stencil_radius == "auto" else min(int(stencil_radius), limit) + radius = min(radius, max_radius) if stencil_radius == "auto" else radius + + kappa = torch.fft.fftshift(torch.fft.ifft2(operators), dim=(-2, -1)) + centre = (n_rows // 2, n_cols // 2) + window = ( + slice(centre[0] - radius, centre[0] + radius + 1), + slice(centre[1] - radius, centre[1] + radius + 1), + ) + weights = kappa[:, window[0], window[1]].reshape(2, -1) + + inside = weights.abs().square().sum() + total = kappa.abs().square().sum() + return weights, { + "stencil_radius": radius, + "mean_error": float((1 - inside / total).clamp_min(0).sqrt()), + "max_error": float((1 - inside / total).clamp_min(0).sqrt()), + } + + def _return_kernel_context( + self, + bf, + *, + kernel, + rotation_angle, + aberration_coefs, + canvas_shape, + upsampling_factor, + matched_filter_norm_epsilon, + kernel_batch_size, + probe_oversample=1, + ): + """``(kernel_args, norm, bf_weights)``, everything a Fourier kernel needs but the batch. + + Shared by the stencil and the FFT convolution paths, so the two cannot build + different kernels from the same settings. + """ + upsampled_sampling = tuple(s / upsampling_factor for s in self.scan_sampling) + qxa, qya = spatial_frequencies(canvas_shape, upsampled_sampling, device=self.device) + kxa, kya, k, phi = self._return_k_grid(rotation_angle) + + cmplx_probe_k = self._return_probe_on_grid(k, phi, aberration_coefs) + bf_weights = cmplx_probe_k[bf.bf_mask].abs().square().sum() + + probe = self._return_probe(aberration_coefs) + if probe.array is not None: + # an empirical probe can only be read on its own reciprocal grid, so refine it + # onto the canvas's -- exactly, by zero-padding in real space + q_step = tuple(1 / (n * d) for n, d in zip(canvas_shape, upsampled_sampling)) + oversample = 1 + if not self._probe_rotation_is_exact(rotation_angle): + # the rotation carries q off the lattice, so the probe has to be interpolated + oversample = max(1, int(probe_oversample)) + probe = FourierProbe( + probe.wavelength, + array=probe.array, + reciprocal_sampling=probe.reciprocal_sampling, + interpolation="bilinear", + ) + if oversample < 8: + warnings.warn( + f"rotation_angle={rotation_angle} takes q off the probe's reciprocal " + f"lattice, so psi is interpolated. At probe_oversample={oversample} " + "that costs roughly " + f"{ {1: '20%', 2: '7%', 4: '2%'}.get(oversample, '<1%') } rms on a " + "speckled probe; raise it (memory grows as its square) or use a " + "rotation that maps the lattice onto itself.", + stacklevel=3, + ) + probe = self._resample_probe(probe, q_step, oversample) + + kernel_args = (bf, kernel, qxa, qya, kxa, kya, cmplx_probe_k, probe) + + # obf and mf normalize by a power spectrum summed over every bright-field pixel, so + # they need a pass over all of them before any kernel is final + norm = None + if kernel in ("obf", "mf"): # icom needs no normalization pass + power = torch.zeros(canvas_shape, device=self.device) + batcher = SimpleBatcher( + bf.num_bf, batch_size=kernel_batch_size, shuffle=False, rng=self.rng + ) + for batch_idx in batcher: + _, gamma = self._return_kernel_fourier(batch_idx, *kernel_args, None) + power += gamma.abs().square().sum(0) + power /= bf_weights + if kernel == "obf": + norm = power.sqrt().clamp_min(1e-8) + else: + norm = (power + matched_filter_norm_epsilon * power.max()).clamp_min(1e-8) + + return kernel_args, norm, bf_weights + + def _return_kernel_stencil( + self, + bf, + *, + kernel, + rotation_angle, + aberration_coefs, + canvas_shape, + upsampling_factor, + shift_centers, + stencil_radius, + truncation_tolerance, + max_stencil_radius, + matched_filter_norm_epsilon, + kernel_batch_size, + verbose, + probe_oversample=1, + ): + """``(stencil_offsets, stencil_weights, bf_weights, info)`` for a convolution kernel. + + Builds each bright-field pixel's Fourier kernel on the canvas grid, transforms it to + real space and truncates it to a box stencil. Costs ``num_bf`` canvas FFTs once per + reconstruction, negligible beside the scatter that follows. + + ``shift_centers`` is the integer parallax shift, which is divided out of the kernel + by a phase ramp and added back to the deposit coordinates. Without it the stencil + would have to span the shift itself -- tens of pixels at realistic defocus -- rather + than just the residual chirp and aperture ringing. + """ + kernel_args, norm, bf_weights = self._return_kernel_context( + bf, + kernel=kernel, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + canvas_shape=canvas_shape, + upsampling_factor=upsampling_factor, + matched_filter_norm_epsilon=matched_filter_norm_epsilon, + kernel_batch_size=kernel_batch_size, + probe_oversample=probe_oversample, + ) + + n_rows, n_cols = canvas_shape + radius_limit = max(1, min(n_rows, n_cols) // 2 - 1) + + def batches(): + return SimpleBatcher( + bf.num_bf, batch_size=kernel_batch_size, shuffle=False, rng=self.rng + ) + + # exp(2i.pi.c.m/N) rolls kappa by -c, undoing the parallax shift + freq_row = torch.fft.fftfreq(n_rows, device=self.device).view(-1, 1) + freq_col = torch.fft.fftfreq(n_cols, device=self.device).view(1, -1) + + def recentered_kappa(batch_idx): + kernel_fourier, _ = self._return_kernel_fourier(batch_idx, *kernel_args, norm) + centers = shift_centers[batch_idx].to(torch.float32) + ramp = torch.exp( + 2j + * math.pi + * (centers[:, 0, None, None] * freq_row + centers[:, 1, None, None] * freq_col) + ) + return torch.fft.fftshift(torch.fft.ifft2(kernel_fourier * ramp), dim=(-2, -1)) + + center = (n_rows // 2, n_cols // 2) + rows = torch.arange(n_rows, device=self.device).view(-1, 1) - center[0] + cols = torch.arange(n_cols, device=self.device).view(1, -1) - center[1] + chebyshev = torch.maximum(rows.abs(), cols.abs()) + + # first pass: cumulative energy inside each candidate box, per bright-field pixel. + # only these scalars are kept, so the full-canvas kernels never all exist at once + candidates = list(range(1, min(max_stencil_radius, radius_limit) + 1)) + if stencil_radius != "auto": + candidates = [min(int(stencil_radius), radius_limit)] + + inside = torch.zeros((bf.num_bf, len(candidates)), device=self.device) + total = torch.zeros(bf.num_bf, device=self.device) + for batch_idx in tqdm(list(batches()), disable=not verbose, desc=f"{kernel} kernel"): + energy = recentered_kappa(batch_idx).abs().square() + total[batch_idx] = energy.sum((-2, -1)) + for column, candidate in enumerate(candidates): + inside[batch_idx, column] = (energy * (chebyshev <= candidate)).sum((-2, -1)) + + errors = 1 - inside / total.clamp_min(torch.finfo(total.dtype).tiny)[:, None] + errors = errors.clamp_min(0).sqrt() + + mean_errors = errors.mean(0) + if stencil_radius == "auto": + meets = (mean_errors <= truncation_tolerance).nonzero() + column = int(meets[0]) if meets.numel() else len(candidates) - 1 + else: + column = 0 + radius = candidates[column] + + info = { + "stencil_radius": radius, + "mean_error": float(mean_errors[column]), + "max_error": float(errors[:, column].max()), + } + + if info["mean_error"] > truncation_tolerance and kernel != "icom": + warnings.warn( + f"A stencil radius of {radius} px leaves an estimated " + f"{info['mean_error']:.0%} truncation error (worst bright-field pixel " + f"{info['max_error']:.0%}). The {kernel.upper()} kernel is not compact in " + "real space -- dividing by |gamma| leaves a phase on a hard-edged support, " + "whose transform has tails decaying as r**-1.5, so the error falls only like " + "1/radius and being in focus does not help. DirectPtychography computes the " + "same kernel exactly by FFT; prefer it unless the scan is ungridded. This " + "estimate assumes a white object spectrum, so it is pessimistic for a real " + "one.", + stacklevel=3, + ) + # iCoM is excluded above: truncating it is riCOM, not an approximation of iCoM. + + # second pass: crop to the chosen box + window = ( + slice(center[0] - radius, center[0] + radius + 1), + slice(center[1] - radius, center[1] + radius + 1), + ) + side = 2 * radius + 1 + stencil_weights = torch.empty( + (bf.num_bf, side * side), device=self.device, dtype=torch.complex64 + ) + for batch_idx in batches(): + stencil_weights[batch_idx] = recentered_kappa(batch_idx)[ + :, window[0], window[1] + ].reshape(len(batch_idx), -1) + + span = torch.arange(-radius, radius + 1, device=self.device) + offsets = torch.stack(torch.meshgrid(span, span, indexing="ij"), dim=-1).reshape(-1, 2) + + return offsets, stencil_weights, bf_weights, info + + def _return_canvas( + self, + shifts_px, + upsampling_factor, + boundary, + pad_px, + defocus_rate_px=None, + delta_c10=None, + obj_origin=None, + obj_fov=None, + ): + """``(canvas_shape, canvas_origin_px, canvas_fov)`` for the requested boundary. + + The field of view is returned alongside the shape, rather than recomputed later, + so the two cannot disagree about the upsampling factor. + """ + positions_up = self._positions_px * upsampling_factor + shift_lo, shift_hi = self._return_shift_extrema(shifts_px, defocus_rate_px, delta_c10) + + def with_fov(canvas_shape, origin): + canvas_fov = tuple( + n * s / upsampling_factor for n, s in zip(canvas_shape, self.scan_sampling) + ) + return canvas_shape, origin, canvas_fov + + if boundary == "wrap" and obj_origin is None and obj_fov is None: + # spans exactly the scan field of view, at any upsampling factor + canvas_shape = tuple(int(n) * upsampling_factor for n in self.scan_gpts) + origin = torch.zeros(2, device=self.device, dtype=self._float_dtype) + return with_fov(canvas_shape, origin) + + if boundary not in ("wrap", "pad"): + raise ValueError(f"`boundary` must be 'wrap' or 'pad', got {boundary!r}") + + if pad_px is not None and obj_fov is not None: + raise ValueError("`pad_px` and `obj_fov` both size the canvas; pass one or the other.") + + if pad_px is None: + lo = torch.floor(_snap_to_integer(positions_up.amin(0) + shift_lo)) + hi = torch.ceil(_snap_to_integer(positions_up.amax(0) + shift_hi)) + else: + lo = torch.floor(_snap_to_integer(positions_up.amin(0))) - pad_px + hi = torch.ceil(_snap_to_integer(positions_up.amax(0))) + pad_px + + # +2 leaves room for the upper bilinear corner at the far edge + canvas_shape = tuple(int(v) + 2 for v in (hi - lo)) + + if obj_origin is not None: + # Angstrom in the caller's frame -> upsampled canvas pixels, deliberately *not* + # rounded: each frame anchors its grid at its own bounding box, so snapping + # would leave the same window a fraction of a pixel apart per frame -- the very + # misregistration this removes. Fractional coordinates are fine for the splat. + offset = np.asarray(obj_origin, dtype=np.float64) - np.asarray(self.scan_origin) + lo_np = offset / np.asarray(self.scan_sampling) * upsampling_factor + lo = torch.as_tensor(lo_np, device=self.device, dtype=self._float_dtype) + + if obj_fov is not None: + fov = np.asarray(obj_fov, dtype=np.float64) + if fov.size != 2 or np.any(fov <= 0): + raise ValueError(f"`obj_fov` must be a positive (row, col) pair, got {obj_fov!r}") + canvas_shape = tuple( + max(1, int(round(f / s * upsampling_factor))) + for f, s in zip(fov, self.scan_sampling) + ) + + if boundary == "wrap": + self._warn_if_positions_wrap(positions_up, lo, canvas_shape, upsampling_factor) + + return with_fov(canvas_shape, lo) + + def _warn_if_positions_wrap(self, positions_up, lo, canvas_shape, upsampling_factor): + """Warn when ``"wrap"`` folds part of the scan back over the object. + + Wrapping a canvas that spans the whole scan is harmless, and is what makes the + montage reproduce ``DirectPtychography``, which is periodic in the scan. Wrapping a + smaller one is not: the positions outside come back on the opposite side and lay a + second, offset copy of the specimen over the first. + """ + shape = torch.as_tensor(canvas_shape, device=positions_up.device, dtype=lo.dtype) + below = (lo - positions_up.amin(0)).clamp_min(0) + above = (positions_up.amax(0) - (lo + shape)).clamp_min(0) + outside = torch.maximum(below, above) + if not bool((outside > 0.5).any()): + return + + overhang = to_numpy(outside) / upsampling_factor * np.asarray(self.scan_sampling) + warnings.warn( + f"boundary='wrap' with a canvas of {tuple(canvas_shape)} px leaves scan positions " + f"up to {np.round(overhang, 1)} Angstrom outside it, which wrap around and lay a " + "second copy of the specimen over the reconstruction. Either widen `obj_fov` to " + "cover the scan -- for an empirical probe, to the next whole multiple of the " + "probe's field of view -- or use boundary='pad', which drops them instead.", + stacklevel=4, + ) + + def reconstruct( + self, + bf_mask=None, + override_aberration_coefs=None, + upsampling_factor=None, + override_rotation_angle=None, + max_batch_size=None, + deconvolution_kernel="parallax", + q_highpass=None, + q_lowpass=None, + butterworth_order=12, + parallax_flip_phase=True, + verbose=None, + use_initial_state=False, + boundary=None, + defocus_gradient=None, + interpolation="nearest", + weight_normalize=None, + weight_threshold=1e-2, + pad_px=None, + obj_origin=None, + obj_fov=None, + compute_variance=True, + suppress_nyquist=False, + convolution_mode="auto", + probe_oversample=8, + stencil_radius="auto", + truncation_tolerance=0.1, + max_stencil_radius=32, + matched_filter_norm_epsilon=1e-1, + kernel_batch_size=16, + ): + """ + Accumulate the canvas and apply the post-hoc Fourier filters. + + Parameters + ---------- + bf_mask : torch.Tensor, optional + Subset of the bright-field mask to use. Must be strictly smaller than the mask + used at initialization. + override_aberration_coefs : dict, optional + Aberration coefficients, overriding the hyperparameter state. + upsampling_factor : int, optional + Integer factor by which to refine the canvas relative to the scan sampling. + override_rotation_angle : float, optional + Detector rotation in degrees, overriding the hyperparameter state. + max_batch_size : int, optional + Number of bright-field pixels splatted at once. Defaults to a memory-bounded + chunk of roughly four million ``(BF pixel, scan position)`` points. + deconvolution_kernel : str + ``"prlx"`` (and its aliases) is a pure translation and is exact. + + ``"ssb"``, ``"obf"``, ``"mf"`` and ``"icom"`` are convolutions. They are exact + by default, evaluated by FFT on the canvas; ``convolution_mode="stencil"`` + truncates them to a box instead. + + Truncating ``"icom"`` is the exception that is not an approximation: it gives + riCOM, whose real-space kernel is ``r / (2 * pi * |r|**2)`` and whose radius is a + high-pass cutoff, so no truncation warning is raised for it. Being linear in + ``k``, it also collapses the detector sum into two convolutions of the + centre-of-mass shift. + q_highpass, q_lowpass : float, optional + Butterworth filter cutoffs, applied once to the finished image. + parallax_flip_phase : bool + Apply the ``sign(sin(chi(q)))`` phase-flip filter. + boundary : {"wrap", "pad"}, optional + ``"wrap"`` wraps the montage periodically over the scan grid and reproduces + ``DirectPtychography``; ``"pad"`` grows the canvas to cover the shifted positions + and drops nothing. Defaults to the value chosen at construction. + defocus_gradient : tuple of float, optional + ``(d C10 / d row, d C10 / d col)`` in Angstrom per Angstrom, for a tilted sample + whose defocus varies across the field of view. Defaults to + :attr:`defocus_gradient`; pass ``(0.0, 0.0)`` to disable it for one call. + + Each scan position is then shifted by its own local defocus, which a Fourier + formulation cannot express: ``exp(-1j * grad_chi . q)`` is global over the scan. + The post-hoc phase flip still uses the aplanatic ``C10``; a space-variant + contrast-transfer correction is not attempted. + interpolation : {"nearest", "bilinear"} + Sub-pixel deposition scheme. + + ``"nearest"`` (default) snaps each shift to the closest canvas pixel, a roll of + the bright-field image by ``round(shift)`` on a raster scan. The quantization + error is ``1/(2*upsampling_factor)`` scan pixels, so it shrinks as you upsample: + against the exact Fourier shift on a 4 mrad apoferritin dataset it retains + 0.81 / 0.95 / 0.99 / 1.00 of the in-band power at ``upsampling_factor`` + 1 / 2 / 4 / 8. + + ``"bilinear"`` spreads each shift over the four neighbouring pixels, trading + some smoothing (0.67 / 0.90 / 0.97 / 0.99 over the same series) for coverage. + Prefer it for positions off a lattice, where snapping leaves parts of the canvas + unvisited -- 17-31% empty against 10-15% on a jittered scan. + weight_normalize : bool, optional + Divide by the accumulated weight rather than by the total bright-field weight. + Defaults to ``True`` for an ungridded scan and ``False`` otherwise. + + This corrects for uneven sampling density, which an ungridded scan needs. On a + raster scan the density is already uniform, so it only rescales the edges of a + padded canvas, amplifying the noise of the few contributions there; leaving it + ``False`` lets those edges fade out instead. + + For ``"wrap"`` with ``upsampling_factor > 1`` the accumulated weight is a comb of + ones and zeros, so normalizing by it is meaningless -- leave it ``False``. + weight_threshold : float + Fraction of the peak weight below which the normalized image is tapered to zero, + following ``bilinear_kde``. Only used when ``weight_normalize`` is true. + pad_px : int, optional + Freeze the ``"pad"`` canvas to the position bounding box plus this many pixels, + instead of sizing it from the (aberration-dependent) shifts. Use this to keep the + canvas a fixed size across hyperparameter trials or a defocus series, where the + automatic size would otherwise change with the shifts. Contributions landing + beyond ``pad_px`` are dropped, so choose it larger than the shifts you expect. + obj_origin : tuple of float, optional + Pin the canvas corner to this ``(row, col)`` coordinate in Angstrom, in the same + frame as the probe positions passed to :meth:`from_dataset3d`. Defaults to + whatever ``pad_px`` or the shifts imply, which follows each acquisition's own + bounding box and so differs between them. + obj_fov : tuple of float, optional + Pin the canvas extent to ``(rows, cols)`` Angstrom, rather than sizing it from + the positions. Mutually exclusive with ``pad_px``. + + Together these name a fixed window in the specimen's coordinates, so separate + acquisitions -- successive frames of a multi-frame scan, say -- reconstruct onto + pixel-identical canvases that can be stacked, differenced or cross-correlated. + Without them each frame's canvas follows its own bounding box, leaving two frames + of the same region a few pixels apart. Both are read back from + :attr:`obj_origin` and :attr:`_obj_sampling`. + + Neither applies to ``boundary="wrap"``, whose canvas is the scan grid. + compute_variance : bool + Accumulate the sum of squares needed by :meth:`variance_loss`. + suppress_nyquist : bool + Zero the Nyquist row and column of the phase-flip filter. Off by default, to + match ``DirectPtychography``; turn it on for odd-order aberrations, where + ``sign(sin(chi))`` is not symmetric and leaves a checkerboard artifact. + convolution_mode : {"auto", "fft", "stencil"} + How the ``ssb`` / ``obf`` / ``mf`` convolutions are evaluated. + + ``"fft"`` splats each bright-field image onto the canvas and multiplies by the + kernel in ``q``. Nothing is truncated, and it is asymptotically cheaper for a + kernel that spans the canvas: one transform per bright-field image against + ``(2 * stencil_radius + 1) ** 2`` taps per scan point. On a gridded scan it + reproduces ``DirectPtychography`` to float precision, where a radius-5 stencil is + 20-34% off and a radius-12 one still 2-5% off. + + ``"stencil"`` keeps the truncated box, evaluated as a grouped convolution. + Worth it when the kernel is local, where it avoids a canvas-sized transform. + + ``"auto"`` reads ``stencil_radius``: naming one takes the stencil, leaving it at + ``"auto"`` takes the FFT. Measuring which is faster would cost a full pass over + the kernels, so this reads intent rather than benchmarking. + + With ``boundary="pad"`` the FFT route doubles the canvas and crops back, since a + Fourier convolution is otherwise circular. That costs four times the transform + area; ``boundary="wrap"`` wants the circular one anyway and pays nothing. + probe_oversample : int + How finely an empirical ``fourier_probe`` is refined before being sampled off its + own reciprocal lattice, which a ``rotation_angle`` that is not a multiple of 90 + degrees forces. Both the accuracy and the memory cost scale as its square: + measured on a speckled X-ray probe at a generic sub-pixel offset, 21% unrefined, + 7.5% at 2, 1.7% at 4, 0.52% at 8, 0.064% at 16. Ignored for an analytic probe, + which is evaluated rather than sampled, and for a rotation that maps the lattice + onto itself. + stencil_radius : int or "auto" + Half-width of the box stencil used by the ``ssb`` / ``obf`` / ``mf`` / ``icom`` + kernels, in canvas pixels. ``"auto"`` grows it until the estimated truncation + error meets ``truncation_tolerance``, capped at ``max_stencil_radius``. For + ``icom`` this is riCOM's ``(n - 1) / 2``, where setting it is the intent rather + than a compromise. + + The box carries no taper: tapering measures worse at equal radius (0.40 against + 0.29 relative error at radius 8, 20 mrad in focus), since it discards mid-radius + content that matters more than the ringing it suppresses. + truncation_tolerance : float + Target relative operator error for ``stencil_radius="auto"``. A warning reports + the achieved error whenever it cannot be met. + max_stencil_radius : int + Cap on the automatic radius. Cost scales with its square. + matched_filter_norm_epsilon : float + Regularization of the ``mf`` power normalization, as in ``DirectPtychography``. + kernel_batch_size : int + Bright-field pixels whose real-space kernels are built at once. + + Returns + ------- + self + + Notes + ----- + The convolution kernels cost more than the parallax one either way: an FFT per + bright-field image, or ``(2 * stencil_radius + 1) ** 2`` deposits per point. On a + gridded scan ``DirectPtychography`` computes the same thing with a single scan FFT + and is faster; see the class docstring for when to prefer which. + """ + state = self.hyperparameter_state + + if verbose is None: + verbose = self.verbose + + if use_initial_state: + if verbose: + print("Reconstructing with:\n\n", state.summarize(which="initial")) + aberration_coefs = state.initial_aberrations + rotation_angle = state.initial_rotation_angle + else: + if verbose: + print( + "Reconstructing with:\n\n", + state.summarize( + which="current", + override_aberration_coefs=override_aberration_coefs, + override_rotation_angle=override_rotation_angle, + ), + ) + aberration_coefs = state.current_aberrations(override_aberration_coefs) + rotation_angle = state.current_rotation_angle(override_rotation_angle) + + kernel = self._normalize_kernel_name(deconvolution_kernel) + if kernel == "prlx": + # zero aberrations would give zero shifts and quietly sum the bright-field stack + # into a plain incoherent image, which is not a parallax reconstruction + self._require_analytic_probe("The parallax kernel") + if upsampling_factor is None: + upsampling_factor = 1 + upsampling_factor = math.ceil(upsampling_factor) + + if bf_mask is None: + bf_mask = self.bf_mask + bf = self._return_bf_context(bf_mask) + + if boundary is None: + boundary = self.boundary + if defocus_gradient is None: + defocus_gradient = self.defocus_gradient + if verbose and defocus_gradient is not None: + print(f" defocus_gradient={tuple(defocus_gradient)!r} A/A,") + if weight_normalize is None: + # density correction matters for an ungridded scan; on a raster it would only + # amplify noise at the low-weight edges of a padded canvas + weight_normalize = not self.gridded_scan + + shifts_px, bf_weights = self._return_shifts_px( + rotation_angle, aberration_coefs, bf.bf_mask, upsampling_factor + ) + + delta_c10 = self._return_delta_c10(defocus_gradient) + if delta_c10 is None: + defocus_rate_px = None + else: + defocus_rate_px = self._return_defocus_rate_px( + rotation_angle, bf.bf_mask, upsampling_factor + ) + + canvas_shape, canvas_origin, canvas_fov = self._return_canvas( + shifts_px, + upsampling_factor, + boundary, + pad_px, + defocus_rate_px, + delta_c10, + obj_origin=obj_origin, + obj_fov=obj_fov, + ) + + if max_batch_size is None: + max_batch_size = max(1, _DEFAULT_POINTS_PER_BATCH // max(self.num_positions, 1)) + + coords_base = self._positions_px * upsampling_factor - canvas_origin + self._reset_reconstruction() + self._kernel = kernel + + mode = self._resolve_convolution_mode(convolution_mode, kernel, stencil_radius) + stencil_offsets = stencil_weights = kernel_args = norm = None + fft_shape = canvas_shape + + # riCOM: the iCoM kernel is linear in k, so summing the detector first turns the + # whole reconstruction into two convolutions of the centre-of-mass shift. Needs + # every bright-field pixel to deposit at the same place, which a per-position + # defocus breaks. + collapse_icom = kernel == "icom" and delta_c10 is None + + if kernel == "prlx": + deposit_shifts = shifts_px + self._stencil_info = None + elif collapse_icom: + deposit_shifts = torch.zeros_like(shifts_px) + self._stencil_info = None + if boundary == "pad" and mode == "fft": + fft_shape = (canvas_shape[0] * 2, canvas_shape[1] * 2) + upsampled_sampling = tuple(s / upsampling_factor for s in self.scan_sampling) + qxa, qya = spatial_frequencies( + fft_shape if mode == "fft" else canvas_shape, + upsampled_sampling, + device=self.device, + ) + icom_operators = self._return_icom_operators(qxa, qya) + icom_values = self._return_com_shift(bf, rotation_angle, upsampling_factor) + _, _, bf_weights = self._return_kernel_context( + bf, + kernel=kernel, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + canvas_shape=canvas_shape, + upsampling_factor=upsampling_factor, + matched_filter_norm_epsilon=matched_filter_norm_epsilon, + kernel_batch_size=kernel_batch_size, + probe_oversample=probe_oversample, + ) + if mode == "stencil": + icom_operators, self._stencil_info = self._truncate_icom_operators( + icom_operators, canvas_shape, stencil_radius, max_stencil_radius + ) + elif mode == "fft": + # nothing is truncated, so there is no reason to divide the shift out of the + # kernel and add it back to the deposits + deposit_shifts = torch.zeros_like(shifts_px) + self._stencil_info = None + if boundary == "pad": + # a Fourier convolution is circular; doubling the canvas and cropping back + # makes it linear, which is what "pad" asks for + fft_shape = (canvas_shape[0] * 2, canvas_shape[1] * 2) + kernel_args, norm, bf_weights = self._return_kernel_context( + bf, + kernel=kernel, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + canvas_shape=fft_shape, + upsampling_factor=upsampling_factor, + matched_filter_norm_epsilon=matched_filter_norm_epsilon, + kernel_batch_size=kernel_batch_size, + probe_oversample=probe_oversample, + ) + # each bright-field pixel now needs a canvas of its own, so size the batch by + # canvas area rather than by scan positions + max_batch_size = max(1, _DEFAULT_POINTS_PER_BATCH // max(np.prod(fft_shape), 1)) + else: + # the kernel carries the parallax shift in its phase; deposit at its integer part + # and leave only the residual in the stencil + deposit_shifts = shifts_px.round() + stencil_offsets, stencil_weights, bf_weights, self._stencil_info = ( + self._return_kernel_stencil( + bf, + kernel=kernel, + rotation_angle=rotation_angle, + aberration_coefs=aberration_coefs, + canvas_shape=canvas_shape, + upsampling_factor=upsampling_factor, + shift_centers=deposit_shifts, + stencil_radius=stencil_radius, + truncation_tolerance=truncation_tolerance, + max_stencil_radius=max_stencil_radius, + matched_filter_norm_epsilon=matched_filter_norm_epsilon, + kernel_batch_size=kernel_batch_size, + verbose=verbose, + probe_oversample=probe_oversample, + ) + ) + # a stencil of S taps costs S deposits per point, so shrink the batch to match + max_batch_size = max(1, max_batch_size // max(len(stencil_offsets), 1)) + + buffers = ( + allocate_splat_buffers(canvas_shape, self.device, accumulate_squares=compute_variance) + if kernel == "prlx" + else None + ) + accumulator = ( + None + if kernel == "prlx" + else torch.zeros(fft_shape, device=self.device, dtype=torch.complex64) + ) + + n_components = 0 if collapse_icom else bf.num_bf + pbar = tqdm(range(n_components), disable=not verbose) + batcher = SimpleBatcher( + n_components, batch_size=max_batch_size, shuffle=False, rng=self.rng + ) + + for batch_idx in batcher: + mapped_idx = bf.vbf_index_mapping[batch_idx] + values = self._vbf_stack[mapped_idx] # (B, N_pos) + coords = coords_base[None] + deposit_shifts[batch_idx][:, None] # (B, N_pos, 2) + if delta_c10 is not None: + # each position gets its own defocus, hence its own shift; the shift is + # exactly linear in C10, so this is one broadcast add rather than a re-fit + coords = coords + defocus_rate_px[batch_idx][:, None, :] * delta_c10[None, :, None] + + if kernel == "prlx": + scatter_add_splat( + values, + coords, + canvas_shape, + boundary=boundary, + interpolation=interpolation, + out=buffers, + ) + elif collapse_icom: + pass # handled in one shot below, outside the bright-field loop + elif mode == "fft": + stack = splat_stack( + values, + coords, + fft_shape, + boundary="pad" if fft_shape != canvas_shape else boundary, + interpolation=interpolation, + ) + kernel_fourier, _ = self._return_kernel_fourier(batch_idx, *kernel_args, norm) + accumulator += convolve_stack_fourier(stack, kernel_fourier) + else: + accumulator += splat_and_convolve( + values, + coords, + canvas_shape, + stencil_weights[batch_idx], + self._stencil_info["stencil_radius"], + boundary=boundary, + interpolation=interpolation, + ).sum(0) + pbar.update(len(batch_idx)) + pbar.close() + + if collapse_icom: + coords = coords_base[None].expand(2, -1, -1) + if mode == "fft": + stack = splat_stack( + icom_values, + coords, + fft_shape, + boundary="pad" if fft_shape != canvas_shape else boundary, + interpolation=interpolation, + ) + accumulator = convolve_stack_fourier(stack, icom_operators) + else: + accumulator = splat_and_convolve( + icom_values, + coords, + canvas_shape, + icom_operators, + self._stencil_info["stencil_radius"], + boundary=boundary, + interpolation=interpolation, + ).sum(0) + + self._canvas_shape = canvas_shape + self._canvas_origin_px = canvas_origin + self._canvas_fov = canvas_fov + self._bf_weights = bf_weights + + if kernel != "prlx": + if mode == "fft": + accumulator = torch.fft.ifft2(accumulator) + # matches DirectPtychography, which takes the real part of the summed stack + obj = accumulator[: canvas_shape[0], : canvas_shape[1]].real / bf_weights + else: + self._sum_w, self._sum_wv, self._sum_wv2 = buffers + # normalization must precede filtering: dividing by the (spatially varying) + # weight map is not linear, so it does not commute with the Fourier filters below + if weight_normalize: + mean, _, support = self._weighted_moments(weight_threshold) + obj = (mean * support).reshape(canvas_shape) + else: + obj = self._sum_wv.reshape(canvas_shape) / bf_weights + + obj = self._apply_fourier_filters( + obj, + aberration_coefs=aberration_coefs, + upsampling_factor=upsampling_factor, + q_lowpass=q_lowpass, + q_highpass=q_highpass, + butterworth_order=butterworth_order, + # only `prlx` needs the phase flip -- the deconvolution kernels already invert + # the contrast transfer, and DirectPtychography draws the same line + parallax_flip_phase=parallax_flip_phase and kernel == "prlx", + suppress_nyquist=suppress_nyquist, + ) + self._corrected_bf = obj.to(torch.float32) + + # memory management + gc.collect() + torch.cuda.empty_cache() + if hasattr(torch, "mps") and torch.backends.mps.is_available(): + torch.mps.empty_cache() + gc.collect() + + return self + + def _apply_fourier_filters( + self, + obj, + *, + aberration_coefs, + upsampling_factor, + q_lowpass, + q_highpass, + butterworth_order, + parallax_flip_phase, + suppress_nyquist, + ): + """ + Apply the bright-field-index-independent filters once, on the summed image. + + ``DirectPtychography`` multiplies these into every bright-field image before its + inverse transform; because they do not depend on the bright-field index, doing it + once on the sum is exactly equivalent. + """ + if not (parallax_flip_phase or q_lowpass or q_highpass or suppress_nyquist): + return obj + + upsampled_sampling = tuple(s / upsampling_factor for s in self.scan_sampling) + qxa, qya = spatial_frequencies(obj.shape, upsampled_sampling, device=self.device) + q, theta = polar_coordinates(qxa, qya) + + # built at the grid's native precision, not the accumulator's: chi(q) reaches tens + # of radians, so sign(sin(chi)) is ill-conditioned at its zero crossings and float64 + # would flip a handful of pixels relative to DirectPtychography + filt = torch.ones_like(q) + if parallax_flip_phase: + chi_q = aberration_surface( + q * self.wavelength, + theta, + self.wavelength, + aberration_coefs=aberration_coefs, + ) + filt = filt * torch.sign(torch.sin(chi_q)) + if q_lowpass: + filt = filt / (1 + (q / q_lowpass) ** (2 * butterworth_order)) + if q_highpass: + filt = filt * (1 - 1 / (1 + (q / q_highpass) ** (2 * butterworth_order))) + if suppress_nyquist: + n_rows, n_cols = obj.shape + if n_rows % 2 == 0: + filt[n_rows // 2, :] = 0.0 + if n_cols % 2 == 0: + filt[:, n_cols // 2] = 0.0 + + return torch.fft.ifft2(torch.fft.fft2(obj) * filt.to(obj.dtype)).real + + @staticmethod + def _moments_from_buffers(sum_w, sum_wv, sum_wv2, weight_threshold: float = 1e-2): + """``(mean, variance, support)`` per canvas pixel, from raw splat accumulators.""" + w = sum_w + tiny = torch.finfo(w.dtype).tiny + inv_w = 1 / w.clamp_min(tiny) + + mean = sum_wv * inv_w + if sum_wv2 is None: + var = torch.zeros_like(mean) + else: + var = (sum_wv2 * inv_w - mean.square()).clamp_min(0) + + w_max = w.max() + if w_max <= 0: + support = torch.zeros_like(w) + else: + support = (w / (weight_threshold * w_max)).clamp(max=1.0) + + return mean * (w > 0), var * (w > 0), support + + @classmethod + def _variance_loss_from_buffers(cls, sum_w, sum_wv, sum_wv2): + """Weight-averaged per-pixel variance across bright-field images.""" + _, var, _ = cls._moments_from_buffers(sum_w, sum_wv, sum_wv2) + denom = sum_w.sum() + if denom <= 0: + return torch.tensor(torch.inf, dtype=sum_w.dtype, device=sum_w.device) + return (var * sum_w).sum() / denom + + def _weighted_moments(self, weight_threshold: float = 1e-2): + """``(mean, variance, support)`` per canvas pixel, as flat tensors.""" + if self._sum_w is None or self._sum_wv is None: + raise RuntimeError("Run reconstruct() before asking for the accumulated moments.") + return self._moments_from_buffers( + self._sum_w, self._sum_wv, self._sum_wv2, weight_threshold + ) + + def variance_loss(self): + """ + Weight-averaged variance across bright-field images, without storing the stack. + + Accumulating ``sum(w)``, ``sum(w*v)`` and ``sum(w*v**2)`` during the splat gives the + per-pixel population variance across bright-field images directly, so no + ``(N_bf, Ny, Nx)`` stack is needed. + + It differs from ``DirectPtychography.variance_loss`` in four documented ways: + + 1. It is a weight-averaged mean over pixels rather than an unweighted one. The two + coincide for ``boundary="wrap"`` on a complete grid with ``upsampling_factor=1``, + where the accumulated weight is exactly ``num_bf`` everywhere. + 2. It lacks the ``1 / bf_weights**2`` scale, being computed on the raw values. + 3. It is computed *before* the phase-flip and Butterworth filters, which are applied + post-hoc here but per bright-field image there. With ``parallax_flip_phase=True`` + this is a genuinely different objective. + 4. For ``upsampling_factor > 1`` it ignores unvisited canvas pixels instead of + counting them as zeros. + + With ``interpolation="bilinear"`` and non-integer shifts it also folds in the + within-interpolation spread, since ``sum(w*v**2)`` averages ``v**2`` over the + neighbouring source pixels. Prefer ``upsampling_factor=1`` and + ``parallax_flip_phase=False`` when driving a hyperparameter search. + """ + if self._kernel != "prlx": + raise NotImplementedError( + f"variance_loss is only defined for the parallax kernel, not {self._kernel!r}: " + "the convolution kernels deposit complex weights that are not a partition of " + "unity, so there is no per-pixel spread across bright-field images to take. " + "To drive a hyperparameter search on this kernel pass loss='rms_gradient', " + "which measures the reconstructed image instead and is defined for every " + "kernel." + ) + if self._sum_w is None or self._sum_wv2 is None: + return None + return self._variance_loss_from_buffers(self._sum_w, self._sum_wv, self._sum_wv2) + + # ------------------------------------------------------------------ + # position-dependent defocus + # ------------------------------------------------------------------ + + def _return_patch_indices(self, patch_grid) -> list[torch.Tensor]: + """Position indices for each tile of a ``patch_grid`` partition of the scan bbox.""" + pos = self._positions_px + lo = pos.amin(0) + span = (pos.amax(0) - lo).clamp_min(torch.finfo(pos.dtype).tiny) + + tile = torch.stack( + [ + ((pos[:, i] - lo[i]) / span[i] * patch_grid[i]).floor().clamp(0, patch_grid[i] - 1) + for i in range(2) + ], + dim=-1, + ).to(torch.int64) + + flat = tile[:, 0] * patch_grid[1] + tile[:, 1] + return [ + torch.nonzero(flat == p, as_tuple=True)[0] + for p in range(int(patch_grid[0]) * int(patch_grid[1])) + ] + + def _patch_canvas(self, pos_idx, margin_px, upsampling_factor): + """``(canvas_shape, origin)`` for a patch, sized independently of the trial defocus.""" + positions_up = self._positions_px[pos_idx] * upsampling_factor + lo = torch.floor(positions_up.amin(0) - margin_px) + hi = torch.ceil(positions_up.amax(0) + margin_px) + return tuple(int(v) + 2 for v in (hi - lo)), lo + + def _patch_variance_loss( + self, + pos_idx, + canvas_shape, + canvas_origin, + *, + bf, + rotation_angle, + aberration_coefs, + upsampling_factor, + interpolation, + max_batch_size, + ): + """Variance loss of a montage built from a spatial subset of the scan positions. + + Deliberately does not go through :meth:`reconstruct`: threading a position subset + through the public signature would make its canvas logic branch for a purely internal + use. It is otherwise the same functional as :meth:`variance_loss` -- the same + weight-averaged per-pixel spread over the same accumulators -- so the patch fit and + a hyperparameter search cannot disagree about what a good defocus is. + + The canvas is passed in frozen, rather than sized from the shifts, which is what + makes the value comparable *across trial defocus values*: a canvas that grew with + the defocus would add low-weight edge pixels, whose variance is small simply because + few images reach them, and the loss would then fall monotonically with defocus + rather than have a minimum at the right one. Weighting by ``sum_w`` is what keeps + those edges from dominating once the canvas is fixed. + """ + shifts_px, _ = self._return_shifts_px( + rotation_angle, aberration_coefs, bf.bf_mask, upsampling_factor + ) + + buffers = allocate_splat_buffers(canvas_shape, self.device, accumulate_squares=True) + coords_base = self._positions_px[pos_idx] * upsampling_factor - canvas_origin + + batcher = SimpleBatcher(bf.num_bf, batch_size=max_batch_size, shuffle=False, rng=self.rng) + for batch_idx in batcher: + values = self._vbf_stack[bf.vbf_index_mapping[batch_idx]][:, pos_idx] + coords = coords_base[None] + shifts_px[batch_idx][:, None] + scatter_add_splat( + values, + coords, + canvas_shape, + boundary="pad", + interpolation=interpolation, + out=buffers, + ) + + # Score the whole patch, not its densest spots. An ungridded weight map is uneven + # everywhere, so a 90%-of-peak cut kept 14% of the canvas and moved with the + # defocus -- putting the apoferritin minimum at C10 = 9.5 kA against a true 13.0 kA. + return float(self._variance_loss_from_buffers(*buffers)) + + @staticmethod + def _refine_minimum(values, losses): + """Sub-grid minimum by a 3-point parabolic fit, or ``None`` if pinned at an end.""" + i = int(np.argmin(losses)) + if i == 0 or i == len(losses) - 1: + return None + + y0, y1, y2 = losses[i - 1], losses[i], losses[i + 1] + denom = y0 - 2 * y1 + y2 + if denom <= 0: # not a minimum -- flat or concave + return float(values[i]) + + # vertex offset in units of the (possibly uneven) local step + offset = 0.5 * (y0 - y2) / denom + step = values[i + 1] - values[i] if offset > 0 else values[i] - values[i - 1] + return float(values[i] + offset * abs(step)) + + def defocus_map( + self, + c10_values, + patch_grid: Tuple[int, int] = (3, 3), + bf_mask=None, + override_aberration_coefs=None, + override_rotation_angle=None, + upsampling_factor: int = 1, + interpolation: str = "bilinear", + min_patch_positions: int = 64, + max_batch_size=None, + verbose=None, + ) -> dict: + """ + Best-fit defocus in each of a grid of spatial patches. + + Reconstructs each patch on its own small canvas over a range of trial ``C10`` values + and takes the variance-loss minimum, refined by a parabolic fit so a coarse + ``c10_values`` grid still resolves sub-step differences. This is the measurement + :meth:`fit_defocus_gradient` fits a plane to; call it directly to inspect the loss + curves before trusting the fit. + + Parameters + ---------- + c10_values : array-like + Trial defocus values in Angstrom, ascending. Must bracket the true local defocus + in every patch -- a patch whose minimum sits on an endpoint is flagged invalid. + patch_grid : tuple of int + Number of patches along each scan axis. Any positive pair; a 1D grid such as + ``(4, 1)`` gives a defocus profile along one axis. Fitting a *plane* needs at + least three patches -- :meth:`fit_defocus_gradient` enforces that. + interpolation : {"bilinear", "nearest"} + Defaults to ``"bilinear"`` here, unlike :meth:`reconstruct`. Snapping to the + nearest pixel makes the loss a staircase in ``C10``, so small defocus changes + produce no change at all and the minimum cannot be located. + min_patch_positions : int + Patches with fewer positions than this are flagged invalid. + + Returns + ------- + dict with keys + ``centers_A`` ``(P, 2)`` patch centers in Angstrom, measured from the position + centroid so they share the frame of :attr:`defocus_gradient`; ``c10_best`` + ``(P,)``; ``losses`` ``(P, n_c10)``; ``valid`` ``(P,)`` bool; ``c10_values``. + + Notes + ----- + The estimator carries a small offset that is uniform across patches -- a few percent + of ``C10`` on synthetic data. That cancels out of the *gradient*, which is a + difference between patches, but it does bias the absolute defocus, so treat the + offset from :meth:`fit_defocus_gradient` as approximate. + + Patches need enough positions for the loss to have a clear minimum, and how many is + data dependent -- the warning here only catches patches smaller than the shifts + themselves. The reliable check is to run this on a region you believe is flat and + confirm ``c10_best`` comes back constant. + """ + if verbose is None: + verbose = self.verbose + + c10_values = np.asarray(c10_values, dtype=np.float64).ravel() + if c10_values.size < 3: + raise ValueError( + f"`c10_values` needs at least 3 points to bracket a minimum, got " + f"{c10_values.size}." + ) + if min(int(patch_grid[0]), int(patch_grid[1])) < 1: + raise ValueError(f"`patch_grid` entries must be positive, got {patch_grid!r}.") + + state = self.hyperparameter_state + base_coefs = state.current_aberrations(override_aberration_coefs) + rotation_angle = state.current_rotation_angle(override_rotation_angle) + + bf = self._return_bf_context(self.bf_mask if bf_mask is None else bf_mask) + if max_batch_size is None: + max_batch_size = max(1, _DEFAULT_POINTS_PER_BATCH // max(self.num_positions, 1)) + + patches = self._return_patch_indices(patch_grid) + scan_sampling = torch.as_tensor( + tuple(self.scan_sampling), device=self.device, dtype=self._float_dtype + ) + + # size every patch canvas for the largest trial defocus, so it stays fixed across + # the scan below -- see _patch_variance_loss + max_shift = 0.0 + for c10 in (c10_values.min(), c10_values.max()): + trial, _ = self._return_shifts_px( + rotation_angle, {**base_coefs, "C10": float(c10)}, bf.bf_mask, upsampling_factor + ) + max_shift = max(max_shift, float(trial.abs().max())) + margin_px = math.ceil(max_shift) + 1 + + centers, best, all_losses, valid = [], [], [], [] + pbar = tqdm(patches, disable=not verbose, desc="defocus map") + for pos_idx in pbar: + center = ( + (self._positions_px[pos_idx].mean(0) - self.positions_centroid_px) * scan_sampling + if pos_idx.numel() + else torch.zeros(2, device=self.device, dtype=self._float_dtype) + ) + centers.append(center.cpu().numpy()) + + if pos_idx.numel() < min_patch_positions: + all_losses.append(np.full(c10_values.size, np.nan)) + best.append(np.nan) + valid.append(False) + continue + + extent = self._positions_px[pos_idx].amax(0) - self._positions_px[pos_idx].amin(0) + if float(extent.min()) < 2 * max_shift: + warnings.warn( + f"A patch spans {float(extent.min()):.0f} scan pixels but the parallax " + f"shifts reach {max_shift:.0f}, so its bright-field images barely " + "overlap and the variance loss has little to compare. Use a coarser " + "`patch_grid`.", + stacklevel=2, + ) + + canvas_shape, canvas_origin = self._patch_canvas(pos_idx, margin_px, upsampling_factor) + losses = np.array( + [ + self._patch_variance_loss( + pos_idx, + canvas_shape, + canvas_origin, + bf=bf, + rotation_angle=rotation_angle, + aberration_coefs={**base_coefs, "C10": float(c10)}, + upsampling_factor=upsampling_factor, + interpolation=interpolation, + max_batch_size=max_batch_size, + ) + for c10 in c10_values + ] + ) + all_losses.append(losses) + + refined = self._refine_minimum(c10_values, losses) + best.append(np.nan if refined is None else refined) + valid.append(refined is not None) + pbar.close() + + results = { + "centers_A": np.stack(centers), + "c10_best": np.array(best), + "losses": np.stack(all_losses), + "valid": np.array(valid), + "c10_values": c10_values, + } + self._defocus_map_results = results + return results + + def fit_defocus_gradient( + self, + c10_values, + patch_grid: Tuple[int, int] = (3, 3), + update_defocus: bool = True, + verbose=None, + **defocus_map_kwargs, + ): + """ + Fit a defocus plane across the field of view and store it as :attr:`defocus_gradient`. + + Runs :meth:`defocus_map` and least-squares fits ``C10 = offset + g . r`` through the + valid patch centers. Because the centers are measured from the position centroid, the + offset is the mean defocus and the gradient is exactly orthogonal to it -- so a + subsequent :meth:`grid_search_hyperparameters` over ``C10`` stays well posed. + + Parameters + ---------- + update_defocus : bool + Also write the fitted offset into the optimized ``C10``. On by default: the offset + and the gradient are fit jointly, so keeping a stale global ``C10`` alongside a + fresh gradient would be inconsistent. + + Returns + ------- + self + """ + if verbose is None: + verbose = self.verbose + + results = self.defocus_map( + c10_values, patch_grid=patch_grid, verbose=verbose, **defocus_map_kwargs + ) + + valid = results["valid"] + if valid.sum() < 3: + raise RuntimeError( + f"Only {int(valid.sum())} of {valid.size} patches gave a bracketed minimum, " + "and a plane needs 3. Widen `c10_values`, or use a coarser `patch_grid` so " + "each patch has more positions." + ) + + centers = results["centers_A"][valid] + design = np.column_stack([np.ones(len(centers)), centers]) + offset, g_row, g_col = np.linalg.lstsq(design, results["c10_best"][valid], rcond=None)[0] + + self.defocus_gradient = (float(g_row), float(g_col)) + + if update_defocus: + state = self.hyperparameter_state + coefs = state.current_aberrations() + coefs["C10"] = float(offset) + state.optimized_aberrations = validate_aberration_coefficients(coefs) + state.optimized_keys.add("C10") + + if verbose: + residual = results["c10_best"][valid] - design @ [offset, g_row, g_col] + print( + f"Fitted defocus plane over {int(valid.sum())}/{valid.size} patches:\n" + f" C10 offset = {offset:.1f} A\n" + f" defocus_gradient= ({g_row:.4g}, {g_col:.4g}) A/A " + f"(|g| = {math.hypot(g_row, g_col):.4g}, " + f"tilt = {math.degrees(math.atan(math.hypot(g_row, g_col))):.2f} deg)\n" + f" residual RMS = {float(np.sqrt((residual**2).mean())):.1f} A" + ) + + return self diff --git a/src/quantem/diffractive_imaging/optimize_hyperparameters.py b/src/quantem/diffractive_imaging/optimize_hyperparameters.py index ea73bf02..7eca0342 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 -from dataclasses import dataclass from typing import Any, Callable, Dict, Mapping, Optional import matplotlib.pyplot as plt @@ -16,28 +15,12 @@ PtychographyDatasetBase, PtychographyDatasetRaster, ) +from quantem.diffractive_imaging.ptycho_utils import ( + OptimizationParameter as OptimizationParameter, # re-export: the documented path +) from quantem.diffractive_imaging.ptychography_lite import PtychoLite, PtychoLiteDIP -@dataclass -class OptimizationParameter: - """Specification for a parameter to optimize.""" - - low: float - high: float - log: bool = False - n_points: int | None = None - - def grid_values(self): - """Return an array of grid values for this parameter.""" - if self.n_points is None: - raise ValueError("n_points must be specified for grid search parameters.") - if self.log: - return np.geomspace(self.low, self.high, self.n_points) - else: - return np.linspace(self.low, self.high, self.n_points) - - def _suggest_from_spec(trial: optuna.trial.Trial, spec: OptimizationParameter, name: str) -> float: """Sample a value from an OptimizationParameter using Optuna trial.""" if spec.log: diff --git a/src/quantem/diffractive_imaging/origin_models.py b/src/quantem/diffractive_imaging/origin_models.py index b74b2f27..30fdda7c 100644 --- a/src/quantem/diffractive_imaging/origin_models.py +++ b/src/quantem/diffractive_imaging/origin_models.py @@ -160,9 +160,13 @@ def fit_origin_background( xa, ya = torch.meshgrid(x, y, indexing="ij") probe_positions = torch.stack([xa, ya], -1).view((-1, 2)) else: - probe_positions = validate_tensor( - probe_positions, "probe positions", dtype=torch.float - ).view((-1, 2)) + # `.to(self.device)` matters: positions usually arrive as a numpy array, and a + # CPU tensor meeting `origin_measured` on MPS raises rather than promoting + probe_positions = ( + validate_tensor(probe_positions, "probe positions", dtype=torch.float) + .to(self.device) + .view((-1, 2)) + ) if probe_positions.shape != self.origin_measured.shape: raise ValueError("probe positions shape must match the measured origins.") diff --git a/src/quantem/diffractive_imaging/ptycho_utils.py b/src/quantem/diffractive_imaging/ptycho_utils.py index 8a102767..fad4c514 100644 --- a/src/quantem/diffractive_imaging/ptycho_utils.py +++ b/src/quantem/diffractive_imaging/ptycho_utils.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from math import ceil from typing import Literal, Union, overload @@ -10,6 +11,32 @@ ArrayLike = Union[np.ndarray, "torch.Tensor"] +@dataclass +class OptimizationParameter: + """Specification for a parameter to optimize. + + Shared by the iterative and direct hyperparameter searches, which both test candidate + specifications with ``isinstance``. It lives here, rather than in either of them, so + that there is only ever one class and a value built for one search is accepted by the + other. It is re-exported from ``direct_ptychography`` and ``optimize_hyperparameters``, + which are the paths callers already use. + """ + + low: float + high: float + log: bool = False + n_points: int | None = None + + def grid_values(self): + """Return an array of grid values for this parameter.""" + if self.n_points is None: + raise ValueError("n_points must be specified for grid search parameters.") + if self.log: + return np.geomspace(self.low, self.high, self.n_points) + else: + return np.linspace(self.low, self.high, self.n_points) + + # TODO: figure out what here should be put into ptycho base vs kept in a utilities file diff --git a/tests/diffractive_imaging/conftest.py b/tests/diffractive_imaging/conftest.py new file mode 100644 index 00000000..83d3eea1 --- /dev/null +++ b/tests/diffractive_imaging/conftest.py @@ -0,0 +1,594 @@ +"""Shared synthetic 4D-STEM data for the direct-ptychography test modules. + +Mirrors the simulation idiom in ``test_ptychography.py`` (white-noise phase object, +soft-aperture defocused probe, ``|FFT(obj_patch * probe)|**2``) but on a smaller grid and +with a probe small enough that parallax shifts stay well inside the scan. + +Imported by ``test_direct_ptychography.py`` and ``test_direct_ptychography_montage.py`` via +``from .conftest import ...``, matching the pattern used by the tomography suite. +""" + +import numpy as np +import pytest +import torch + +from quantem.core.datastructures import Dataset3d, Dataset4d +from quantem.core.utils.utils import electron_wavelength_angstrom + +N = 32 # detector / object gridpoints +Q_MAX = 0.5 # inverse Angstroms +Q_PROBE = Q_MAX / 4 # inverse Angstroms -> BF disk radius of N/8 = 4 px +PROBE_ENERGY = 300e3 # eV +SCAN_STEP_SIZE = 1 # pixels + +SAMPLING = 1 / Q_MAX / 2 # Angstroms +RECIPROCAL_SAMPLING = 2 * Q_MAX / N # inverse Angstroms +SEMIANGLE_CUTOFF = Q_PROBE * electron_wavelength_angstrom(PROBE_ENERGY) * 1e3 # mrad +SCAN_SAMPLING = SAMPLING * SCAN_STEP_SIZE # Angstroms + +#: fixing the fitted origin keeps the center-of-mass step exact and deterministic +ORIGIN = (N // 2, N // 2) + +DECONVOLUTION_KERNELS = ("ssb", "obf", "mf", "prlx", "icom") + + +def integer_shift_defocus(pixel_shift_per_k_step: int, upsampling_factor: int = 1) -> float: + """``C10`` (in Angstrom) placing every BF pixel's parallax shift on an exact pixel. + + For pure defocus the aberration gradient is ``dchi/dk = 2*pi*wavelength*C10*k``, so the + lateral shift is ``wavelength*C10*k`` Angstrom. Because ``k = m * reciprocal_sampling`` + exactly (``spatial_frequencies`` uses ``fftfreq(n, 1/(dk*n))``), choosing + + C10 = p * scan_sampling / (U * wavelength * dk) + + gives a shift of exactly ``p * m`` upsampled pixels for integer detector index ``m``. + """ + wavelength = electron_wavelength_angstrom(PROBE_ENERGY) + return ( + pixel_shift_per_k_step + * SCAN_SAMPLING + / (upsampling_factor * wavelength * RECIPROCAL_SAMPLING) + ) + + +def make_complex_obj(seed: int = 42) -> np.ndarray: + """White-noise pure-phase object.""" + rng = np.random.default_rng(seed) + arr = rng.random((N, N)) + arr -= arr.mean() + return np.exp(1.0j * arr.astype(np.float32)) + + +def band_limited_phase(seed: int = 42, cutoff: float = 2 * Q_PROBE) -> np.ndarray: + """Ground-truth phase, low-passed to what the bright-field disk can actually transfer. + + Correlating a reconstruction against the raw white-noise phase caps out near 0.4 simply + because most of its power sits above the aperture cutoff; band-limiting first makes + "did this recover the object" a meaningful question. + """ + phase = np.angle(make_complex_obj(seed)) + qx = qy = np.fft.fftfreq(N, SCAN_SAMPLING) + q = np.hypot(qx[:, None], qy[None, :]) + limited = np.fft.ifft2(np.fft.fft2(phase) * (q <= cutoff)).real + return limited - limited.mean() + + +def make_probe_array(defocus: float) -> np.ndarray: + """Soft-aperture probe with ``C10 = defocus`` Angstrom.""" + qx = qy = np.fft.fftfreq(N, SAMPLING) + q = np.sqrt(qx[:, None] ** 2 + qy[None, :] ** 2) + + aperture_fourier = np.sqrt( + np.clip((Q_PROBE - q) / RECIPROCAL_SAMPLING + 0.5, 0, 1), + ) + chi = q**2 * electron_wavelength_angstrom(PROBE_ENERGY) * np.pi * defocus + probe_fourier = aperture_fourier * np.exp(-1j * chi) + probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2)) + return np.fft.ifft2(probe_fourier) * N + + +def scan_positions_px() -> np.ndarray: + """``(N_pos, 2)`` raster positions in object pixels, "ij" (row, col) ordering.""" + n = N // SCAN_STEP_SIZE + ii, jj = np.meshgrid( + np.arange(n) * SCAN_STEP_SIZE, + np.arange(n) * SCAN_STEP_SIZE, + indexing="ij", + ) + return np.stack((ii.ravel(), jj.ravel()), axis=-1).astype(np.float64) + + +def simulate_intensities(complex_obj: np.ndarray, probe: np.ndarray) -> np.ndarray: + """``(N_pos, N, N)`` diffraction intensities, corner-centered in reciprocal space. + + ``probe`` is either a single ``(N, N)`` array or a ``(N_pos, N, N)`` stack, which + broadcasts against the extracted object patches and gives every scan position its own + probe -- how a tilted sample is simulated. + """ + positions_px = scan_positions_px() + x0 = np.round(positions_px[:, 0]).astype(int) + y0 = np.round(positions_px[:, 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]) % N + col = (y0[:, None, None] + y_ind[None, None, :]) % N + + exit_waves = complex_obj[row, col] * probe + return np.abs(np.fft.fft2(exit_waves)) ** 2 + + +def make_dataset4d(defocus: float | None = None, seed: int = 42) -> Dataset4d: + """``(n_scan, n_scan, N, N)`` 4D-STEM dataset, reciprocal axes fftshifted.""" + if defocus is None: + defocus = integer_shift_defocus(1) + + complex_obj = make_complex_obj(seed) + probe = make_probe_array(defocus) + intensities = simulate_intensities(complex_obj, probe) + + n = N // SCAN_STEP_SIZE + array = np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((n, n, N, N)) + + return Dataset4d.from_array( + array.astype(np.float32), + name="synthetic 4D-STEM", + sampling=(SCAN_SAMPLING, SCAN_SAMPLING, RECIPROCAL_SAMPLING, RECIPROCAL_SAMPLING), + units=("A", "A", "A^-1", "A^-1"), + ) + + +def defocus_per_position(mean_defocus: float, defocus_gradient: tuple[float, float]) -> np.ndarray: + """``(N_pos,)`` local defocus for a tilted sample, mean-zero about the scan centroid. + + Matches ``DirectPtychographyMontage._return_delta_c10``: the offset is measured from the + centroid of the scan positions in Angstrom, in the unrotated scan frame. + """ + positions_ang = scan_positions_px() * SCAN_SAMPLING + offsets = positions_ang - positions_ang.mean(axis=0) + return mean_defocus + offsets @ np.asarray(defocus_gradient, dtype=np.float64) + + +def make_tilted_dataset4d( + mean_defocus: float, + defocus_gradient: tuple[float, float], + seed: int = 42, +) -> Dataset4d: + """4D-STEM dataset whose defocus varies linearly across the field of view. + + The gradients used in the tests look enormous as tilts -- at this fixture's 32 Angstrom + field of view a gradient of 12.7 A/A is needed to swing the parallax shift by a single + scan pixel. That is an artifact of the tiny synthetic scan, not of the method: it is only + a +/-12% swing on the ~1625 A baseline defocus. + """ + complex_obj = make_complex_obj(seed) + defocus = defocus_per_position(mean_defocus, defocus_gradient) + probes = np.stack([make_probe_array(float(d)) for d in defocus]) + intensities = simulate_intensities(complex_obj, probes) + + n = N // SCAN_STEP_SIZE + array = np.fft.fftshift(intensities * 100, axes=(-2, -1)).reshape((n, n, N, N)) + + return Dataset4d.from_array( + array.astype(np.float32), + name="synthetic tilted 4D-STEM", + sampling=(SCAN_SAMPLING, SCAN_SAMPLING, RECIPROCAL_SAMPLING, RECIPROCAL_SAMPLING), + units=("A", "A", "A^-1", "A^-1"), + ) + + +def _bilinear_sample_periodic(image: np.ndarray, coords: np.ndarray) -> np.ndarray: + """Sample ``image`` at fractional ``(..., 2)`` coordinates, wrapping at the edges.""" + n_rows, n_cols = image.shape + base = np.floor(coords) + frac = coords - base + base = base.astype(np.int64) + + out = np.zeros(coords.shape[:-1], dtype=np.float64) + for d_row, d_col in ((0, 0), (1, 0), (0, 1), (1, 1)): + w_row = frac[..., 0] if d_row else 1 - frac[..., 0] + w_col = frac[..., 1] if d_col else 1 - frac[..., 1] + out += ( + w_row * w_col * image[(base[..., 0] + d_row) % n_rows, (base[..., 1] + d_col) % n_cols] + ) + return out + + +def make_model_vbf_stack( + mean_defocus: float, + defocus_gradient: tuple[float, float] = (0.0, 0.0), + scan_gpts: tuple[int, int] = (96, 96), + seed: int = 42, +): + """Virtual bright-field stack built directly from the parallax model, at any scan size. + + Returns ``(vbf_dataset, bf_mask_dataset, object_phase)``. + + Every image is ``v_m(r) = object[r + shift_m(C10(r))]``, the relation the montage inverts + -- the ``+`` sign was measured against the 4D pipeline, where a bright-field image + correlates with ``obj[r + shift]`` at 0.30 versus 0.09 for ``obj[r - shift]``. + + The full 4D simulator cannot be used for the position-dependent tests: its 32x32 scan + only affords ~16-position patches, and at that size the per-patch defocus estimator has a + ~430 Angstrom bias -- larger than the signal, as the zero-gradient control shows. This + builds the same relation at whatever scan size the estimator needs, cheaply. + + Bright-field pixels are ordered by ``np.nonzero`` over the corner-centered mask, which is + the order ``torch.nonzero`` gives in ``_return_bf_context``. Pass ``crop_bf_mask=False`` + so the class uses this mask verbatim and the ordering is preserved. + """ + from quantem.core.datastructures import Dataset2d, Dataset3d + + wavelength = electron_wavelength_angstrom(PROBE_ENERGY) + + # detector grid just large enough to hold the bright-field disk + radius_px = int(round(Q_PROBE / RECIPROCAL_SAMPLING)) + n_det = 2 * radius_px + 3 + centered = np.hypot(*np.meshgrid(*(np.arange(n_det) - n_det // 2,) * 2, indexing="ij")) + bf_mask = np.fft.ifftshift(centered <= radius_px) + + # k = m * reciprocal_sampling exactly, independent of the grid size + freqs = np.fft.fftfreq(n_det, 1 / (RECIPROCAL_SAMPLING * n_det)) + inds_i, inds_j = np.nonzero(bf_mask) + k_vec = np.stack((freqs[inds_i], freqs[inds_j]), axis=-1) # (num_bf, 2) + + rng = np.random.default_rng(seed) + phase = rng.random(scan_gpts) + qx = np.fft.fftfreq(scan_gpts[0], SCAN_SAMPLING) + qy = np.fft.fftfreq(scan_gpts[1], SCAN_SAMPLING) + q = np.hypot(qx[:, None], qy[None, :]) + obj = np.fft.ifft2(np.fft.fft2(phase) * (q <= 2 * Q_PROBE)).real + obj -= obj.mean() + + ii, jj = np.meshgrid(np.arange(scan_gpts[0]), np.arange(scan_gpts[1]), indexing="ij") + positions_px = np.stack((ii.ravel(), jj.ravel()), axis=-1).astype(np.float64) + + offsets_ang = (positions_px - positions_px.mean(0)) * SCAN_SAMPLING + c10 = mean_defocus + offsets_ang @ np.asarray(defocus_gradient, dtype=np.float64) + + # shift_px[m, n] = wavelength * C10[n] * k[m] / scan_sampling + shifts = ( + wavelength * c10[None, :, None] * k_vec[:, None, :] / SCAN_SAMPLING + ) # (num_bf, N_pos, 2) + stack = _bilinear_sample_periodic(obj, positions_px[None] + shifts) + + vbf_dataset = Dataset3d.from_array( + stack.reshape(len(k_vec), *scan_gpts).astype(np.float32), + name="model virtual BF stack", + sampling=(1.0, SCAN_SAMPLING, SCAN_SAMPLING), + units=("index", "A", "A"), + ) + bf_mask_dataset = Dataset2d.from_array( + bf_mask, + name="BF mask", + sampling=(RECIPROCAL_SAMPLING, RECIPROCAL_SAMPLING), + units=("A^-1", "A^-1"), + ) + return vbf_dataset, bf_mask_dataset, obj + + +def model_vbf_kwargs(defocus: float) -> dict: + """Constructor kwargs for a stack from :func:`make_model_vbf_stack`.""" + return dict( + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + aberration_coefs={"C10": defocus}, + crop_bf_mask=False, + verbose=False, + ) + + +def direct_ptycho_kwargs(defocus: float) -> dict: + """Constructor kwargs shared by both direct-ptychography classes.""" + return dict( + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + aberration_coefs={"C10": defocus}, + force_fitted_origin=ORIGIN, + verbose=False, + ) + + +def correlation(image: np.ndarray, reference: np.ndarray) -> float: + """Pearson correlation between a reconstruction and a reference, means removed.""" + a = np.asarray(image, dtype=np.float64).ravel() + b = np.asarray(reference, dtype=np.float64).ravel() + a = a - a.mean() + b = b - b.mean() + return float(np.corrcoef(a, b)[0, 1]) + + +@pytest.fixture(scope="module") +def dataset4d(): + """Read-only synthetic dataset at the integer-shift defocus.""" + return make_dataset4d() + + +# --------------------------------------------------------------------------- +# white-noise / analytical-CTF harness +# +# The object has constant Fourier amplitude, so |FFT(reconstruction)| *is* the contrast +# transfer function and can be compared against an analytical one. This is the only fixture +# here that can see whether `upsampling_factor` extends the CTF ("unfolding") or merely +# replicates it: the scan pitch is deliberately coarser than the aperture's transfer limit, +# leaving a genuine factor of 2 above the scan Nyquist to recover. Ported from +# ipynb-playground/white-noise-ptycho.ipynb. +# --------------------------------------------------------------------------- + +CTF_N = 64 # object gridpoints +CTF_K_MAX = 2.0 # inverse Angstroms +CTF_K_PROBE = 1.0 # inverse Angstroms -> transfers to 2 * K_PROBE +CTF_SCAN_STEP = 2 # object pixels -> scan Nyquist is half the transfer limit +CTF_SAMPLING = 1 / CTF_K_MAX / 2 +CTF_RECIPROCAL_SAMPLING = 2 * CTF_K_MAX / CTF_N +CTF_SCAN_SAMPLING = CTF_SAMPLING * CTF_SCAN_STEP +CTF_SEMIANGLE = CTF_K_PROBE * electron_wavelength_angstrom(PROBE_ENERGY) * 1e3 +CTF_ABERRATIONS = {"C10": 100, "C12": 50, "phi12": np.deg2rad(11)} + + +def white_noise_object_2D(n: int = CTF_N, phi0: float = 1.0, seed: int = 0) -> np.ndarray: + """Real 2D array whose FFT has random phase and *constant* amplitude.""" + rng = np.random.default_rng(seed) + even = n % 2 == 0 + pos = np.arange(1, (n if even else n + 1) // 2) + neg = np.flip(np.arange(n // 2 + 1, n)) + + arr = rng.standard_normal((n, n)) + arr[pos[:, None], pos[None, :]] = -arr[neg[:, None], neg[None, :]] + arr[pos[:, None], neg[None, :]] = -arr[neg[:, None], pos[None, :]] + arr[0, pos] = -arr[0, neg] + arr[pos, 0] = -arr[neg, 0] + if even: + arr[n // 2, :] = 0 + arr[:, n // 2] = 0 + arr[0, 0] = 0 + + return np.fft.ifft2(np.exp(2j * np.pi * arr) * phi0).real + + +def _ctf_grids(): + kx = ky = np.fft.fftfreq(CTF_N, CTF_SAMPLING) + k2 = kx[:, None] ** 2 + ky[None, :] ** 2 + phi = np.arctan2(ky[None, :], kx[:, None]) + aperture = np.clip((CTF_K_PROBE - np.sqrt(k2)) / CTF_RECIPROCAL_SAMPLING + 0.5, 0, 1) + return k2, phi, aperture + + +def ctf_probe(aberrations: dict = CTF_ABERRATIONS): + """``(chi, probe_array)`` for the soft-aperture probe used by the CTF harness.""" + k2, phi, aperture = _ctf_grids() + wavelength = electron_wavelength_angstrom(PROBE_ENERGY) + + chi = k2 * wavelength * np.pi * aberrations.get("C10", 0.0) + chi = chi + ( + k2 + * wavelength + * np.pi + * aberrations.get("C12", 0.0) + * np.cos(2 * (phi - aberrations.get("phi12", 0.0))) + ) + probe_fourier = aperture * np.exp(-1j * chi) + probe_fourier /= np.sqrt(np.sum(np.abs(probe_fourier) ** 2)) + return chi, np.fft.ifft2(probe_fourier) * CTF_N + + +def analytical_parallax_ctf(chi: np.ndarray): + """``(full_band, scan_band)`` analytical parallax CTF, fftshifted.""" + _, _, aperture = _ctf_grids() + aperture_0 = aperture / np.sqrt(np.sum(np.abs(aperture) ** 2)) + autocorr = np.real(np.fft.ifft2(np.abs(np.fft.fft2(aperture_0)) ** 2)) + full = np.fft.fftshift(np.abs(autocorr * -np.sin(chi))) + return full, full[CTF_N // 4 : -CTF_N // 4, CTF_N // 4 : -CTF_N // 4] + + +def ctf_raster_positions() -> np.ndarray: + """``(N_pos, 2)`` raster positions in *object* pixels.""" + axis = np.arange(0.0, CTF_N, CTF_SCAN_STEP) + xx, yy = np.meshgrid(axis, axis, indexing="ij") + return np.stack((xx.ravel(), yy.ravel()), axis=-1) + + +def ctf_jittered_positions(amplitude_scan_px: float, seed: int = 1) -> np.ndarray: + """Raster positions scattered off the lattice, with the outer ring pinned. + + Pinning the boundary keeps the bounding box exactly the scan field of view, so a + ``boundary="wrap"`` montage and a regridded reconstruction stay directly comparable. + """ + positions = ctf_raster_positions() + rng = np.random.default_rng(seed) + amplitude = amplitude_scan_px * CTF_SCAN_STEP + jitter = rng.uniform(-amplitude, amplitude, positions.shape) + edge = ( + (positions[:, 0] == 0) + | (positions[:, 1] == 0) + | (positions[:, 0] == CTF_N - CTF_SCAN_STEP) + | (positions[:, 1] == CTF_N - CTF_SCAN_STEP) + ) + jitter[edge] = 0.0 + return positions + jitter + + +def ctf_simulate(complex_obj: np.ndarray, probe: np.ndarray, positions_px: np.ndarray): + """``(N_pos, n, n)`` intensities, using the Fourier shift theorem off-lattice.""" + n = CTF_N + if np.allclose(positions_px, np.round(positions_px)): + x0 = np.round(positions_px[:, 0]).astype(int) + y0 = np.round(positions_px[:, 1]).astype(int) + idx = np.fft.fftfreq(n, d=1 / n).astype(int) + patches = complex_obj[ + (x0[:, None, None] + idx[None, :, None]) % n, + (y0[:, None, None] + idx[None, None, :]) % n, + ] + else: + fx = fy = np.fft.fftfreq(n) + obj_fourier = np.fft.fft2(complex_obj) + ramp = np.exp( + 2j + * np.pi + * ( + fx[None, :, None] * positions_px[:, 0, None, None] + + fy[None, None, :] * positions_px[:, 1, None, None] + ) + ) + patches = np.fft.ifft2(obj_fourier[None] * ramp, axes=(-2, -1)) + + return (np.abs(np.fft.fft2(patches * probe)) ** 2).astype(np.float32) + + +def ctf_dataset4d(intensities: np.ndarray) -> Dataset4d: + side = int(round(np.sqrt(len(intensities)))) + return Dataset4d.from_array( + np.fft.fftshift(intensities.reshape(side, side, CTF_N, CTF_N), axes=(-1, -2)), + name="white noise 4D-STEM", + sampling=(CTF_SCAN_SAMPLING,) * 2 + (CTF_RECIPROCAL_SAMPLING,) * 2, + units=("A", "A", "A^-1", "A^-1"), + ) + + +def ctf_dataset3d(intensities: np.ndarray) -> Dataset3d: + return Dataset3d.from_array( + np.fft.fftshift(intensities, axes=(-1, -2)), + name="white noise ungridded", + sampling=(1.0, CTF_RECIPROCAL_SAMPLING, CTF_RECIPROCAL_SAMPLING), + units=("index", "A^-1", "A^-1"), + ) + + +def measured_ctf(obj: np.ndarray) -> np.ndarray: + return np.fft.fftshift(np.abs(np.fft.fft2(obj)) * 2) + + +def ctf_band_scores(measured: np.ndarray, analytic: np.ndarray, pixel_size: float): + """Correlation with the analytical CTF, inside and beyond the scan Nyquist.""" + if measured.shape != analytic.shape: + raise ValueError(f"shape mismatch {measured.shape} vs {analytic.shape}") + + freq = np.fft.fftshift(np.fft.fftfreq(measured.shape[0], pixel_size)) + q = np.hypot(freq[:, None], freq[None, :]) + nyquist = 1 / (2 * CTF_SCAN_SAMPLING) + + def corr(mask): + a, b = measured[mask].astype(np.float64), analytic[mask].astype(np.float64) + if a.std() == 0 or b.std() == 0: + return float("nan") + return float(np.corrcoef(a, b)[0, 1]) + + extension = (q > nyquist) & (q <= 2 * CTF_K_PROBE) + return corr(q <= nyquist), corr(extension) + + +def ctf_kwargs() -> dict: + return dict( + energy=PROBE_ENERGY, + semiangle_cutoff=CTF_SEMIANGLE, + rotation_angle=0.0, + aberration_coefs=CTF_ABERRATIONS, + verbose=False, + ) + + +@pytest.fixture(scope="module") +def ctf_scene(): + """``(complex_obj, chi, probe, ctf_full, ctf_sub)`` for the white-noise harness.""" + chi, probe = ctf_probe() + full, sub = analytical_parallax_ctf(chi) + return np.exp(1j * white_noise_object_2D()), chi, probe, full, sub + + +def ctf_radial_correlation(measured: np.ndarray, analytic: np.ndarray, pixel_size: float) -> float: + """Correlation of radially averaged CTF profiles, independent of grid shape. + + A grid sized to cover scattered positions rarely lands on the same shape as the + analytical CTF, so compare profiles rather than pixels. + """ + + def profile(image, px, bins=40, q_max=2 * CTF_K_PROBE): + freq = np.fft.fftshift(np.fft.fftfreq(image.shape[0], px)) + q = np.hypot(freq[:, None], freq[None, :]) + index = np.clip((q / q_max * bins).astype(int), 0, bins - 1) + inside = q <= q_max + return np.array( + [ + image[inside & (index == i)].mean() if (inside & (index == i)).any() else 0.0 + for i in range(bins) + ] + ) + + a = profile(measured, pixel_size) + b = profile(analytic, CTF_SAMPLING) + return float(np.corrcoef(a, b)[0, 1]) + + +def analytic_probe_array(reconstruction, aberration_coefs): + """The analytic probe of a reconstruction, materialized as a complex array. + + Feeding this back in as a `FourierProbe.from_array` must reproduce the reconstruction it + came from -- which is what pins the empirical path against the analytic one. + """ + from quantem.diffractive_imaging.complex_probe import ( + evaluate_probe, + polar_spatial_frequencies, + ) + + k, phi = polar_spatial_frequencies( + reconstruction.gpts, reconstruction.sampling, device=reconstruction.device + ) + return evaluate_probe( + k * reconstruction.wavelength, + phi, + reconstruction.semiangle_cutoff, + reconstruction.angular_sampling, + reconstruction.wavelength, + aberration_coefs=aberration_coefs, + ) + + +def ctf_interleaved_frames( + complex_obj: np.ndarray, + probe: np.ndarray, + drift: np.ndarray, + seed: int = 3, +): + """Split the CTF raster into interleaved frames, each displaced by a known drift. + + Models a multi-frame acquisition: every frame samples the whole field of view sparsely, + the frames interleave to fill it, and the specimen moves between them. The recorded + positions stay nominal while the *object* is imaged displaced, which is what specimen + drift does and what ``estimate_frame_drift`` has to undo. + + The sign follows the specimen: a feature at object coordinate ``x`` moves to ``x + + drift`` in the microscope's frame, so it is measured when the probe is nominally at + ``x + drift`` -- equivalently the probe illuminates ``nominal - drift``. The + reconstruction then shows the object translated by ``+drift``, and ``positions - drift`` + puts it back. + + Returns ``(datasets, positions_A)``, both lists of length ``len(drift)``, with positions + in Angstrom -- what the ungridded constructors take. + """ + positions = ctf_raster_positions() + rng = np.random.default_rng(seed) + assignment = rng.permutation(len(positions)) % len(drift) + + datasets, positions_A = [], [] + for frame, offset in enumerate(np.asarray(drift, dtype=np.float64)): + nominal = positions[assignment == frame] + illuminated = nominal - offset / CTF_SAMPLING + intensities = ctf_simulate(complex_obj, probe, illuminated) + datasets.append(ctf_dataset3d(intensities)) + positions_A.append(nominal * CTF_SAMPLING) + + return datasets, positions_A + + +#: every accelerator present, so the same assertions run on CPU, CUDA and MPS +ACCELERATORS = [ + device + for device, present in ( + ("cuda", torch.cuda.is_available()), + ("mps", torch.backends.mps.is_available()), + ) + if present +] diff --git a/tests/diffractive_imaging/test_direct_ptychography.py b/tests/diffractive_imaging/test_direct_ptychography.py new file mode 100644 index 00000000..1796a768 --- /dev/null +++ b/tests/diffractive_imaging/test_direct_ptychography.py @@ -0,0 +1,1117 @@ +"""Tests for the Fourier-space direct-ptychography reconstruction. + +Covers all five deconvolution kernels, the aberration/rotation conventions they share, the +four hyperparameter-fitting routines, and save/load. Synthetic data comes from +``conftest.py``: a white-noise pure-phase object imaged with a defocused soft aperture. +""" + +import warnings + +import numpy as np +import pytest +import torch + +from quantem.core.datastructures import Dataset2d, Dataset3d +from quantem.core.io.serialize import load +from quantem.core.utils.utils import electron_wavelength_angstrom, to_numpy +from quantem.diffractive_imaging import DirectPtychography, OptimizationParameter +from quantem.diffractive_imaging.complex_probe import spatial_frequencies + +from .conftest import ( + CTF_SAMPLING, + CTF_SCAN_SAMPLING, + DECONVOLUTION_KERNELS, + PROBE_ENERGY, + Q_PROBE, + SCAN_SAMPLING, + SEMIANGLE_CUTOFF, + N, + band_limited_phase, + correlation, + ctf_band_scores, + ctf_dataset3d, + ctf_dataset4d, + ctf_jittered_positions, + ctf_kwargs, + ctf_radial_correlation, + ctf_raster_positions, + ctf_simulate, + direct_ptycho_kwargs, + integer_shift_defocus, + make_dataset4d, + measured_ctf, + scan_positions_px, +) + +TRUE_C10 = integer_shift_defocus(1) + + +def _build(dataset4d, defocus=TRUE_C10, **overrides): + kwargs = dict(direct_ptycho_kwargs(defocus), edge_blend_pixels=0) + kwargs.update(overrides) + return DirectPtychography.from_dataset4d(dataset4d, **kwargs) + + +@pytest.fixture(scope="module") +def recon(dataset4d): + """A reconstruction seeded with the true defocus. Reconstruct before asserting.""" + return _build(dataset4d) + + +class TestConstruction: + def test_geometry_matches_the_dataset(self, recon): + assert recon.bf_mask.shape == recon.gpts + assert recon.vbf_stack.shape == (recon.num_bf, N, N) + assert recon.num_bf == int(recon.bf_mask.sum()) + assert recon.scan_gpts == (N, N) + assert recon.scan_sampling[0] == pytest.approx(SCAN_SAMPLING) + + def test_bf_mask_is_about_the_aperture_size(self, recon): + """The mask is thresholded from the mean pattern, so it should track the BF disk.""" + expected_area = np.pi * (Q_PROBE / recon.reciprocal_sampling[0]) ** 2 + assert recon.num_bf == pytest.approx(expected_area, rel=0.35) + + def test_preprocess_zeroes_the_dc_bin(self, recon): + assert torch.allclose( + recon._vbf_fourier[..., 0, 0], torch.zeros_like(recon._vbf_fourier[..., 0, 0]) + ) + + def test_from_virtual_bfs_reproduces_from_dataset4d(self, dataset4d, recon): + """Re-wrapping the stored stack must give a bit-identical reconstruction.""" + vbf_dataset = Dataset3d.from_array( + to_numpy(recon.vbf_stack), + name="vBF stack", + units=("index", "A", "A"), + sampling=(1, SCAN_SAMPLING, SCAN_SAMPLING), + ) + bf_mask_dataset = Dataset2d.from_array( + to_numpy(recon.bf_mask), + name="BF mask", + units=("A^-1", "A^-1"), + sampling=tuple(recon.reciprocal_sampling), + ) + rebuilt = DirectPtychography.from_virtual_bfs( + vbf_dataset, + bf_mask_dataset, + energy=PROBE_ENERGY, + rotation_angle=0.0, + semiangle_cutoff=SEMIANGLE_CUTOFF, + aberration_coefs={"C10": TRUE_C10}, + crop_bf_mask=False, # the stored mask is already cropped + verbose=False, + ) + + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + rebuilt.reconstruct(deconvolution_kernel="ssb", verbose=False) + assert np.allclose(rebuilt.obj, recon.obj, rtol=1e-5, atol=1e-8) + + def test_cropping_the_bf_mask_preserves_the_reconstruction(self, dataset4d): + """Cropping shrinks the detector grid but must not permute the vBF stack order.""" + cropped = _build(dataset4d, crop_bf_mask=True) + uncropped = _build(dataset4d, crop_bf_mask=False) + + assert cropped.gpts[0] < uncropped.gpts[0] + assert cropped.num_bf == uncropped.num_bf + + cropped.reconstruct(deconvolution_kernel="ssb", verbose=False) + uncropped.reconstruct(deconvolution_kernel="ssb", verbose=False) + assert correlation(cropped.obj, uncropped.obj) > 0.99 + + def test_direct_instantiation_is_blocked(self): + with pytest.raises(RuntimeError, match="from_virtual_bfs"): + DirectPtychography( + vbf_dataset=None, + bf_mask_dataset=None, + energy=PROBE_ENERGY, + rotation_angle=0.0, + aberration_coefs={}, + semiangle_cutoff=SEMIANGLE_CUTOFF, + soft_edges=True, + crop_bf_mask=False, + bf_mask_padding_px=1, + rng=None, + device="cpu", + verbose=False, + ) + + +class TestDeconvolutionKernels: + @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) + def test_recovers_the_band_limited_object(self, recon, kernel): + """Every kernel must recover the object over the band the aperture transfers.""" + recon.reconstruct(deconvolution_kernel=kernel, verbose=False) + + assert recon.obj.shape == (N, N) + assert np.isfinite(recon.obj).all() + assert abs(correlation(recon.obj, band_limited_phase())) > 0.7 + + @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) + def test_aliases_resolve(self, recon, kernel): + aliases = { + "ssb": "single-sideband", + "obf": "optimum-bright-field", + "mf": "matched-filter", + "prlx": "tilt-corrected-bright-field", + "icom": "center-of-mass", + } + by_short = recon.reconstruct(deconvolution_kernel=kernel, verbose=False).obj.copy() + by_alias = recon.reconstruct(deconvolution_kernel=aliases[kernel], verbose=False).obj + + assert np.array_equal(by_short, by_alias) + + @pytest.mark.parametrize("upsampling_factor", [1, 2, 3]) + def test_upsampling_preserves_the_field_of_view(self, recon, upsampling_factor): + recon.reconstruct( + deconvolution_kernel="ssb", upsampling_factor=upsampling_factor, verbose=False + ) + + assert recon.obj.shape == (N * upsampling_factor, N * upsampling_factor) + assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx(N * SCAN_SAMPLING) + + def test_corrected_stack_sums_to_corrected_bf(self, recon): + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + + assert recon.corrected_stack.shape == (recon.num_bf, N, N) + assert torch.allclose(recon.corrected_stack.sum(0), recon.corrected_bf) + + def test_unknown_kernel_raises(self, recon): + with pytest.raises(ValueError, match="Unknown deconvolution kernel"): + recon.reconstruct(deconvolution_kernel="wiener", verbose=False) + + +class TestLateralShifts: + """`_return_lateral_shifts` underpins both the parallax kernel and the montage class.""" + + def test_pure_defocus_matches_the_analytic_shift(self, recon): + """For pure C10 the lateral shift is exactly `wavelength * C10 * k` Angstrom.""" + shifts = recon._return_lateral_shifts(0.0, {"C10": TRUE_C10}, recon.bf_mask) + + kxa, kya = spatial_frequencies(recon.gpts, recon.sampling, device=recon.device) + expected = ( + torch.stack((kxa[recon.bf_mask], kya[recon.bf_mask]), -1) + * electron_wavelength_angstrom(PROBE_ENERGY) + * TRUE_C10 + ) + + assert torch.allclose(shifts, expected, rtol=1e-5, atol=1e-6) + + def test_no_aberrations_means_no_shift(self, recon): + shifts = recon._return_lateral_shifts(0.0, {}, recon.bf_mask) + + assert torch.count_nonzero(shifts) == 0 + + def test_rotation_rotates_the_shifts(self, recon): + """`_passively_rotate_grid` sends (kx, ky) -> (-ky, kx) at 90 degrees.""" + coefs = {"C10": TRUE_C10} + unrotated = recon._return_lateral_shifts(0.0, coefs, recon.bf_mask) + rotated = recon._return_lateral_shifts(90.0, coefs, recon.bf_mask) + + expected = torch.stack((-unrotated[:, 1], unrotated[:, 0]), dim=-1) + assert torch.allclose(rotated, expected, atol=1e-5) + + def test_astigmatism_breaks_the_radial_symmetry(self, recon): + radial = recon._return_lateral_shifts(0.0, {"C10": TRUE_C10}, recon.bf_mask) + astigmatic = recon._return_lateral_shifts( + 0.0, {"C10": TRUE_C10, "C12": 0.3 * TRUE_C10}, recon.bf_mask + ) + + assert not torch.allclose(radial, astigmatic) + + +class TestBrightFieldSubsets: + def test_checkerboard_halves_are_additive(self, recon): + """Reconstructions are linear in the BF sum once the per-subset weight is undone.""" + recon.reconstruct(deconvolution_kernel="prlx", parallax_flip_phase=False, verbose=False) + full = recon.corrected_bf.clone() + + halves = [] + for mask in recon._make_checkerboard_bf_masks(recon.gpts, recon.bf_mask): + recon.reconstruct( + bf_mask=mask, + deconvolution_kernel="prlx", + parallax_flip_phase=False, + verbose=False, + ) + halves.append(recon.corrected_bf.clone() * recon.corrected_stack.shape[0]) + + # each half is normalized by its own BF weight; undo that before comparing + combined = to_numpy(halves[0] + halves[1]) / recon.num_bf + assert correlation(combined, to_numpy(full)) > 0.99 + + def test_halfsets_helper_returns_two_images(self, recon): + first, second = recon._reconstruct_with_halfsets(deconvolution_kernel="ssb") + + assert first.shape == second.shape == (N, N) + assert correlation(to_numpy(first), to_numpy(second)) > 0.5 + + def test_subset_uses_fewer_bright_field_pixels(self, recon): + mask = recon._make_checkerboard_bf_masks(recon.gpts, recon.bf_mask)[0] + recon.reconstruct(bf_mask=mask, deconvolution_kernel="ssb", verbose=False) + + assert recon.corrected_stack.shape[0] == int(mask.sum()) + assert int(mask.sum()) < recon.num_bf + + +class TestFilters: + def test_lowpass_suppresses_high_frequencies(self, recon): + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + unfiltered = np.abs(np.fft.fft2(recon.obj)) + recon.reconstruct(deconvolution_kernel="ssb", q_lowpass=0.05, verbose=False) + filtered = np.abs(np.fft.fft2(recon.obj)) + + qx = qy = np.fft.fftfreq(N, SCAN_SAMPLING) + high = np.hypot(qx[:, None], qy[None, :]) > 0.15 + + assert filtered[high].sum() < 0.05 * unfiltered[high].sum() + + def test_highpass_suppresses_low_frequencies(self, recon): + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + unfiltered = np.abs(np.fft.fft2(recon.obj)) + recon.reconstruct(deconvolution_kernel="ssb", q_highpass=0.15, verbose=False) + filtered = np.abs(np.fft.fft2(recon.obj)) + + qx = qy = np.fft.fftfreq(N, SCAN_SAMPLING) + low = (np.hypot(qx[:, None], qy[None, :]) < 0.05) & ( + np.hypot(qx[:, None], qy[None, :]) > 0 + ) + + assert filtered[low].sum() < 0.2 * unfiltered[low].sum() + + def test_parallax_phase_flip_vanishes_at_zero_defocus(self, recon): + """`sign(sin(chi))` is identically zero when chi is, so the image is too. + + Not a defect: an unaberrated probe transfers no phase contrast in this formulation. + Worth pinning because the all-zero output is otherwise startling. + """ + recon.reconstruct( + deconvolution_kernel="prlx", + override_aberration_coefs={"C10": 0.0}, + parallax_flip_phase=True, + verbose=False, + ) + assert np.ptp(recon.obj) == 0.0 + + recon.reconstruct( + deconvolution_kernel="prlx", + override_aberration_coefs={"C10": 0.0}, + parallax_flip_phase=False, + verbose=False, + ) + assert np.ptp(recon.obj) > 0.0 + + +class TestVarianceLoss: + def test_is_positive_after_reconstructing(self, recon): + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + + assert float(recon.variance_loss()) > 0 + + def test_is_minimized_at_the_true_defocus(self, recon): + losses = {} + for scale in (0.5, 0.8, 1.0, 1.2, 1.5): + recon.reconstruct( + deconvolution_kernel="prlx", + override_aberration_coefs={"C10": scale * TRUE_C10}, + parallax_flip_phase=False, + verbose=False, + ) + losses[scale] = float(recon.variance_loss()) + + assert min(losses, key=losses.get) == pytest.approx(1.0) + + +class TestHyperparameterFitting: + """All four routines must recover a seeded defocus from a blind start.""" + + def test_grid_search(self, dataset4d): + recon = _build(dataset4d, aberration_coefs={}) + recon.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) + }, + deconvolution_kernel="prlx", + verbose=False, + ) + fitted = recon.hyperparameter_state.optimized_aberrations["C10"] + + step = (1.6 - 0.4) * TRUE_C10 / 6 + assert abs(fitted - TRUE_C10) <= step + + def test_least_squares(self, dataset4d): + recon = _build(dataset4d, aberration_coefs={}) + recon.fit_hyperparameters_least_squares( + cartesian_basis="defocus", fit_method="global", verbose=False + ) + fitted = recon.hyperparameter_state.optimized_aberrations["C10"] + + assert fitted == pytest.approx(TRUE_C10, rel=0.15) + + def test_cross_correlation(self, dataset4d): + recon = _build(dataset4d, aberration_coefs={}) + recon.fit_hyperparameters_cross_correlation(bin_factors=(2, 1), verbose=False) + state = recon.hyperparameter_state + + assert state.optimized_aberrations["C10"] == pytest.approx(TRUE_C10, rel=0.3) + assert state.optimized_rotation_angle == pytest.approx(0.0, abs=5.0) + + @pytest.mark.slow + def test_optuna(self, dataset4d): + recon = _build(dataset4d, aberration_coefs={}) + recon.optimize_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10) + }, + n_trials=25, + deconvolution_kernel="prlx", + verbose=False, + ) + fitted = recon.hyperparameter_state.optimized_aberrations["C10"] + + assert fitted == pytest.approx(TRUE_C10, rel=0.15) + assert recon.hyperparameter_state.study is not None + + def test_fitting_leaves_the_initial_state_intact(self, dataset4d): + """`use_initial_state=True` must ignore whatever a fit wrote back.""" + recon = _build(dataset4d) + recon.hyperparameter_state.optimized_aberrations = {"C10": 5.0 * TRUE_C10} + + recon.reconstruct(deconvolution_kernel="ssb", use_initial_state=True, verbose=False) + from_initial = recon.obj.copy() + recon.reconstruct( + deconvolution_kernel="ssb", + override_aberration_coefs={"C10": TRUE_C10}, + verbose=False, + ) + + assert np.allclose(from_initial, recon.obj) + + +class TestLossFunctions: + """The searches accept an objective by name or as a callable, and minimize it.""" + + def test_rms_gradient_is_none_before_reconstruct(self, dataset4d): + assert _build(dataset4d).rms_gradient_loss() is None + + def test_rms_gradient_peaks_at_the_true_defocus(self, dataset4d): + """It is a sharpness metric, so the correct deconvolution must maximize it.""" + recon = _build(dataset4d) + + losses = {} + for scale in (0.5, 1.0, 1.5): + recon.reconstruct( + deconvolution_kernel="prlx", + override_aberration_coefs={"C10": scale * TRUE_C10}, + verbose=False, + ) + losses[scale] = recon.rms_gradient_loss() + + # negated, so the true defocus is the *smallest* + assert losses[1.0] < losses[0.5] + assert losses[1.0] < losses[1.5] + + def test_rms_gradient_follows_the_object_sampling(self, dataset4d): + """Per Angstrom, not per pixel, so upsampling does not rescale it.""" + recon = _build(dataset4d) + recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=1, verbose=False) + coarse = recon.rms_gradient_loss() + recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) + fine = recon.rms_gradient_loss() + + # upsampling tiles the spectrum rather than adding detail, so the physical gradient + # is the same quantity; a per-pixel metric would differ by the sampling ratio + assert fine == pytest.approx(coarse, rel=0.35) + + def test_the_two_classes_agree(self, dataset4d): + """`prlx` is the same operator in both, so the objective must match.""" + from quantem.diffractive_imaging import DirectPtychographyMontage + + kwargs = dict(direct_ptycho_kwargs(TRUE_C10), edge_blend_pixels=0) + fourier = DirectPtychography.from_dataset4d(dataset4d, **kwargs) + montage = DirectPtychographyMontage.from_dataset4d(dataset4d, boundary="wrap", **kwargs) + for recon in (fourier, montage): + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + + assert montage.rms_gradient_loss() == pytest.approx(fourier.rms_gradient_loss(), rel=1e-5) + + def test_grid_search_with_the_rms_gradient(self, dataset4d): + recon = _build(dataset4d, aberration_coefs={}) + recon.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) + }, + loss="rms_gradient", + deconvolution_kernel="prlx", + verbose=False, + ) + fitted = recon.hyperparameter_state.optimized_aberrations["C10"] + + step = (1.6 - 0.4) * TRUE_C10 / 6 + assert abs(fitted - TRUE_C10) <= step + + def test_grid_search_accepts_a_callable(self, dataset4d): + """Anything that scores a reconstruction can drive a search.""" + recon = _build(dataset4d, aberration_coefs={}) + recon.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=7) + }, + loss=lambda r: -float(np.std(r.obj)), + deconvolution_kernel="prlx", + verbose=False, + ) + fitted = recon.hyperparameter_state.optimized_aberrations["C10"] + + step = (1.6 - 0.4) * TRUE_C10 / 6 + assert abs(fitted - TRUE_C10) <= step + + def test_recorded_losses_are_the_requested_objective(self, dataset4d): + """`_grid_search_results` must hold the loss actually optimized, not the default.""" + recon = _build(dataset4d, aberration_coefs={}) + recon.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter(low=0.4 * TRUE_C10, high=1.6 * TRUE_C10, n_points=3) + }, + loss="rms_gradient", + deconvolution_kernel="prlx", + verbose=False, + ) + + # the RMS gradient loss is negative; the variance loss is positive + assert all(loss < 0 for _, loss in recon._grid_search_results) + + @pytest.mark.parametrize("bad", ["sharpness", 3, None]) + def test_unknown_losses_are_rejected(self, dataset4d, bad): + recon = _build(dataset4d) + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + with pytest.raises(ValueError, match="must be a callable or one of"): + recon._return_loss_value(bad) + + +class TestHyperparameterState: + def test_optimized_overrides_initial(self, recon): + state = recon.hyperparameter_state + state.clear_optimized() + + assert state.current_rotation_angle() == 0.0 + state.optimized_rotation_angle = 12.0 + assert state.current_rotation_angle() == 12.0 + assert state.current_rotation_angle(override_fixed=3.0) == 3.0 + + state.clear_optimized() + assert state.current_aberrations()["C10"] == pytest.approx(TRUE_C10) + + def test_defocus_alias_is_negated(self): + from quantem.diffractive_imaging.direct_ptychography_base import HyperparameterState + + state = HyperparameterState(initial_aberrations={"defocus": 100.0}) + + assert state.initial_aberrations == {"C10": -100.0} + + def test_rejects_unknown_aberrations(self): + from quantem.diffractive_imaging.direct_ptychography_base import HyperparameterState + + with pytest.raises(ValueError): + HyperparameterState(initial_aberrations={"C99": 1.0}) + + +class TestSerialization: + def test_round_trip_preserves_the_reconstruction(self, dataset4d, tmp_path): + recon = _build(dataset4d) + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + before = recon.obj.copy() + + path = str(tmp_path / "direct.zip") + recon.save(path, mode="o") + restored = load(path) + + assert isinstance(restored, DirectPtychography) + assert np.array_equal(restored.obj, before) + assert restored.num_bf == recon.num_bf + assert restored.gpts == recon.gpts + + restored.reconstruct(deconvolution_kernel="ssb", verbose=False) + assert np.allclose(restored.obj, before) + + +class TestRotationSensitivity: + def test_wrong_rotation_degrades_the_reconstruction(self, dataset4d): + """The data is simulated unrotated, so 0 degrees must beat a large rotation.""" + recon = _build(dataset4d) + reference = band_limited_phase() + + recon.reconstruct(deconvolution_kernel="prlx", override_rotation_angle=0.0, verbose=False) + aligned = abs(correlation(recon.obj, reference)) + recon.reconstruct(deconvolution_kernel="prlx", override_rotation_angle=60.0, verbose=False) + misaligned = abs(correlation(recon.obj, reference)) + + assert aligned > misaligned + + def test_detector_rotation_is_estimated_when_not_supplied(self, dataset4d): + recon = DirectPtychography.from_dataset4d( + dataset4d, + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=None, + aberration_coefs={"C10": TRUE_C10}, + edge_blend_pixels=0, + verbose=False, + ) + + # simulated without rotation; the curl-minimizing estimate should land near 0 or 180 + estimated = abs(recon.rotation_angle) % 180 + assert min(estimated, 180 - estimated) < 20 + + +class TestReconstructAllPermutations: + def test_returns_one_image_per_kernel(self, recon): + images = recon._reconstruct_all_permutations(verbose=False) + + assert len(images) == len(DECONVOLUTION_KERNELS) + assert all(image.shape == (N, N) for image in images) + assert all(np.isfinite(image).all() for image in images) + + +class TestNormalizationOrder: + def test_linear_background_normalization_runs(self, dataset4d): + recon = _build(dataset4d, normalization_order=1) + recon.reconstruct(deconvolution_kernel="ssb", verbose=False) + + assert np.isfinite(recon.obj).all() + assert abs(correlation(recon.obj, band_limited_phase())) > 0.5 + + def test_rejects_unknown_order(self, dataset4d): + with pytest.raises(ValueError, match="normalization_order"): + _build(dataset4d, normalization_order=2) + + def test_edge_blending_tapers_the_stack(self, dataset4d): + """A nonzero blend pulls the scan-edge vBF values toward unity.""" + blended = _build(dataset4d, edge_blend_pixels=4) + sharp = _build(dataset4d, edge_blend_pixels=0) + + edge_blended = blended.vbf_stack[:, 0, :] + edge_sharp = sharp.vbf_stack[:, 0, :] + + assert (edge_blended - 1).abs().mean() < (edge_sharp - 1).abs().mean() + + +class TestDatasetVariants: + def test_reconstruction_tracks_the_simulated_defocus(self): + """Data simulated at a different defocus must prefer that defocus.""" + other_C10 = 2 * TRUE_C10 + dataset = make_dataset4d(defocus=other_C10) + recon = DirectPtychography.from_dataset4d( + dataset, edge_blend_pixels=0, **direct_ptycho_kwargs(other_C10) + ) + + losses = {} + for scale in (0.5, 1.0, 1.5): + recon.reconstruct( + deconvolution_kernel="prlx", + override_aberration_coefs={"C10": scale * other_C10}, + parallax_flip_phase=False, + verbose=False, + ) + losses[scale] = float(recon.variance_loss()) + + assert min(losses, key=losses.get) == pytest.approx(1.0) + + +class TestFromDataset3d: + """Ungridded scans, by resampling the bright-field stack onto a grid first. + + Every kernel here is a scan-space Fourier multiplier and so needs a regular grid. Once + regridded they all run unchanged, which is what makes this exact for positions that were + already on a lattice. + """ + + @staticmethod + def _dataset3d(dataset4d): + return Dataset3d.from_array( + np.asarray(dataset4d.array).reshape(-1, N, N), + name="ungridded patterns", + sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), + units=("index", "A^-1", "A^-1"), + ) + + @staticmethod + def _build(dataset3d, positions, **overrides): + kwargs = dict( + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + aberration_coefs={"C10": TRUE_C10}, + force_fitted_origin=(N // 2, N // 2), + verbose=False, + ) + kwargs.update(overrides) + return DirectPtychography.from_dataset3d(dataset3d, positions, **kwargs) + + @pytest.mark.parametrize("kernel", DECONVOLUTION_KERNELS) + def test_lattice_positions_reproduce_from_dataset4d(self, dataset4d, kernel): + """On a lattice the splat is an identity map, so this must be exact.""" + gridded = _build(dataset4d) + ungridded = self._build(self._dataset3d(dataset4d), scan_positions_px() * SCAN_SAMPLING) + + assert ungridded.scan_gpts == gridded.scan_gpts + assert np.allclose( + ungridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, + gridded.reconstruct(deconvolution_kernel=kernel, verbose=False).obj, + atol=1e-6, + ) + + def test_position_axis_order_is_row_col(self, dataset4d): + """Swapping the columns must transpose the regridded stack, not scramble it. + + Checked with no aberrations, so the parallax shifts vanish. With a shift present the + relation does not hold: the shifts come from the detector k-grid, which transposing + the *positions* leaves alone. + """ + dataset3d = self._dataset3d(dataset4d) + positions = scan_positions_px() * SCAN_SAMPLING + flat = dict(aberration_coefs={}, scan_gpts=(N, N)) + + row_col = self._build(dataset3d, positions, **flat) + col_row = self._build(dataset3d, positions[:, ::-1].copy(), **flat) + + assert np.allclose( + to_numpy(col_row.vbf_stack), + to_numpy(row_col.vbf_stack).transpose(0, 2, 1), + atol=1e-5, + ) + + def test_jittered_positions_recover_the_object(self, dataset4d): + rng = np.random.default_rng(0) + positions = scan_positions_px() * SCAN_SAMPLING + positions = positions + rng.uniform(-0.3, 0.3, positions.shape) * SCAN_SAMPLING + + reconstruction = self._build(self._dataset3d(dataset4d), positions) + obj = reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + + assert correlation(obj[:N, :N], band_limited_phase()) > 0.5 + + def test_warns_when_positions_are_clustered(self, dataset4d): + """Enough positions to cover the grid, yet most of it empty: genuinely uneven.""" + positions = scan_positions_px() * SCAN_SAMPLING + corner = positions[(positions[:, 0] < 10) & (positions[:, 1] < 10)] + # more positions than grid pixels, but all piled onto 100 distinct spots + repeated = np.repeat(corner, 10, axis=0) + + with pytest.warns(UserWarning, match="clustered rather than merely sparse"): + DirectPtychography.from_dataset3d( + Dataset3d.from_array( + np.asarray(dataset4d.array).reshape(-1, N, N)[: len(repeated)], + name="clustered", + sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), + units=("index", "A^-1", "A^-1"), + ), + repeated, + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + scan_sampling=(SCAN_SAMPLING / 2, SCAN_SAMPLING / 2), + aberration_coefs={"C10": TRUE_C10}, + force_fitted_origin=(N // 2, N // 2), + verbose=False, + ) + + def test_a_finer_grid_does_not_warn_about_its_gaps(self, dataset4d): + """Fewer positions than pixels is a comb, not missing data.""" + positions = scan_positions_px() * SCAN_SAMPLING + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + self._build( + self._dataset3d(dataset4d), + positions, + scan_sampling=(SCAN_SAMPLING / 2, SCAN_SAMPLING / 2), + ) + + def test_no_warning_on_a_fully_covered_grid(self, dataset4d): + positions = scan_positions_px() * SCAN_SAMPLING + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + self._build(self._dataset3d(dataset4d), positions) + + def test_scan_gpts_pads_without_rescaling_the_positions(self, dataset4d): + """`scan_sampling` stays authoritative: a bigger grid pads, it does not rescale.""" + positions = scan_positions_px() * SCAN_SAMPLING + plain = self._build(self._dataset3d(dataset4d), positions) + padded = self._build(self._dataset3d(dataset4d), positions, scan_gpts=(N + 8, N + 8)) + + assert padded.scan_gpts == (N + 8, N + 8) + assert padded.scan_sampling[0] == pytest.approx(plain.scan_sampling[0]) + + def test_scan_gpts_too_small_raises(self, dataset4d): + positions = scan_positions_px() * SCAN_SAMPLING + with pytest.raises(ValueError, match="positions would be dropped"): + self._build(self._dataset3d(dataset4d), positions, scan_gpts=(N // 2, N // 2)) + + def test_auto_scan_sampling_warns_and_infers(self, dataset4d): + positions = scan_positions_px() * SCAN_SAMPLING + + with pytest.warns(UserWarning, match="Inferred scan_sampling"): + reconstruction = self._build( + self._dataset3d(dataset4d), positions, scan_sampling="auto" + ) + + assert reconstruction.scan_sampling[0] == pytest.approx(SCAN_SAMPLING) + + def test_accepts_a_dataset2d_of_positions(self, dataset4d): + positions = Dataset2d.from_array( + scan_positions_px() * SCAN_SAMPLING, + name="positions", + sampling=(1.0, 1.0), + units=("A", "A"), + ) + reconstruction = self._build(self._dataset3d(dataset4d), positions) + + assert reconstruction.scan_gpts == (N, N) + + def test_rejects_positions_in_the_wrong_units(self, dataset4d): + positions = Dataset2d.from_array( + scan_positions_px() * SCAN_SAMPLING, + name="positions", + sampling=(1.0, 1.0), + units=("nm", "nm"), + ) + with pytest.raises(ValueError, match="must be given in 'A'"): + self._build(self._dataset3d(dataset4d), positions) + + def test_rejects_mismatched_position_count(self, dataset4d): + with pytest.raises(ValueError, match="rows but `dataset` has"): + self._build(self._dataset3d(dataset4d), scan_positions_px()[:10] * SCAN_SAMPLING) + + def test_rejects_linear_normalization(self, dataset4d): + with pytest.raises(ValueError, match="needs a scan grid"): + self._build( + self._dataset3d(dataset4d), + scan_positions_px() * SCAN_SAMPLING, + normalization_order=1, + ) + + def test_survives_a_serialization_round_trip(self, dataset4d, tmp_path): + reconstruction = self._build( + self._dataset3d(dataset4d), scan_positions_px() * SCAN_SAMPLING + ) + reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False) + path = tmp_path / "ungridded.zip" + reconstruction.save(path, mode="o") + + assert np.allclose(load(path).obj, reconstruction.obj) + + @staticmethod + def _disk_masked(dataset4d, radius=14.0): + """A non-rectangular scan subset -- the case that exposes hole handling.""" + rows_cols = scan_positions_px() + center = (N - 1) / 2 + keep = ((rows_cols[:, 0] - center) ** 2 + (rows_cols[:, 1] - center) ** 2) < radius**2 + dataset3d = Dataset3d.from_array( + np.asarray(dataset4d.array).reshape(-1, N, N)[keep], + name="masked patterns", + sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), + units=("index", "A^-1", "A^-1"), + ) + return dataset3d, rows_cols[keep] * SCAN_SAMPLING, keep + + def test_mean_hole_fill_beats_zero_on_a_masked_scan(self, dataset4d): + """Regression guard: zero-filled holes wreck the reconstruction, mean-filled do not. + + `_preprocess` zeroes the DC bin, subtracting the mean over the whole grid including + holes, so zero-filled holes sit at `-mean` -- a hard-edged step the deconvolution + smears everywhere. Measured 0.25 vs 0.69 correlation with ground truth. + """ + dataset3d, positions, keep = self._disk_masked(dataset4d) + rows_cols = scan_positions_px()[keep].astype(int) + low, high = rows_cols.min(0), rows_cols.max(0) + 1 + truth = band_limited_phase()[low[0] : high[0], low[1] : high[1]] + + def score(hole_fill): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + recon = self._build(dataset3d, positions, hole_fill=hole_fill) + obj = recon.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + return correlation(obj[: truth.shape[0], : truth.shape[1]], truth) + + assert score("mean") > 0.6 + assert score("mean") > score("zero") + 0.2 + + def test_mean_fill_matches_the_montage_on_a_masked_scan(self, dataset4d): + """With holes filled, the two formulations agree on data neither was built for.""" + from quantem.diffractive_imaging import DirectPtychographyMontage + + dataset3d, positions, keep = self._disk_masked(dataset4d) + rows_cols = scan_positions_px()[keep].astype(int) + low, high = rows_cols.min(0), rows_cols.max(0) + 1 + truth = band_limited_phase()[low[0] : high[0], low[1] : high[1]] + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + fourier = self._build(dataset3d, positions) + montage = DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + aberration_coefs={"C10": TRUE_C10}, + force_fitted_origin=(N // 2, N // 2), + verbose=False, + ) + + fourier_obj = fourier.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + montage.reconstruct(deconvolution_kernel="prlx", weight_normalize=False, verbose=False) + origin = to_numpy(montage._canvas_origin_px) + row0, col0 = int(round(-origin[0])), int(round(-origin[1])) + montage_obj = montage.obj[row0 : row0 + truth.shape[0], col0 : col0 + truth.shape[1]] + + fourier_corr = correlation(fourier_obj[: truth.shape[0], : truth.shape[1]], truth) + montage_corr = correlation(montage_obj, truth) + + assert fourier_corr > 0.6 + assert abs(fourier_corr - montage_corr) < 0.1 + + def test_rejects_an_unknown_hole_fill(self, dataset4d): + dataset3d, positions, _ = self._disk_masked(dataset4d) + with pytest.raises(ValueError, match="`hole_fill` must be"): + self._build(dataset3d, positions, hole_fill="interpolate") + + def test_upsampling_an_irregular_regrid_warns(self, dataset4d): + """Regridding discards the sub-pixel positions upsampling needs to unfold.""" + rng = np.random.default_rng(3) + positions = scan_positions_px() * SCAN_SAMPLING + positions = positions + rng.uniform(-0.4, 0.4, positions.shape) * SCAN_SAMPLING + recon = self._build(self._dataset3d(dataset4d), positions) + + assert recon._regrid_info["lattice_rms_px"] > 0.1 + with pytest.warns(UserWarning, match="cannot unfold"): + recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) + + def test_no_unfolding_warning_on_a_lattice_or_without_upsampling(self, dataset4d): + dataset3d = self._dataset3d(dataset4d) + on_lattice = self._build(dataset3d, scan_positions_px() * SCAN_SAMPLING) + + assert on_lattice._regrid_info["lattice_rms_px"] == pytest.approx(0.0, abs=1e-9) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + on_lattice.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) + + rng = np.random.default_rng(3) + positions = scan_positions_px() * SCAN_SAMPLING + positions = positions + rng.uniform(-0.4, 0.4, positions.shape) * SCAN_SAMPLING + irregular = self._build(dataset3d, positions) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + irregular.reconstruct(deconvolution_kernel="prlx", verbose=False) + + def test_gridded_reconstructions_never_warn_about_unfolding(self, dataset4d): + """`from_dataset4d` has no regrid info, so the check must be a no-op there.""" + gridded = _build(dataset4d) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + gridded.reconstruct(deconvolution_kernel="prlx", upsampling_factor=4, verbose=False) + + +class TestUnfolding: + """Does `upsampling_factor` extend the CTF, or merely replicate it? + + Uses the white-noise object, whose Fourier amplitude is constant, so + `|FFT(reconstruction)|` is the contrast transfer function and can be compared against an + analytical one. The scan pitch is half the aperture's transfer limit, so there is a + genuine factor of two above the scan Nyquist to recover -- without that, an upsampling + test measures nothing. + """ + + @staticmethod + def _scores(obj, analytic, upsampling_factor): + return ctf_band_scores(measured_ctf(obj), analytic, CTF_SCAN_SAMPLING / upsampling_factor) + + @staticmethod + def _ungridded(scene, positions, **overrides): + _, _, probe, _, _ = scene + kwargs = dict( + ctf_kwargs(), + scan_sampling=(CTF_SCAN_SAMPLING, CTF_SCAN_SAMPLING), + ) + kwargs.update(overrides) + return DirectPtychography.from_dataset3d( + ctf_dataset3d(ctf_simulate(scene[0], probe, positions)), + positions * CTF_SAMPLING, + **kwargs, + ) + + def test_gridded_upsampling_unfolds_the_ctf(self, ctf_scene): + """The baseline the ungridded path is judged against.""" + obj, _, probe, ctf_full, ctf_sub = ctf_scene + recon = DirectPtychography.from_dataset4d( + ctf_dataset4d(ctf_simulate(obj, probe, ctf_raster_positions())), + edge_blend_pixels=0, + **ctf_kwargs(), + ) + + in_band, _ = self._scores( + recon.reconstruct(deconvolution_kernel="prlx", verbose=False).obj, ctf_sub, 1 + ) + up = recon.reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False).obj + up_in, up_ext = self._scores(up, ctf_full, 2) + + assert in_band > 0.95 + assert up_in > 0.95 + assert up_ext > 0.95 # genuinely extended, not replicated + + def test_lattice_positions_unfold_through_the_ungridded_path(self, ctf_scene): + _, _, _, ctf_full, _ = ctf_scene + recon = self._ungridded(ctf_scene, ctf_raster_positions()) + + obj = recon.reconstruct( + deconvolution_kernel="prlx", upsampling_factor=2, verbose=False + ).obj + + assert self._scores(obj, ctf_full, 2)[1] > 0.95 + + def test_upsampling_factor_degrades_on_an_irregular_scan(self, ctf_scene): + """Binning first discards the sub-pixel positions upsampling needs. + + Compared against the same call on a lattice rather than an absolute threshold: what + matters is that scattering the positions costs the extension band, not where the + number happens to land. + """ + _, _, _, ctf_full, _ = ctf_scene + + on_lattice = ( + self._ungridded(ctf_scene, ctf_raster_positions()) + .reconstruct(deconvolution_kernel="prlx", upsampling_factor=2, verbose=False) + .obj + ) + scattered_recon = self._ungridded(ctf_scene, ctf_jittered_positions(1.0)) + with pytest.warns(UserWarning, match="cannot unfold"): + scattered = scattered_recon.reconstruct( + deconvolution_kernel="prlx", upsampling_factor=2, verbose=False + ).obj + + assert ( + self._scores(scattered, ctf_full, 2)[1] + < self._scores(on_lattice, ctf_full, 2)[1] - 0.1 + ) + + def test_a_finer_scan_sampling_recovers_it(self, ctf_scene): + """Splatting straight onto a finer grid keeps the probes on it.""" + _, _, _, ctf_full, _ = ctf_scene + positions = ctf_jittered_positions(1.0) + + binned = self._ungridded(ctf_scene, positions) + with pytest.warns(UserWarning, match="cannot unfold"): + coarse = binned.reconstruct( + deconvolution_kernel="prlx", upsampling_factor=2, verbose=False + ).obj + + refined = self._ungridded(ctf_scene, positions, scan_sampling=(CTF_SCAN_SAMPLING / 2,) * 2) + fine = refined.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + + fine_corr = ctf_radial_correlation(measured_ctf(fine), ctf_full, refined.scan_sampling[0]) + coarse_corr = ctf_radial_correlation( + measured_ctf(coarse), ctf_full, binned.scan_sampling[0] / 2 + ) + assert fine_corr > 0.99 + assert fine_corr > coarse_corr + + def test_a_finer_scan_sampling_refines_the_grid(self, ctf_scene): + plain = self._ungridded(ctf_scene, ctf_raster_positions()) + finer = self._ungridded( + ctf_scene, ctf_raster_positions(), scan_sampling=(CTF_SCAN_SAMPLING / 2,) * 2 + ) + + assert finer.scan_sampling[0] == pytest.approx(plain.scan_sampling[0] / 2) + assert finer.scan_gpts[0] >= 2 * plain.scan_gpts[0] - 2 + + def test_nearest_beats_bilinear_on_a_finer_grid(self, ctf_scene): + """Bilinear smears each measurement over four pixels and blurs out the detail.""" + _, _, _, ctf_full, _ = ctf_scene + positions = ctf_jittered_positions(1.0) + fine = dict(scan_sampling=(CTF_SCAN_SAMPLING / 2,) * 2) + + scores = {} + for scheme in ("nearest", "bilinear"): + recon = self._ungridded(ctf_scene, positions, interpolation=scheme, **fine) + obj = recon.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + scores[scheme] = ctf_radial_correlation( + measured_ctf(obj), ctf_full, recon.scan_sampling[0] + ) + + assert scores["nearest"] > scores["bilinear"] + + def test_zero_fill_inverts_a_finer_grid(self, ctf_scene): + """`hole_fill="mean"` centers the comb's gaps; leaving them at zero inverts it.""" + _, _, _, ctf_full, _ = ctf_scene + positions = ctf_jittered_positions(1.0) + fine = dict(scan_sampling=(CTF_SCAN_SAMPLING / 2,) * 2) + + scores = {} + for fill in ("mean", "zero"): + recon = self._ungridded(ctf_scene, positions, hole_fill=fill, **fine) + obj = recon.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + scores[fill] = ctf_radial_correlation( + measured_ctf(obj), ctf_full, recon.scan_sampling[0] + ) + + assert scores["mean"] > 0.99 + assert scores["zero"] < scores["mean"] - 0.1 + + def test_montage_unfolds_without_any_regridding(self, ctf_scene): + """The reference: the montage never bins, so irregularity costs it nothing.""" + from quantem.diffractive_imaging import DirectPtychographyMontage + + obj, _, probe, ctf_full, _ = ctf_scene + positions = ctf_jittered_positions(1.0) + montage = DirectPtychographyMontage.from_dataset3d( + ctf_dataset3d(ctf_simulate(obj, probe, positions)), + positions * CTF_SAMPLING, + scan_sampling=(CTF_SCAN_SAMPLING, CTF_SCAN_SAMPLING), + boundary="wrap", + **ctf_kwargs(), + ) + + result = montage.reconstruct( + deconvolution_kernel="prlx", upsampling_factor=2, verbose=False + ).obj + + assert self._scores(result, ctf_full, 2)[1] > 0.95 + + +class TestOptimizationParameterIsOneClass: + """`OptimizationParameter` used to be defined twice, identically. + + One copy lived in ``direct_ptychography`` and the other in + ``optimize_hyperparameters``. Both searches test candidate specifications with + ``isinstance``, so a value built from one module was silently ignored by the other + rather than raising. These pin the single definition in place. + """ + + def test_every_import_path_is_the_same_class(self): + from quantem.diffractive_imaging import OptimizationParameter as package + from quantem.diffractive_imaging.direct_ptychography import ( + OptimizationParameter as direct, + ) + from quantem.diffractive_imaging.direct_ptychography_base import ( + OptimizationParameter as base, + ) + from quantem.diffractive_imaging.optimize_hyperparameters import ( + OptimizationParameter as iterative, + ) + from quantem.diffractive_imaging.ptycho_utils import ( + OptimizationParameter as canonical, + ) + + assert package is direct is base is iterative is canonical + assert canonical.__module__ == "quantem.diffractive_imaging.ptycho_utils" + + def test_a_spec_built_anywhere_is_accepted_everywhere(self): + """The failure the duplicate caused: isinstance across the two searches.""" + from quantem.diffractive_imaging import OptimizationParameter + from quantem.diffractive_imaging.optimize_hyperparameters import ( + OptimizationParameter as iterative, + ) + + assert isinstance(OptimizationParameter(0.0, 1.0), iterative) diff --git a/tests/diffractive_imaging/test_direct_ptychography_montage.py b/tests/diffractive_imaging/test_direct_ptychography_montage.py new file mode 100644 index 00000000..cdd00a79 --- /dev/null +++ b/tests/diffractive_imaging/test_direct_ptychography_montage.py @@ -0,0 +1,2530 @@ +"""Tests for the real-space (shadow montage) direct-ptychography reconstruction. + +The headline check is equivalence with the Fourier-space parallax kernel of +``DirectPtychography``: the two are the same linear operator written in different domains, +so on a raster scan with periodic wraparound they must agree. +""" + +import numpy as np +import pytest +import torch + +from quantem.core.datastructures import Dataset2d, Dataset3d +from quantem.core.io.serialize import load +from quantem.core.utils.utils import to_numpy +from quantem.diffractive_imaging import ( + DirectPtychography, + DirectPtychographyMontage, + OptimizationParameter, +) +from quantem.diffractive_imaging.complex_probe import FourierProbe, spatial_frequencies +from quantem.diffractive_imaging.direct_ptycho_utils import ( + allocate_splat_buffers, + estimate_frame_drift, + scatter_add_convolve, + scatter_add_splat, + splat_and_convolve, + splat_stack, +) + +from .conftest import ( + ACCELERATORS, + CTF_SAMPLING, + CTF_SCAN_SAMPLING, + ORIGIN, + PROBE_ENERGY, + RECIPROCAL_SAMPLING, + SCAN_SAMPLING, + SEMIANGLE_CUTOFF, + N, + analytic_probe_array, + band_limited_phase, + correlation, + ctf_interleaved_frames, + ctf_kwargs, + integer_shift_defocus, + make_model_vbf_stack, + make_tilted_dataset4d, + model_vbf_kwargs, + scan_positions_px, +) +from .conftest import ( + direct_ptycho_kwargs as _common_kwargs, +) + +#: defocus and scan size at which the per-patch estimator is well conditioned; see +#: `make_model_vbf_stack` for why the 32x32 4D fixture is not +MODEL_DEFOCUS = 3000.0 +MODEL_SCAN_GPTS = (96, 96) +MODEL_C10_GRID = np.linspace(1500.0, 4500.0, 13) + + +def _model_montage(defocus_gradient, defocus=MODEL_DEFOCUS): + """A montage over a model vBF stack with a seeded defocus plane, plus the ground truth.""" + vbf, bf_mask, obj = make_model_vbf_stack(defocus, defocus_gradient, scan_gpts=MODEL_SCAN_GPTS) + montage = DirectPtychographyMontage.from_virtual_bfs(vbf, bf_mask, **model_vbf_kwargs(defocus)) + return montage, obj + + +def _build_pair(dataset4d, defocus): + """A `DirectPtychography` and a `DirectPtychographyMontage` over the same data.""" + fourier = DirectPtychography.from_dataset4d( + dataset4d, edge_blend_pixels=0, **_common_kwargs(defocus) + ) + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, edge_blend_pixels=0, boundary="wrap", **_common_kwargs(defocus) + ) + return fourier, montage + + +def _relative_error(a, b): + return float(np.abs(a - b).max() / np.abs(b).max()) + + +class TestSplatKernel: + """Unit tests for `scatter_add_splat`, independent of any reconstruction.""" + + def test_integer_coordinates_deposit_one_pixel(self): + values = torch.tensor([[2.0]]) + coords = torch.tensor([[[3.0, 4.0]]]) + sum_w, sum_wv, sum_wv2 = scatter_add_splat(values, coords, (8, 8)) + + assert sum_w.sum().item() == pytest.approx(1.0) + assert sum_w.reshape(8, 8)[3, 4].item() == pytest.approx(1.0) + assert sum_wv.reshape(8, 8)[3, 4].item() == pytest.approx(2.0) + assert sum_wv2.reshape(8, 8)[3, 4].item() == pytest.approx(4.0) + + def test_half_pixel_splits_four_ways(self): + values = torch.tensor([[1.0]]) + coords = torch.tensor([[[3.5, 4.5]]]) + sum_w, _, _ = scatter_add_splat(values, coords, (8, 8)) + + nonzero = sum_w[sum_w > 0] + assert nonzero.numel() == 4 + assert torch.allclose(nonzero, torch.full((4,), 0.25, dtype=nonzero.dtype)) + + def test_wrap_is_periodic(self): + values = torch.tensor([[1.0]]) + coords = torch.tensor([[[-0.5, 8.25]]]) + sum_w, _, _ = scatter_add_splat(values, coords, (8, 8), boundary="wrap") + + assert sum_w.sum().item() == pytest.approx(1.0) + touched = {tuple(ij) for ij in torch.nonzero(sum_w.reshape(8, 8)).tolist()} + assert touched == {(7, 0), (7, 1), (0, 0), (0, 1)} + + def test_pad_drops_out_of_bounds_without_corrupting_the_edge(self): + values = torch.tensor([[100.0, 2.0]]) + coords = torch.tensor([[[-5.0, 4.0], [3.0, 4.0]]]) + sum_w, sum_wv, _ = scatter_add_splat(values, coords, (8, 8), boundary="pad") + + # only the in-bounds point contributes, and the clamped row picks up nothing + assert sum_w.sum().item() == pytest.approx(1.0) + assert sum_wv.sum().item() == pytest.approx(2.0) + assert sum_w.reshape(8, 8)[0].max().item() == pytest.approx(0.0) + + def test_weights_are_a_partition_of_unity(self): + generator = torch.Generator().manual_seed(0) + values = torch.randn(4, 100, generator=generator) + coords = torch.rand(4, 100, 2, generator=generator) * 8 + sum_w, _, _ = scatter_add_splat(values, coords, (8, 8), boundary="wrap") + + assert sum_w.sum().item() == pytest.approx(400.0) + + def test_matches_brute_force(self): + generator = torch.Generator().manual_seed(1) + values = torch.randn(3, 50, generator=generator) + coords = torch.rand(3, 50, 2, generator=generator) * 10 - 1 + + sum_w, sum_wv, sum_wv2 = scatter_add_splat(values, coords, (10, 10), boundary="pad") + + # the splat accumulates in float64, so the reference must too + values_np = to_numpy(values).astype(np.float64) + coords_np = to_numpy(coords).astype(np.float64) + ref_w, ref_wv, ref_wv2 = (np.zeros(100) for _ in range(3)) + for b in range(3): + for t in range(50): + r0 = int(np.floor(coords_np[b, t, 0])) + c0 = int(np.floor(coords_np[b, t, 1])) + fr = coords_np[b, t, 0] - r0 + fc = coords_np[b, t, 1] - c0 + for d_row, d_col in ((0, 0), (1, 0), (0, 1), (1, 1)): + row, col = r0 + d_row, c0 + d_col + if not (0 <= row < 10 and 0 <= col < 10): + continue + weight = (fr if d_row else 1 - fr) * (fc if d_col else 1 - fc) + value = values_np[b, t] + ref_w[row * 10 + col] += weight + ref_wv[row * 10 + col] += weight * value + ref_wv2[row * 10 + col] += weight * value**2 + + assert np.allclose(to_numpy(sum_w), ref_w) + assert np.allclose(to_numpy(sum_wv), ref_wv) + assert np.allclose(to_numpy(sum_wv2), ref_wv2) + + def test_nearest_rounds(self): + values = torch.tensor([[1.0]]) + coords = torch.tensor([[[3.4, 4.6]]]) + sum_w, _, _ = scatter_add_splat(values, coords, (8, 8), interpolation="nearest") + + assert torch.nonzero(sum_w.reshape(8, 8)).tolist() == [[3, 5]] + + def test_accumulates_into_provided_buffers(self): + buffers = allocate_splat_buffers((8, 8), "cpu") + values = torch.tensor([[1.0]]) + coords = torch.tensor([[[3.0, 4.0]]]) + + for _ in range(3): + scatter_add_splat(values, coords, (8, 8), out=buffers) + + assert buffers[0].reshape(8, 8)[3, 4].item() == pytest.approx(3.0) + + def test_rejects_unknown_modes(self): + values = torch.tensor([[1.0]]) + coords = torch.tensor([[[3.0, 4.0]]]) + with pytest.raises(ValueError, match="boundary"): + scatter_add_splat(values, coords, (8, 8), boundary="reflect") + with pytest.raises(ValueError, match="interpolation"): + scatter_add_splat(values, coords, (8, 8), interpolation="cubic") + + def test_convolve_with_a_unit_tap_is_a_plain_splat(self): + values = torch.tensor([[1.0, 2.0]]) + coords = torch.tensor([[[3.0, 4.0], [5.0, 6.0]]]) + + out = scatter_add_convolve( + values, + coords, + (8, 8), + torch.tensor([[0, 0]]), + torch.ones(1, 1, dtype=torch.complex64), + ).reshape(8, 8) + + assert out[3, 4].item() == pytest.approx(1.0) + assert out[5, 6].item() == pytest.approx(2.0) + + def test_convolve_places_each_tap_at_its_offset(self): + offsets = torch.tensor([[0, 0], [1, 0], [0, -2]]) + weights = torch.tensor([[1.0, 2.0j, -3.0]], dtype=torch.complex64) + + out = scatter_add_convolve( + torch.tensor([[1.0]]), torch.tensor([[[4.0, 4.0]]]), (8, 8), offsets, weights + ).reshape(8, 8) + + assert out[4, 4].item() == pytest.approx(1.0) + assert out[5, 4].item() == pytest.approx(2.0j) + assert out[4, 2].item() == pytest.approx(-3.0) + assert int((out != 0).sum()) == 3 + + def test_convolve_wraps_taps_at_the_boundary(self): + out = scatter_add_convolve( + torch.tensor([[1.0]]), + torch.tensor([[[0.0, 0.0]]]), + (8, 8), + torch.tensor([[-1, -1]]), + torch.ones(1, 1, dtype=torch.complex64), + boundary="wrap", + ).reshape(8, 8) + + assert out[7, 7].item() == pytest.approx(1.0) + + +class TestIntegerShiftConstruction: + """The defocus used below must put every BF pixel on an exact canvas pixel.""" + + @pytest.mark.parametrize("upsampling_factor", [1, 2]) + @pytest.mark.parametrize("pixel_shift", [1, 2]) + def test_shifts_land_on_pixel_centers(self, dataset4d, upsampling_factor, pixel_shift): + defocus = integer_shift_defocus(pixel_shift, upsampling_factor) + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, boundary="wrap", **_common_kwargs(defocus) + ) + shifts, _ = montage._return_shifts_px( + 0.0, {"C10": defocus}, montage.bf_mask, upsampling_factor + ) + residual = (shifts - shifts.round()).abs().max().item() + + assert residual < 1e-4, f"shifts are not integral: max residual {residual:.2e} px" + assert shifts.abs().max().item() > 1.0, "test would be vacuous with zero shifts" + + +class TestFourierEquivalence: + """`DirectPtychographyMontage` must reproduce `DirectPtychography`'s parallax kernel.""" + + @pytest.mark.parametrize("upsampling_factor", [1, 2]) + @pytest.mark.parametrize("pixel_shift", [1, 2]) + def test_matches_parallax_kernel(self, dataset4d, upsampling_factor, pixel_shift): + defocus = integer_shift_defocus(pixel_shift, upsampling_factor) + fourier, montage = _build_pair(dataset4d, defocus) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + upsampling_factor=upsampling_factor, + verbose=False, + ) + fourier.reconstruct(**recon_kwargs) + montage.reconstruct( + boundary="wrap", interpolation="bilinear", weight_normalize=False, **recon_kwargs + ) + + assert montage.obj.shape == fourier.obj.shape + assert _relative_error(montage.obj, fourier.obj) < 1e-4 + + def test_matches_with_phase_flip(self, dataset4d): + """The phase-flip filter is BF-independent, so post-hoc application is exact.""" + defocus = integer_shift_defocus(1) + fourier, montage = _build_pair(dataset4d, defocus) + + fourier.reconstruct(deconvolution_kernel="prlx", parallax_flip_phase=True, verbose=False) + montage.reconstruct( + deconvolution_kernel="prlx", + parallax_flip_phase=True, + weight_normalize=False, + verbose=False, + ) + + assert _relative_error(montage.obj, fourier.obj) < 1e-4 + + def test_matches_with_butterworth_filters(self, dataset4d): + defocus = integer_shift_defocus(1) + fourier, montage = _build_pair(dataset4d, defocus) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + q_lowpass=0.2, + q_highpass=0.02, + verbose=False, + ) + fourier.reconstruct(**recon_kwargs) + montage.reconstruct(weight_normalize=False, **recon_kwargs) + + assert _relative_error(montage.obj, fourier.obj) < 1e-4 + + def test_weight_normalization_differs_only_by_a_constant(self, dataset4d): + """On a full grid at U=1 the accumulated weight is exactly num_bf everywhere.""" + defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, defocus) + + montage.reconstruct( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=False, + verbose=False, + ) + unnormalized = montage.obj + bf_weights = float(montage._bf_weights) + + montage.reconstruct( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=True, + verbose=False, + ) + normalized = montage.obj + + assert np.allclose(to_numpy(montage.weights), montage.num_bf, rtol=1e-6) + rescaled = normalized * montage.num_bf / bf_weights + assert _relative_error(rescaled, unnormalized) < 1e-4 + + def test_variance_loss_tracks_the_fourier_one(self, dataset4d): + defocus = integer_shift_defocus(1) + fourier, montage = _build_pair(dataset4d, defocus) + + recon_kwargs = dict(deconvolution_kernel="prlx", parallax_flip_phase=False, verbose=False) + fourier.reconstruct(**recon_kwargs) + montage.reconstruct(weight_normalize=False, **recon_kwargs) + + bf_weights = float(montage._bf_weights) + expected = float(fourier.variance_loss()) * bf_weights**2 + assert float(montage.variance_loss()) == pytest.approx(expected, rel=1e-3) + + def test_bf_mask_subsets_are_additive(self, dataset4d): + """Checkerboard half-sets, rescaled by their BF weights, sum to the whole.""" + defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, defocus) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=False, + verbose=False, + ) + montage.reconstruct(**recon_kwargs) + full = montage.obj * float(montage._bf_weights) + + halves = [] + for mask in montage._make_checkerboard_bf_masks(montage.gpts, montage.bf_mask): + montage.reconstruct(bf_mask=mask, **recon_kwargs) + halves.append(montage.obj * float(montage._bf_weights)) + + assert _relative_error(halves[0] + halves[1], full) < 1e-4 + + @pytest.mark.parametrize("kernel", ["ssb", "obf", "mf"]) + def test_accepts_the_deconvolution_kernels(self, dataset4d, kernel): + """Available as truncated real-space convolutions; see `TestRealSpaceKernels`.""" + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + obj = montage.reconstruct(deconvolution_kernel=kernel, stencil_radius=6, verbose=False).obj + + assert np.isfinite(obj).all() + + def test_icom_matches_the_fourier_class(self, dataset4d): + """iCoM has no compact real-space form, but the FFT route never truncates.""" + fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + fourier.reconstruct(deconvolution_kernel="icom", verbose=False) + montage.reconstruct(deconvolution_kernel="icom", convolution_mode="fft", verbose=False) + + assert _relative_error(montage.obj, fourier.obj) < 1e-5 + + +class TestRotationConvention: + """`spatial_frequencies` rotates the k-grid; shifts must follow it, unflipped.""" + + def test_ninety_degrees_rotates_the_shifts(self, dataset4d): + defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, defocus) + coefs = {"C10": defocus} + + unrotated, _ = montage._return_shifts_px(0.0, coefs, montage.bf_mask, 1) + rotated, _ = montage._return_shifts_px(90.0, coefs, montage.bf_mask, 1) + + # _passively_rotate_grid sends (kx, ky) -> (kx cos a - ky sin a, kx sin a + ky cos a), + # so at 90 deg the shifts, which are parallel to k for pure defocus, map (r, c) -> (-c, r) + expected = torch.stack((-unrotated[:, 1], unrotated[:, 0]), dim=-1) + assert torch.allclose(rotated, expected, atol=1e-4) + + def test_rotation_changes_the_reconstruction(self, dataset4d): + defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, defocus) + + montage.reconstruct(override_rotation_angle=0.0, verbose=False) + unrotated = montage.obj.copy() + montage.reconstruct(override_rotation_angle=30.0, verbose=False) + + assert not np.allclose(unrotated, montage.obj) + + +class TestReconstructDefaults: + """The defaults are behaviour; pin them.""" + + def test_interpolation_defaults_to_nearest(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + default = montage.reconstruct(upsampling_factor=2, verbose=False).obj.copy() + nearest = montage.reconstruct( + upsampling_factor=2, interpolation="nearest", verbose=False + ).obj + + assert np.array_equal(default, nearest) + + def test_nearest_is_a_roll_of_the_bright_field_images(self, dataset4d): + """On a raster scan, snapping moves every position of a BF image by one integer. + + `positions_px * U` is an exact integer, so `round(n + s) == n + round(s)`. + """ + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + upsampling_factor = 4 + shifts, _ = montage._return_shifts_px( + 0.0, montage.aberration_coefs, montage.bf_mask, upsampling_factor + ) + positions = montage.positions_px * upsampling_factor + coords = positions[None] + shifts[:, None] + + offsets = coords.round() - positions[None] + # every position of a given BF image moves by the same integer + assert torch.equal(offsets, shifts.round()[:, None].expand_as(offsets)) + + def test_gridded_constructors_flag_the_scan_as_gridded(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + assert montage.gridded_scan is True + + def test_weight_normalize_defaults_off_for_a_raster_scan(self, dataset4d): + """Uniform density needs no correction, and normalizing amplifies edge noise.""" + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + default = montage.reconstruct(boundary="pad", verbose=False).obj.copy() + unnormalized = montage.reconstruct( + boundary="pad", weight_normalize=False, verbose=False + ).obj + + assert np.array_equal(default, unnormalized) + + def test_weight_normalize_defaults_on_for_an_ungridded_scan(self, dataset4d): + dataset3d, positions = TestNonGridScan._dataset3d_and_positions(dataset4d) + recon = DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + **_common_kwargs(integer_shift_defocus(1)), + ) + + assert recon.gridded_scan is False + default = recon.reconstruct(verbose=False).obj.copy() + normalized = recon.reconstruct(weight_normalize=True, verbose=False).obj + assert np.array_equal(default, normalized) + + +class TestPadBoundary: + """`boundary="pad"` grows the canvas instead of wrapping.""" + + def test_interior_matches_wrap(self, dataset4d): + defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, defocus) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=False, + verbose=False, + ) + montage.reconstruct(boundary="wrap", **recon_kwargs) + wrapped = montage.obj + + shifts, _ = montage._return_shifts_px(0.0, {"C10": defocus}, montage.bf_mask, 1) + margin = int(np.ceil(float(shifts.abs().max()))) + 1 + + montage.reconstruct(boundary="pad", **recon_kwargs) + padded = montage.obj + row0, col0 = (-to_numpy(montage._canvas_origin_px)).astype(int) + + # far enough from every edge, the wrap modulo is a no-op and the two agree exactly + n_rows, n_cols = wrapped.shape + interior_wrap = wrapped[margin : n_rows - margin, margin : n_cols - margin] + interior_pad = padded[ + row0 + margin : row0 + n_rows - margin, col0 + margin : col0 + n_cols - margin + ] + + assert interior_pad.shape == interior_wrap.shape + assert _relative_error(interior_pad, interior_wrap) < 1e-5 + + def test_canvas_covers_the_shifted_positions(self, dataset4d): + defocus = integer_shift_defocus(2) + _, montage = _build_pair(dataset4d, defocus) + montage.reconstruct(boundary="pad", verbose=False) + + # nothing was dropped: every (BF pixel, position) pair landed on the canvas + assert float(montage.weights.sum()) == pytest.approx( + montage.num_bf * montage.num_positions, rel=1e-6 + ) + assert montage.obj.shape[0] > N and montage.obj.shape[1] > N + + def test_pad_px_freezes_the_canvas(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + shapes = set() + for pixel_shift in (1, 2): + montage.reconstruct( + override_aberration_coefs={"C10": integer_shift_defocus(pixel_shift)}, + boundary="pad", + pad_px=12, + verbose=False, + ) + shapes.add(montage.obj.shape) + + assert len(shapes) == 1, f"canvas resized across trials: {shapes}" + + +class TestNonGridScan: + """`from_dataset3d` with raster positions must reproduce the gridded path.""" + + @staticmethod + def _dataset3d_and_positions(dataset4d): + n_scan = dataset4d.shape[0] + dataset3d = Dataset3d.from_array( + dataset4d.array.reshape(-1, N, N), + name="synthetic 3D stack", + units=("index", "A^-1", "A^-1"), + sampling=(1, RECIPROCAL_SAMPLING, RECIPROCAL_SAMPLING), + ) + positions = scan_positions_px()[: n_scan * n_scan] * SCAN_SAMPLING + return dataset3d, positions + + def test_matches_from_dataset4d(self, dataset4d): + defocus = integer_shift_defocus(1) + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + + gridded = DirectPtychographyMontage.from_dataset4d( + dataset4d, edge_blend_pixels=0, boundary="wrap", **_common_kwargs(defocus) + ) + ungridded = DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + boundary="wrap", + **_common_kwargs(defocus), + ) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=False, + verbose=False, + ) + gridded.reconstruct(**recon_kwargs) + ungridded.reconstruct(**recon_kwargs) + + assert ungridded.obj.shape == gridded.obj.shape + assert _relative_error(ungridded.obj, gridded.obj) < 1e-5 + + def test_position_axis_order_is_row_col(self, dataset4d): + """Swapping the position columns must transpose the reconstruction.""" + defocus = integer_shift_defocus(1) + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + + recon_kwargs = dict( + deconvolution_kernel="prlx", + parallax_flip_phase=False, + weight_normalize=False, + verbose=False, + ) + common = dict( + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + boundary="wrap", + **_common_kwargs(defocus), + ) + + straight = DirectPtychographyMontage.from_dataset3d(dataset3d, positions, **common) + swapped = DirectPtychographyMontage.from_dataset3d( + dataset3d, positions[:, ::-1].copy(), **common + ) + straight.reconstruct(**recon_kwargs) + swapped.reconstruct(**recon_kwargs) + + assert not np.allclose(straight.obj, swapped.obj) + + def test_scattered_positions_reconstruct(self, dataset4d): + """A jittered, shuffled scan still produces a supported montage.""" + defocus = integer_shift_defocus(1) + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + + rng = np.random.default_rng(0) + order = rng.permutation(len(positions)) + jittered = positions[order] + rng.normal( + scale=0.25 * SCAN_SAMPLING, size=(len(positions), 2) + ) + shuffled = Dataset3d.from_array( + dataset3d.array[order], + name="shuffled", + units=dataset3d.units, + sampling=dataset3d.sampling, + ) + + recon = DirectPtychographyMontage.from_dataset3d( + shuffled, + jittered, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + boundary="pad", + **_common_kwargs(defocus), + ) + recon.reconstruct(parallax_flip_phase=False, verbose=False) + + assert np.isfinite(recon.obj).all() + assert float(recon.weights.max()) > 0 + + def test_auto_scan_sampling_warns_and_infers(self, dataset4d): + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + + with pytest.warns(UserWarning, match="Inferred scan_sampling"): + recon = DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + scan_sampling="auto", + **_common_kwargs(integer_shift_defocus(1)), + ) + + assert recon.scan_sampling[0] == pytest.approx(SCAN_SAMPLING, rel=1e-6) + + def test_accepts_a_dataset2d_of_positions(self, dataset4d): + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + positions_dataset = Dataset2d.from_array(positions, name="positions", units=("A", "A")) + + recon = DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions_dataset, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + **_common_kwargs(integer_shift_defocus(1)), + ) + assert recon.num_positions == len(positions) + + def test_rejects_positions_in_the_wrong_units(self, dataset4d): + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + positions_dataset = Dataset2d.from_array( + positions, name="positions", units=("pixels", "pixels") + ) + + with pytest.raises(ValueError, match="must be given in 'A'"): + DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions_dataset, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + **_common_kwargs(integer_shift_defocus(1)), + ) + + def test_rejects_mismatched_position_count(self, dataset4d): + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + + with pytest.raises(ValueError, match="rows but `dataset` has"): + DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions[:-1], + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + **_common_kwargs(integer_shift_defocus(1)), + ) + + def test_requires_an_explicit_rotation_angle(self, dataset4d): + dataset3d, positions = self._dataset3d_and_positions(dataset4d) + kwargs = _common_kwargs(integer_shift_defocus(1)) + kwargs["rotation_angle"] = None + + with pytest.raises(ValueError, match="must be given for non-raster scans"): + DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + **kwargs, + ) + + +class TestHyperparameterSearch: + def test_grid_search_recovers_the_seeded_defocus(self, dataset4d): + true_defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, true_defocus) + + montage.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter( + low=0.4 * true_defocus, high=1.6 * true_defocus, n_points=7 + ) + }, + parallax_flip_phase=False, + verbose=False, + ) + best = montage.hyperparameter_state.optimized_aberrations["C10"] + + step = (1.6 - 0.4) * true_defocus / 6 + assert abs(best - true_defocus) <= step + + def test_variance_loss_is_minimized_at_the_true_defocus(self, dataset4d): + true_defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, true_defocus) + + losses = {} + for scale in (0.5, 0.8, 1.0, 1.2, 1.5): + montage.reconstruct( + override_aberration_coefs={"C10": scale * true_defocus}, + parallax_flip_phase=False, + verbose=False, + ) + losses[scale] = float(montage.variance_loss()) + + assert all(value > 0 for value in losses.values()) + assert min(losses, key=losses.get) == pytest.approx(1.0) + + def test_rms_gradient_drives_a_search_over_a_convolution_kernel(self, dataset4d): + """`variance_loss` is undefined for `ssb` here; the sharpness objective is not. + + That is the reason it exists on this class: without it there is no way to tune + aberrations for the real-space SSB/OBF/MF kernels at all. + """ + true_defocus = integer_shift_defocus(1) + _, montage = _build_pair(dataset4d, true_defocus) + + montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=6, verbose=False) + with pytest.raises(NotImplementedError): + montage.variance_loss() + + montage.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter( + low=0.4 * true_defocus, high=1.6 * true_defocus, n_points=5 + ) + }, + loss="rms_gradient", + deconvolution_kernel="ssb", + stencil_radius=6, + verbose=False, + ) + best = montage.hyperparameter_state.optimized_aberrations["C10"] + + step = (1.6 - 0.4) * true_defocus / 4 + assert abs(best - true_defocus) <= step + + +class TestSerialization: + """Both classes must survive a save/load round-trip and stay usable afterwards.""" + + @pytest.mark.parametrize("cls_name", ["fourier", "montage"]) + def test_round_trip(self, dataset4d, tmp_path, cls_name): + defocus = integer_shift_defocus(1) + fourier, montage = _build_pair(dataset4d, defocus) + recon = fourier if cls_name == "fourier" else montage + + recon.hyperparameter_state.optimized_aberrations = {"C10": 123.0} + recon.hyperparameter_state.optimized_rotation_angle = 7.5 + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + before = recon.obj.copy() + + path = str(tmp_path / f"{cls_name}.zip") + recon.save(path, mode="o") + restored = load(path) + + assert type(restored) is type(recon) + assert np.array_equal(restored.obj, before) + assert restored.hyperparameter_state.optimized_aberrations == {"C10": 123.0} + assert restored.hyperparameter_state.optimized_rotation_angle == 7.5 + assert restored.gpts == recon.gpts + assert float(restored.variance_loss()) == pytest.approx(float(recon.variance_loss())) + + # and it must still be able to reconstruct + restored.reconstruct(deconvolution_kernel="prlx", verbose=False) + assert np.allclose(restored.obj, before) + + def test_torch_size_attributes_round_trip_as_tuples(self, dataset4d): + """`torch.Size` is a tuple subclass; AutoSerialize used to choke on the subclass name.""" + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + assert type(montage.gpts) is tuple + assert type(montage.scan_gpts) is tuple + + +class TestVisualization: + """The object sampling, and hence `visualize`'s scalebar, must follow the upsampling.""" + + @pytest.mark.parametrize("cls_name", ["fourier", "montage"]) + @pytest.mark.parametrize("upsampling_factor", [1, 2, 3]) + def test_scalebar_follows_the_upsampling(self, dataset4d, cls_name, upsampling_factor): + fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1, upsampling_factor)) + recon = fourier if cls_name == "fourier" else montage + recon.reconstruct( + deconvolution_kernel="prlx", upsampling_factor=upsampling_factor, verbose=False + ) + + expected = SCAN_SAMPLING / upsampling_factor + assert recon._obj_sampling[0] == pytest.approx(expected) + # the reported sampling must span the same field of view as the image itself + assert recon.obj.shape[0] * recon._obj_sampling[0] == pytest.approx(N * SCAN_SAMPLING) + + def test_sampling_resets_between_reconstructions(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + montage.reconstruct(upsampling_factor=3, verbose=False) + assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / 3) + montage.reconstruct(upsampling_factor=1, verbose=False) + assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING) + + def test_sampling_is_defined_before_reconstructing(self, dataset4d): + fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + for recon in (fourier, montage): + assert recon._obj_sampling[0] == pytest.approx(SCAN_SAMPLING) + + def test_fov_matches_the_scan_extent(self, dataset4d): + fourier, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + + for recon in (fourier, montage): + assert recon.fov == pytest.approx((N * SCAN_SAMPLING, N * SCAN_SAMPLING)) + + @pytest.mark.parametrize("upsampling_factor", [1, 2]) + def test_padded_canvas_keeps_the_scan_sampling(self, dataset4d, upsampling_factor): + """A padded canvas spans more than the scan, so its sampling is not fov/shape. + + Deriving from the *scan* field of view would under-report the pixel size by the + padding fraction; `_obj_fov` reports the canvas extent instead. + """ + _, montage = _build_pair(dataset4d, integer_shift_defocus(1, upsampling_factor)) + montage.reconstruct(boundary="pad", upsampling_factor=upsampling_factor, verbose=False) + + assert montage.obj.shape[0] > N * upsampling_factor # canvas really did grow + assert montage._obj_sampling[0] == pytest.approx(SCAN_SAMPLING / upsampling_factor) + assert montage._obj_fov[0] > montage.fov[0] + # and the reported extent still matches the image it describes + assert montage.obj.shape[0] * montage._obj_sampling[0] == pytest.approx( + montage._obj_fov[0] + ) + + def test_wrapped_canvas_spans_exactly_the_scan(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + montage.reconstruct(boundary="wrap", upsampling_factor=2, verbose=False) + + assert montage._obj_fov == pytest.approx(montage.fov) + + def test_visualize_before_reconstruct_raises(self, dataset4d): + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + with pytest.raises(RuntimeError, match="Run reconstruct"): + montage.visualize() + + +class TestFourierProbe: + """An empirical ``psi(k)`` in place of an aperture plus aberrations.""" + + DEFOCUS = staticmethod(lambda: integer_shift_defocus(1)) + + def _pair(self, dataset4d, cls, **probe_kwargs): + """The same reconstruction twice: analytic, and with its own probe fed back in.""" + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) + if cls is DirectPtychographyMontage: + kwargs["boundary"] = "wrap" + + analytic = cls.from_dataset4d(dataset4d, **kwargs) + psi = analytic_probe_array(analytic, {"C10": defocus}) + probe = FourierProbe.from_array( + psi, analytic.reciprocal_sampling, analytic.wavelength, **probe_kwargs + ) + empirical = cls.from_dataset4d(dataset4d, fourier_probe=probe, **kwargs) + return analytic, empirical + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + @pytest.mark.parametrize("kernel", ["ssb", "obf", "mf"]) + def test_empirical_probe_reproduces_the_analytic_one(self, dataset4d, cls, kernel): + """The headline: the same probe, described two ways, must reconstruct the same. + + This is what makes the empirical path trustworthy -- there is no ground truth for a + measured probe, so the only check available is that a probe we *can* write down + analytically goes through the array path unchanged. + """ + analytic, empirical = self._pair(dataset4d, cls, normalize=False) + extra = {} if cls is DirectPtychography else {"stencil_radius": 5} + + analytic.reconstruct(deconvolution_kernel=kernel, verbose=False, **extra) + expected = analytic.obj.copy() + empirical.reconstruct(deconvolution_kernel=kernel, verbose=False, **extra) + + assert _relative_error(empirical.obj, expected) < 1e-5 + + def test_is_analytic_reports_which_path_is_live(self, dataset4d): + analytic, empirical = self._pair(dataset4d, DirectPtychographyMontage) + + assert analytic.fourier_probe is None + assert empirical.fourier_probe is not None + assert empirical.fourier_probe.is_analytic is False + + def test_zero_outside_the_detector(self, dataset4d): + """Beyond the detector's Nyquist nothing was measured, so `psi` is zero there. + + Wrapping instead would fold the opposite edge of the aperture back in, which moved + the reconstruction by 13% on this fixture before it was fixed -- the bright-field + mask is cropped tight to the disk, so `k -/+ q` leaves the grid constantly. + """ + _, empirical = self._pair(dataset4d, DirectPtychographyMontage) + probe = empirical.fourier_probe + dq_row, dq_col = probe.reciprocal_sampling + n_rows = probe.array.shape[0] + + just_outside = torch.tensor([(n_rows // 2 + 1) * dq_row]) + zero = torch.zeros(1) + + assert probe.at(just_outside, zero).abs().item() == 0.0 + assert probe.at(-just_outside, zero).abs().item() == 0.0 + # and the centre is emphatically not zero + assert probe.at(zero, zero).abs().item() > 0 + + def test_resampling_is_exact_on_the_shared_lattice(self): + """Zero-padding a compact real-space probe refines `psi` without changing it.""" + torch.manual_seed(0) + real = torch.zeros(32, 32, dtype=torch.complex64) + real[:8, :8] = torch.randn(8, 8, dtype=torch.complex64) # confined, so band limited + psi = torch.fft.fft2(real) + probe = FourierProbe.from_array(psi, (0.25, 0.25), 0.02, normalize=False) + + refined = probe.resampled_to((0.0625, 0.0625)) + + assert refined.array.shape == (128, 128) + # every fourth sample of the fine grid is the original, exactly + assert _relative_error(to_numpy(refined.array[::4, ::4]), to_numpy(psi)) < 1e-5 + + def test_resampling_handles_anisotropy_and_odd_sizes(self): + torch.manual_seed(0) + real = torch.zeros(31, 16, dtype=torch.complex64) + real[:6, :4] = torch.randn(6, 4, dtype=torch.complex64) + psi = torch.fft.fft2(real) + probe = FourierProbe.from_array(psi, (0.25, 0.5), 0.02, normalize=False) + + refined = probe.resampled_to((0.125, 0.125)) + + assert refined.array.shape == (62, 64) + assert _relative_error(to_numpy(refined.array[::2, ::4]), to_numpy(psi)) < 1e-5 + + def test_resampling_rejects_a_non_integer_ratio(self): + probe = FourierProbe.from_array( + np.ones((8, 8), dtype=np.complex64), (0.25, 0.25), 0.02, normalize=False + ) + with pytest.raises(ValueError, match="whole number"): + probe.resampled_to((0.1, 0.1)) + + def test_a_canvas_matching_the_probe_needs_no_resampling(self, dataset4d): + """Canvas field of view == probe field of view is the exact, assumption-free case.""" + _, empirical = self._pair(dataset4d, DirectPtychographyMontage, normalize=False) + empirical.reconstruct(deconvolution_kernel="ssb", convolution_mode="fft", verbose=False) + + _, _, refined = empirical._resampled_probe + assert refined.array.shape == empirical.fourier_probe.array.shape + + def test_refined_sampling_converges_on_the_exact_shift(self): + """Bilinear on a refined grid must approach the exact band-limited evaluation. + + Isolated from any reconstruction, and on a probe whose real-space form really is + confined, so that zero-padding is exact and only the interpolation is under test. + Error falls as the square of the refinement. + """ + torch.manual_seed(0) + n = 64 + real = torch.zeros(n, n, dtype=torch.complex64) + real[:10, :10] = torch.randn(10, 10, dtype=torch.complex64) + psi = torch.fft.fft2(real) + step = 0.25 + + # ground truth at a generic sub-pixel offset: a phase ramp in real space + drow, dcol = 0.31, 0.17 + ramp_axis = torch.fft.fftfreq(n) + ramp = torch.exp(-2j * np.pi * (drow * ramp_axis[:, None] + dcol * ramp_axis[None, :])) + truth = torch.fft.fft2(real * ramp) + + index = torch.arange(n) + centered = torch.where(index < n // 2, index, index - n).to(torch.float32) + kx = centered[:, None].expand(n, n) * step + ky = centered[None, :].expand(n, n) * step + + errors = [] + for oversample in (1, 4, 16): + probe = FourierProbe.from_array( + psi, (step, step), 0.02, normalize=False, interpolation="bilinear" + ) + if oversample > 1: + probe = probe.resampled_to((step / oversample, step / oversample)) + sampled = probe.at(kx + drow * step, ky + dcol * step) + errors.append(float((sampled - truth).abs().max() / truth.abs().max())) + + assert errors[0] > errors[1] > errors[2] + assert errors[2] < 0.01 + + @pytest.mark.parametrize("rotation_angle", [15.0, -30.0]) + def test_rotation_off_the_lattice_is_supported(self, dataset4d, rotation_angle): + """A rotation carries `q` off the probe's lattice, which is interpolated, not refused. + + It is a real experimental parameter. What it costs is a sampling error that the + refinement controls -- and, on this fixture, a floor of about 6% that no refinement + removes: the aperture is cropped to an 11x11 grid, so its real-space probe is not + confined and the band-limited interpolant is not the true probe. That floor is a + property of the data, not of the method. + """ + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + kwargs["rotation_angle"] = rotation_angle + + analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + psi = analytic_probe_array(analytic, {"C10": defocus}) + analytic.reconstruct(deconvolution_kernel="ssb", convolution_mode="fft", verbose=False) + + empirical = DirectPtychographyMontage.from_dataset4d( + dataset4d, + fourier_probe=FourierProbe.from_array( + psi, analytic.reciprocal_sampling, analytic.wavelength, normalize=False + ), + **kwargs, + ) + empirical.reconstruct( + deconvolution_kernel="ssb", convolution_mode="fft", probe_oversample=8, verbose=False + ) + + assert np.isfinite(empirical.obj).all() + assert correlation(empirical.obj, analytic.obj) > 0.99 + + def test_low_oversampling_warns(self, dataset4d): + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + kwargs["rotation_angle"] = 15.0 + analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + empirical = DirectPtychographyMontage.from_dataset4d( + dataset4d, + fourier_probe=FourierProbe.from_array( + analytic_probe_array(analytic, {"C10": defocus}), + analytic.reciprocal_sampling, + analytic.wavelength, + ), + **kwargs, + ) + + with pytest.warns(UserWarning, match="off the probe's reciprocal lattice"): + empirical.reconstruct( + deconvolution_kernel="ssb", + convolution_mode="fft", + probe_oversample=2, + verbose=False, + ) + + def test_off_grid_sampling_raises(self, dataset4d): + """A canvas incommensurate with the probe would need interpolation; say so.""" + _, empirical = self._pair(dataset4d, DirectPtychographyMontage) + probe = empirical.fourier_probe + half_pixel = torch.tensor([0.5 * probe.reciprocal_sampling[0]]) + + with pytest.raises(ValueError, match="own reciprocal grid"): + probe.at(half_pixel, torch.zeros(1)) + + def test_bilinear_accepts_off_grid_sampling(self, dataset4d): + _, empirical = self._pair(dataset4d, DirectPtychographyMontage, interpolation="bilinear") + probe = empirical.fourier_probe + dq_row = probe.reciprocal_sampling[0] + + midpoint = probe.at(torch.tensor([0.5 * dq_row]), torch.zeros(1)) + ends = probe.at(torch.tensor([0.0, dq_row]), torch.zeros(2)) + + assert midpoint.item() == pytest.approx(complex(ends.mean()), rel=1e-5) + + @pytest.mark.parametrize( + "kwargs, match", + [ + ({"deconvolution_kernel": "prlx"}, "parallax kernel"), + ( + { + "deconvolution_kernel": "ssb", + "stencil_radius": 5, + "defocus_gradient": (1.0, 0.0), + }, + "defocus gradient", + ), + ], + ) + def test_aberration_only_features_raise(self, dataset4d, kwargs, match): + """Everything that reads chi(k) has no meaning without aberration coefficients.""" + _, empirical = self._pair(dataset4d, DirectPtychographyMontage) + + with pytest.raises(NotImplementedError, match=match): + empirical.reconstruct(verbose=False, **kwargs) + + def test_semiangle_cutoff_becomes_optional(self, dataset4d): + """The empirical probe carries its own aperture, whatever shape it is.""" + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + psi = analytic_probe_array(analytic, {"C10": defocus}) + + kwargs["semiangle_cutoff"] = None + empirical = DirectPtychographyMontage.from_dataset4d( + dataset4d, + fourier_probe=FourierProbe.from_array( + psi, analytic.reciprocal_sampling, analytic.wavelength + ), + **kwargs, + ) + + assert empirical.semiangle_cutoff is None + empirical.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) + assert np.isfinite(empirical.obj).all() + + def test_normalize_gives_unit_intensity(self, dataset4d): + _, empirical = self._pair(dataset4d, DirectPtychographyMontage, normalize=True) + + assert float(empirical.fourier_probe.array.abs().square().sum()) == pytest.approx(1.0) + + @pytest.mark.parametrize( + "array, match", + [ + (np.ones((8, 8)), "must be a complex probe"), + (np.ones((8, 8, 2), dtype=complex), "must be 2D"), + ], + ) + def test_array_validation(self, array, match): + with pytest.raises(ValueError, match=match): + FourierProbe.from_array(array, (0.1, 0.1), 0.02) + + def test_shape_must_match_the_detector(self, dataset4d): + analytic, _ = self._pair(dataset4d, DirectPtychographyMontage) + wrong = FourierProbe.from_array( + np.ones((4, 4), dtype=np.complex64), analytic.reciprocal_sampling, analytic.wavelength + ) + + with pytest.raises(ValueError, match="detector grid"): + analytic.fourier_probe = wrong + + def test_survives_a_round_trip(self, dataset4d, tmp_path): + _, empirical = self._pair(dataset4d, DirectPtychographyMontage) + empirical.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) + before = empirical.obj.copy() + + path = str(tmp_path / "empirical.zip") + empirical.save(path, mode="o") + restored = load(path) + + assert restored.fourier_probe is not None + restored.reconstruct(deconvolution_kernel="ssb", stencil_radius=5, verbose=False) + assert np.allclose(restored.obj, before) + + +class TestWavelength: + """`wavelength` given directly, for anything that is not an electron.""" + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + def test_energy_and_wavelength_agree(self, dataset4d, cls): + """The two routes must land on the same geometry, and hence the same image.""" + common = dict( + semiangle_cutoff=SEMIANGLE_CUTOFF, + aberration_coefs={"C10": integer_shift_defocus(1)}, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + edge_blend_pixels=0, + verbose=False, + ) + from_energy = cls.from_dataset4d(dataset4d, energy=PROBE_ENERGY, **common) + from_wavelength = cls.from_dataset4d( + dataset4d, wavelength=from_energy.wavelength, **common + ) + + assert from_wavelength.wavelength == pytest.approx(from_energy.wavelength) + assert from_wavelength.angular_sampling == pytest.approx(from_energy.angular_sampling) + + for recon in (from_energy, from_wavelength): + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + assert np.array_equal(from_wavelength.obj, from_energy.obj) + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + @pytest.mark.parametrize( + "kwargs", + [ + {}, # neither + {"energy": PROBE_ENERGY, "wavelength": 0.02}, # both + ], + ) + def test_exactly_one_of_energy_or_wavelength(self, dataset4d, cls, kwargs): + with pytest.raises(ValueError, match="exactly one of `energy`"): + cls.from_dataset4d( + dataset4d, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + verbose=False, + **kwargs, + ) + + def test_photon_wavelength_is_not_the_electron_one(self): + """The reason this exists: the de Broglie formula is wrong for photons. + + `electron_wavelength_angstrom` is relativistic and correct for electrons -- 0.019687 + Angstrom at 300 kV -- but it is the wrong physics for a 7.9 keV photon, where the + answer is `hc/E` = 1.5694 rather than 0.1375. + """ + from quantem.core.utils.utils import electron_wavelength_angstrom + + assert electron_wavelength_angstrom(300e3) == pytest.approx(0.019687, abs=1e-6) + + photon_energy_ev = 7900.08 + assert electron_wavelength_angstrom(photon_energy_ev) == pytest.approx(0.1375, abs=1e-3) + assert 12398.42 / photon_energy_ev == pytest.approx(1.5694, abs=1e-3) + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + def test_wavelength_survives_a_round_trip(self, dataset4d, cls, tmp_path): + recon = cls.from_dataset4d( + dataset4d, + wavelength=1.5694, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + edge_blend_pixels=0, + verbose=False, + ) + recon.reconstruct(deconvolution_kernel="prlx", verbose=False) + + path = str(tmp_path / f"{cls.__name__}.zip") + recon.save(path, mode="o") + restored = load(path) + + assert restored.wavelength == pytest.approx(1.5694) + + @pytest.mark.parametrize("bad", [0.0, -1.0]) + def test_non_positive_wavelength_raises(self, dataset4d, bad): + with pytest.raises(ValueError): + DirectPtychographyMontage.from_dataset4d( + dataset4d, + wavelength=bad, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + verbose=False, + ) + + +class TestSemiangleCutoff: + """`semiangle_cutoff` sets the probe aperture and is never optional.""" + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + def test_from_virtual_bfs_requires_it(self, cls, dataset4d): + """Enforced at runtime rather than by the signature. + + It used to be a required positional, but `wavelength` has to sit alongside `energy` + with a default, and Python will not take a defaulted parameter before an + undefaulted one. The guarantee is the error, so test the error. + """ + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + vbf_dataset = Dataset3d.from_array( + to_numpy(montage.vbf_stack).reshape(montage.num_bf, N, N), + name="vBF stack", + units=("index", "A", "A"), + sampling=(1, SCAN_SAMPLING, SCAN_SAMPLING), + ) + bf_mask_dataset = Dataset2d.from_array( + to_numpy(montage.bf_mask), + name="BF mask", + units=("A^-1", "A^-1"), + sampling=tuple(montage.reciprocal_sampling), + ) + + with pytest.raises(ValueError, match="`semiangle_cutoff` is required"): + cls.from_virtual_bfs( + vbf_dataset, + bf_mask_dataset, + energy=PROBE_ENERGY, + rotation_angle=0.0, + crop_bf_mask=False, + verbose=False, + ) + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + def test_none_raises_a_clear_error(self, dataset4d, cls): + with pytest.raises(ValueError, match="`semiangle_cutoff` is required"): + cls.from_dataset4d( + dataset4d, + energy=PROBE_ENERGY, + semiangle_cutoff=None, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + verbose=False, + ) + + @pytest.mark.parametrize("cls", [DirectPtychography, DirectPtychographyMontage]) + def test_non_positive_raises(self, dataset4d, cls): + with pytest.raises(ValueError): + cls.from_dataset4d( + dataset4d, + energy=PROBE_ENERGY, + semiangle_cutoff=-1.0, + rotation_angle=0.0, + force_fitted_origin=ORIGIN, + verbose=False, + ) + + +class TestDefocusGradient: + """Position-dependent defocus, for a tilted sample. + + The montage shifts each scan position by its own local defocus. A Fourier multiplier is + global over the scan by construction, so `DirectPtychography` has no counterpart to + compare against; these check the model relation directly instead. + """ + + def test_none_and_zero_are_the_same_reconstruction(self, dataset4d): + """The gradient must be a no-op when absent -- guards the fast path.""" + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + + without = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + with_zero = DirectPtychographyMontage.from_dataset4d( + dataset4d, defocus_gradient=(0.0, 0.0), **kwargs + ) + + assert np.array_equal( + without.reconstruct(verbose=False).obj, + with_zero.reconstruct(verbose=False).obj, + ) + + def test_defocus_rate_is_the_analytic_lambda_k(self, dataset4d): + """`d shift / d C10 = wavelength * k`, independent of the other aberrations.""" + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) + + kxa, kya = spatial_frequencies(montage.gpts, montage.sampling, device=montage.device) + scan_sampling = torch.as_tensor( + tuple(montage.scan_sampling), dtype=torch.float64, device=montage.device + ) + expected = ( + torch.stack((kxa[montage.bf_mask], kya[montage.bf_mask]), -1).to(torch.float64) + * montage.wavelength + / scan_sampling + ) + + assert torch.allclose(rate, expected, atol=1e-9) + + def test_defocus_rate_ignores_other_aberrations(self, dataset4d): + """chi is linear in every magnitude, so the rate cannot depend on the rest.""" + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) + + base = {"C10": 500.0, "C12": 40.0, "phi12": 0.7, "C30": 1.2e5} + shifted = montage._return_shifts_px(0.0, {**base, "C10": 501.0}, montage.bf_mask, 1)[0] + unshifted = montage._return_shifts_px(0.0, base, montage.bf_mask, 1)[0] + + assert torch.allclose(shifted - unshifted, rate, atol=1e-6) + + def test_delta_defocus_is_mean_zero(self, dataset4d): + """Measuring from the centroid keeps the gradient orthogonal to the global C10.""" + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + delta = montage._return_delta_c10((7.0, -3.0)) + + assert delta is not None + assert float(delta.mean().abs()) < 1e-9 + assert float(delta.abs().max()) > 0 + + def test_zero_gradient_short_circuits(self, dataset4d): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + assert montage._return_delta_c10(None) is None + assert montage._return_delta_c10((0.0, 0.0)) is None + + def test_padded_canvas_covers_the_gradient(self, dataset4d): + """A gradient widens the range of shifts, so `"pad"` must grow to match.""" + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) + + flat = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + tilted = DirectPtychographyMontage.from_dataset4d( + dataset4d, defocus_gradient=(30.0, -10.0), **kwargs + ) + + flat_shape = flat.reconstruct(boundary="pad", verbose=False).obj.shape + tilted_shape = tilted.reconstruct(boundary="pad", verbose=False).obj.shape + + assert tilted_shape[0] > flat_shape[0] + assert tilted_shape[1] > flat_shape[1] + + def test_shift_extrema_match_a_brute_force_scan(self, dataset4d): + """The closed form must bound every (BF pixel, position) pair, with no slack.""" + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + gradient = (30.0, -10.0) + shifts = montage._return_shifts_px(0.0, montage.aberration_coefs, montage.bf_mask, 1)[0] + rate = montage._return_defocus_rate_px(0.0, montage.bf_mask, 1) + delta = montage._return_delta_c10(gradient) + + lo, hi = DirectPtychographyMontage._return_shift_extrema(shifts, rate, delta) + brute = shifts[:, None, :] + rate[:, None, :] * delta[None, :, None] + + assert torch.allclose(lo, brute.amin((0, 1))) + assert torch.allclose(hi, brute.amax((0, 1))) + + def test_sign_convention_on_simulated_tilted_data(self): + """On real 4D data a negated gradient must be worse than the true one. + + Only the *ordering* is asserted. At this fixture's 32 Angstrom field of view a + visible gradient needs a defocus swing so large that the probe size varies threefold + across the scan, and the parallax model itself starts to break down -- so + "better than no correction at all" is not true here, and is checked on the model + stack below instead. + """ + defocus = integer_shift_defocus(1) + gradient = (40.0, 0.0) + dataset = make_tilted_dataset4d(defocus, gradient) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) + + def corr(g): + montage = DirectPtychographyMontage.from_dataset4d( + dataset, defocus_gradient=g, **kwargs + ) + return correlation(montage.reconstruct(verbose=False).obj, band_limited_phase()) + + assert corr(gradient) > corr((-gradient[0], -gradient[1])) + + @pytest.mark.parametrize("gradient", [(20.0, 0.0), (20.0, -10.0), (-15.0, 25.0)]) + def test_correcting_the_gradient_sharpens_the_reconstruction(self, gradient): + montage, obj = _model_montage(gradient) + reconstruct = dict(parallax_flip_phase=False, interpolation="bilinear", verbose=False) + + uncorrected = correlation(montage.reconstruct(**reconstruct).obj, obj) + corrected = correlation( + montage.reconstruct(defocus_gradient=gradient, **reconstruct).obj, obj + ) + + assert corrected > uncorrected + assert corrected > 0.98 + + def test_defocus_map_tracks_a_seeded_plane(self): + gradient = (20.0, -10.0) + montage, _ = _model_montage(gradient) + + results = montage.defocus_map( + MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False + ) + expected = MODEL_DEFOCUS + results["centers_A"] @ np.asarray(gradient) + + assert results["valid"].all() + # the estimator carries a small uniform offset (~170 A, see the flat control below), + # so compare the spatial variation rather than the absolute value. + # 0.97 rather than 0.99: scoring only pixels near the peak accumulated weight used + # to reach 0.995 on this gridded fixture, but on a real ungridded scan the weight is + # uneven everywhere and that cut selects density hot-spots instead of the interior -- + # see `_patch_variance_loss`. The weight-averaged mean is unbiased on both. + recovered = results["c10_best"] - results["c10_best"].mean() + assert np.corrcoef(recovered, expected - expected.mean())[0, 1] > 0.97 + + def test_defocus_map_is_flat_without_a_gradient(self): + """The control that makes the test above meaningful.""" + montage, _ = _model_montage((0.0, 0.0)) + + results = montage.defocus_map( + MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False + ) + + assert np.ptp(results["c10_best"]) < 0.05 * MODEL_DEFOCUS + + @pytest.mark.parametrize("gradient", [(0.0, 0.0), (20.0, -10.0), (-15.0, 25.0)]) + def test_fit_defocus_gradient_recovers_the_seed(self, gradient): + montage, _ = _model_montage(gradient) + + montage.fit_defocus_gradient( + MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False + ) + + assert montage.defocus_gradient is not None + scale = max(np.hypot(*gradient), 1.0) + assert np.hypot(*np.subtract(montage.defocus_gradient, gradient)) < 0.15 * scale + + def test_fit_defocus_gradient_updates_the_global_defocus(self): + montage, _ = _model_montage((20.0, 0.0)) + montage.hyperparameter_state.optimized_aberrations = {} + + montage.fit_defocus_gradient( + MODEL_C10_GRID, patch_grid=(3, 3), interpolation="bilinear", verbose=False + ) + + assert "C10" in montage.hyperparameter_state.optimized_aberrations + assert "C10" in montage.hyperparameter_state.optimized_keys + + def test_fit_defocus_gradient_can_leave_the_defocus_alone(self): + montage, _ = _model_montage((20.0, 0.0)) + + montage.fit_defocus_gradient( + MODEL_C10_GRID, + patch_grid=(3, 3), + interpolation="bilinear", + update_defocus=False, + verbose=False, + ) + + assert montage.hyperparameter_state.optimized_aberrations == {} + + def test_endpoint_pinned_patches_are_invalid(self): + """A grid that does not bracket the local defocus must be reported, not fitted.""" + montage, _ = _model_montage((20.0, 0.0)) + + results = montage.defocus_map( + np.linspace(3400.0, 4500.0, 6), + patch_grid=(3, 3), + interpolation="bilinear", + verbose=False, + ) + + assert not results["valid"].all() + assert np.isnan(results["c10_best"][~results["valid"]]).all() + + def test_fit_raises_when_too_few_patches_bracket(self): + montage, _ = _model_montage((20.0, 0.0)) + + with pytest.raises(RuntimeError, match="bracketed minimum"): + montage.fit_defocus_gradient( + np.linspace(4200.0, 4500.0, 4), + patch_grid=(2, 2), + interpolation="bilinear", + verbose=False, + ) + + def test_defocus_map_allows_a_one_dimensional_grid(self): + """A (P, 1) grid is a profile along one axis -- only the plane fit needs three.""" + montage, _ = _model_montage((20.0, 0.0)) + + results = montage.defocus_map( + MODEL_C10_GRID, patch_grid=(3, 1), interpolation="bilinear", verbose=False + ) + + assert results["c10_best"].shape == (3,) + + def test_defocus_map_rejects_a_degenerate_grid(self): + montage, _ = _model_montage((0.0, 0.0)) + + with pytest.raises(ValueError, match="must be positive"): + montage.defocus_map(MODEL_C10_GRID, patch_grid=(0, 3), verbose=False) + + def test_defocus_map_needs_enough_trial_values(self): + montage, _ = _model_montage((0.0, 0.0)) + + with pytest.raises(ValueError, match="at least 3 points"): + montage.defocus_map([2000.0, 3000.0], patch_grid=(2, 2), verbose=False) + + def test_gradient_is_orthogonal_to_a_global_defocus_search(self): + """The API worry: a grid search over C10 must stay well posed with a gradient set.""" + gradient = (20.0, -10.0) + montage, _ = _model_montage(gradient) + + montage.grid_search_hyperparameters( + aberration_coefs={ + "C10": OptimizationParameter( + low=MODEL_DEFOCUS - 900, high=MODEL_DEFOCUS + 900, n_points=7 + ) + }, + defocus_gradient=gradient, + interpolation="bilinear", + parallax_flip_phase=False, + verbose=False, + ) + + fitted = montage.hyperparameter_state.current_aberrations()["C10"] + assert abs(fitted - MODEL_DEFOCUS) < 0.2 * MODEL_DEFOCUS + + def test_rejects_a_malformed_gradient(self, dataset4d): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, **_common_kwargs(integer_shift_defocus(1)) + ) + with pytest.raises(ValueError, match="must be a \\(row, col\\) pair"): + montage.defocus_gradient = (1.0, 2.0, 3.0) + + def test_survives_a_serialization_round_trip(self, dataset4d, tmp_path): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, defocus_gradient=(7.0, -3.0), **_common_kwargs(integer_shift_defocus(1)) + ) + path = tmp_path / "montage.zip" + montage.save(path, mode="o") + + assert load(path).defocus_gradient == (7.0, -3.0) + + +class TestRealSpaceKernels: + """SSB / OBF / MF as truncated real-space convolutions. + + `obj = sum_m ifft2(G_m * K_m) = sum_m (v_m conv kappa_m)` is an identity, so truncating + `kappa_m` to a box stencil is the *only* approximation: with a large enough stencil these + must reproduce `DirectPtychography` exactly. + + They converge slowly. Dividing by `|gamma|` leaves a unit-magnitude phase on a hard-edged + support, whose transform has `r**-1.5` tails, so the error falls like 1/radius -- and + being in focus does not help, since there is no chirp to concentrate the kernel. + """ + + KERNELS = ("ssb", "obf", "mf") + MAX_RADIUS = 15 # canvas half-width for the 32x32 fixture + + @staticmethod + def _pair(dataset4d): + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) + return ( + DirectPtychography.from_dataset4d(dataset4d, **kwargs), + DirectPtychographyMontage.from_dataset4d(dataset4d, boundary="wrap", **kwargs), + ) + + @pytest.mark.parametrize("kernel", KERNELS) + def test_converges_to_the_fourier_kernel(self, dataset4d, kernel): + fourier, montage = self._pair(dataset4d) + reference = fourier.reconstruct(deconvolution_kernel=kernel, verbose=False).obj + + errors = [] + for radius in (4, 8, self.MAX_RADIUS): + obj = montage.reconstruct( + deconvolution_kernel=kernel, + stencil_radius=radius, + interpolation="nearest", + verbose=False, + ).obj + errors.append(_relative_error(obj, reference)) + + assert errors[0] > errors[1] > errors[2] + assert errors[-1] < 2e-2 + + @pytest.mark.parametrize("kernel", KERNELS) + def test_reported_error_bounds_the_measured_one(self, dataset4d, kernel): + """The estimate assumes a white object spectrum, so it must be conservative.""" + fourier, montage = self._pair(dataset4d) + reference = fourier.reconstruct(deconvolution_kernel=kernel, verbose=False).obj + + obj = montage.reconstruct( + deconvolution_kernel=kernel, + stencil_radius=8, + interpolation="nearest", + verbose=False, + ).obj + + assert _relative_error(obj, reference) < montage._stencil_info["mean_error"] + + def test_a_box_stencil_beats_a_tapered_one(self, dataset4d): + """Pins a counterintuitive choice: tapering the stencil measures consistently worse. + + A Hann taper discards mid-radius content that matters more than the ringing it + suppresses, so do not "improve" the box away. + """ + fourier, montage = self._pair(dataset4d) + reference = fourier.reconstruct(deconvolution_kernel="ssb", verbose=False).obj + radius = 8 + + boxed = montage.reconstruct( + deconvolution_kernel="ssb", + stencil_radius=radius, + interpolation="nearest", + verbose=False, + ).obj + + # rebuild the same stencil, taper it, and run the same accumulation by hand + bf = montage._return_bf_context(montage.bf_mask) + shifts, _ = montage._return_shifts_px(0.0, montage.aberration_coefs, bf.bf_mask, 1) + offsets, weights, bf_weights, _ = montage._return_kernel_stencil( + bf, + kernel="ssb", + rotation_angle=0.0, + aberration_coefs=montage.aberration_coefs, + canvas_shape=montage._canvas_shape, + upsampling_factor=1, + shift_centers=shifts.round(), + stencil_radius=radius, + truncation_tolerance=1.0, + max_stencil_radius=radius, + matched_filter_norm_epsilon=1e-1, + kernel_batch_size=16, + verbose=False, + ) + distance = offsets.to(torch.float64).abs().amax(-1) + taper = 0.5 * (1 + torch.cos(np.pi * distance / (radius + 1))) + canvas_shape = montage._canvas_shape + + accumulator = torch.zeros( + canvas_shape[0] * canvas_shape[1], device=montage.device, dtype=torch.complex64 + ) + coords = montage.positions_px[None] + shifts.round()[:, None] + scatter_add_convolve( + montage.vbf_stack, + coords, + canvas_shape, + offsets, + weights * taper[None, :], + boundary="wrap", + interpolation="nearest", + out=accumulator, + ) + tapered = to_numpy(accumulator.real.reshape(canvas_shape) / bf_weights) + + assert _relative_error(boxed, reference) < _relative_error(tapered, reference) + + def test_warns_when_the_tolerance_cannot_be_met(self, dataset4d): + _, montage = self._pair(dataset4d) + + with pytest.warns(UserWarning, match="truncation error"): + montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=2, verbose=False) + + def test_auto_respects_the_radius_cap(self, dataset4d): + _, montage = self._pair(dataset4d) + + with pytest.warns(UserWarning, match="truncation error"): + montage.reconstruct( + deconvolution_kernel="ssb", + convolution_mode="stencil", + stencil_radius="auto", + max_stencil_radius=3, + verbose=False, + ) + + assert montage._stencil_info["stencil_radius"] <= 3 + + def test_auto_reports_what_it_chose(self, dataset4d): + _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="stencil", verbose=False) + + info = montage._stencil_info + assert info["stencil_radius"] >= 1 + assert 0.0 <= info["mean_error"] <= info["max_error"] + + @pytest.mark.parametrize("kernel", KERNELS) + def test_variance_loss_is_undefined(self, dataset4d, kernel): + _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel=kernel, stencil_radius=4, verbose=False) + + with pytest.raises(NotImplementedError, match="only defined for the parallax"): + montage.variance_loss() + + def test_variance_loss_returns_after_a_parallax_reconstruction(self, dataset4d): + _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=4, verbose=False) + montage.reconstruct(verbose=False) + + assert float(montage.variance_loss()) > 0 + + def test_icom_runs_in_either_mode(self, dataset4d): + _, montage = self._pair(dataset4d) + for mode in ("fft", "stencil"): + montage.reconstruct( + deconvolution_kernel="icom", + convolution_mode=mode, + stencil_radius=8, + verbose=False, + ) + assert np.isfinite(montage.obj).all() + + def test_icom_ignores_an_empirical_probe(self, dataset4d): + """It never reads psi, so an empirical probe changes nothing -- not even a bit. + + `normalize=False` matters only because the probe's total intensity sets + `bf_weights`, which scales the finished object; the image itself is identical + either way. + """ + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + analytic = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + empirical = DirectPtychographyMontage.from_dataset4d( + dataset4d, + fourier_probe=FourierProbe.from_array( + analytic_probe_array(analytic, {"C10": defocus}), + analytic.reciprocal_sampling, + analytic.wavelength, + normalize=False, + ), + **kwargs, + ) + for recon in (analytic, empirical): + recon.reconstruct(deconvolution_kernel="icom", convolution_mode="fft", verbose=False) + + assert np.array_equal(empirical.obj, analytic.obj) + + def test_truncated_icom_is_ricom(self): + """A truncated iCoM stencil *is* riCOM (Yu et al., Microsc Microanal 28, 1526). + + riCOM cross-correlates the centre-of-mass shift map with a kernel + `(r_p - r_xy) / |r_p - r_xy|**2` truncated to an n x n box. That kernel is the + real-space form of the iCoM operator: for `G = ln|r| / 2pi`, `grad G = r / (2 pi + |r|**2)`, whose transform is `-i q / |q|**2`. Since each bright-field pixel's kernel + is linear in `k_m`, summing over the detector collapses the montage's per-pixel + convolutions into one convolution of the COM shift -- which is riCOM exactly. + + Two consequences, both checked below: a kernel spanning the canvas reproduces + untruncated iCoM, and shrinking it acts as a high pass. The latter is riCOM's whole + point -- it is what suppresses the long-range drift that blurs an iCoM image. + """ + vbf, bf_mask, _ = make_model_vbf_stack(MODEL_DEFOCUS, (0.0, 0.0), scan_gpts=(96, 96)) + montage = DirectPtychographyMontage.from_virtual_bfs( + vbf, bf_mask, **model_vbf_kwargs(MODEL_DEFOCUS) + ) + + size = montage.scan_gpts[0] + axis = np.minimum(np.arange(size), size - np.arange(size)) + radius = np.hypot(axis[:, None], axis[None, :]).astype(int) + + def bands(image): + spectrum = np.abs(np.fft.fft2(image - image.mean())) + profile = np.bincount(radius.ravel(), spectrum.ravel()) / np.maximum( + np.bincount(radius.ravel()), 1 + ) + return profile[1:4].mean(), profile[10:25].mean() + + common = dict(deconvolution_kernel="icom", boundary="wrap", verbose=False) + montage.reconstruct(convolution_mode="fft", **common) + untruncated = montage.obj.copy() + + ratios = [] + for stencil_radius in (5, 10, 20, 40): + montage.reconstruct( + convolution_mode="stencil", stencil_radius=stencil_radius, **common + ) + low, high = bands(montage.obj) + ratios.append(high / low) + if stencil_radius == 40: # spans the canvas, so nothing is truncated away + assert correlation(montage.obj, untruncated) > 0.99 + + # smaller kernel -> more weight at high frequency, monotonically + assert ratios == sorted(ratios, reverse=True) + assert ratios[0] > 1.2 * ratios[-1] + + def test_collapsing_the_detector_changes_nothing(self, dataset4d): + """riCOM sums over the detector first; that must be an optimization, not a change. + + The per-detector-pixel route still runs when a defocus gradient gives each position + its own shift, which the collapse cannot represent -- so a gradient small enough to + be physically negligible exercises the slow path and the two must agree. + """ + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + montage = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + + montage.reconstruct(deconvolution_kernel="icom", convolution_mode="fft", verbose=False) + collapsed = montage.obj.copy() + montage.reconstruct( + deconvolution_kernel="icom", + convolution_mode="fft", + defocus_gradient=(1e-9, 0.0), + verbose=False, + ) + + assert _relative_error(montage.obj, collapsed) < 1e-5 + + def test_com_shift_is_the_k_weighted_first_moment(self, dataset4d): + """The collapse itself, against an explicit sum over the detector.""" + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0, boundary="wrap") + montage = DirectPtychographyMontage.from_dataset4d(dataset4d, **kwargs) + bf = montage._return_bf_context(montage.bf_mask) + + collapsed = montage._return_com_shift(bf, 0.0, 1) + + kxa, kya, _, _ = montage._return_k_grid(0.0) + k_vectors = torch.stack((kxa[bf.bf_mask], kya[bf.bf_mask]), dim=-1) + values = montage._vbf_stack[bf.vbf_index_mapping] + expected = torch.stack([(values * k_vectors[:, i : i + 1]).sum(0) for i in (0, 1)]) + + assert _relative_error(to_numpy(collapsed), to_numpy(expected)) < 1e-5 + + def test_the_icom_kernel_is_the_ricom_kernel(self): + """`ifft2(-i q / |q|**2)` is `r / (2 pi |r|**2)`, the kernel riCOM writes down.""" + size = 128 + frequency = np.fft.fftfreq(size) + qx, qy = frequency[:, None], frequency[None, :] + q_square = qx**2 + qy**2 + q_square[0, 0] = 1.0 + operator = -1j * qx / q_square + operator[0, 0] = 0.0 + transformed = np.real(np.fft.ifft2(operator)) + + position = np.fft.fftfreq(size, 1 / size) + rx, ry = position[:, None], position[None, :] + r_square = rx**2 + ry**2 + r_square[0, 0] = 1.0 + analytic = rx / (2 * np.pi * r_square) + + # away from the singular origin and the periodic seam + distance = np.hypot(rx, ry) + inside = (distance > 3) & (distance < size / 4) + assert np.corrcoef(transformed[inside], analytic[inside])[0, 1] > 0.98 + + def test_phase_flip_is_not_applied_to_deconvolution_kernels(self, dataset4d): + """The kernels already invert the contrast transfer, as in `DirectPtychography`.""" + _, montage = self._pair(dataset4d) + common = dict(deconvolution_kernel="ssb", stencil_radius=4, verbose=False) + + flipped = montage.reconstruct(parallax_flip_phase=True, **common).obj + unflipped = montage.reconstruct(parallax_flip_phase=False, **common).obj + + assert np.array_equal(flipped, unflipped) + + def test_stencil_covers_the_parallax_shift_without_growing(self, dataset4d): + """The shift is divided out of the kernel, so the stencil sizes the residual only.""" + _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="stencil", verbose=False) + modest_defocus = montage._stencil_info["stencil_radius"] + + montage.reconstruct( + deconvolution_kernel="ssb", + convolution_mode="stencil", + override_aberration_coefs={"C10": integer_shift_defocus(3)}, + verbose=False, + ) + + assert montage._stencil_info["stencil_radius"] <= modest_defocus + 2 + + +class TestConvolutionModes: + """`convolution_mode` picks how the SSB / OBF / MF convolutions are evaluated. + + `"fft"` multiplies in `q`, which is exact; `"stencil"` truncates to a box, which is not. + Both are the same operator, so the only differences that may appear are the truncation + and the circular-versus-linear boundary. + """ + + KERNELS = ("ssb", "obf", "mf") + + @staticmethod + def _pair(dataset4d, boundary="wrap"): + defocus = integer_shift_defocus(1) + kwargs = dict(_common_kwargs(defocus), edge_blend_pixels=0) + return ( + DirectPtychography.from_dataset4d(dataset4d, **kwargs), + DirectPtychographyMontage.from_dataset4d(dataset4d, boundary=boundary, **kwargs), + ) + + @pytest.mark.parametrize("kernel", KERNELS) + def test_fft_mode_is_exact(self, dataset4d, kernel): + """The headline: no truncation at all, so it must equal the Fourier class. + + For contrast, the truncated stencil on this fixture is 20-34% off at radius 5 and + still 2-5% off at radius 12 -- these kernels are not compact and never converge fast. + """ + fourier, montage = self._pair(dataset4d) + fourier.reconstruct(deconvolution_kernel=kernel, verbose=False) + montage.reconstruct(deconvolution_kernel=kernel, convolution_mode="fft", verbose=False) + + assert _relative_error(montage.obj, fourier.obj) < 1e-5 + + @pytest.mark.parametrize("kernel", KERNELS) + def test_fft_mode_beats_a_truncated_stencil(self, dataset4d, kernel): + fourier, montage = self._pair(dataset4d) + fourier.reconstruct(deconvolution_kernel=kernel, verbose=False) + reference = fourier.obj.copy() + + montage.reconstruct(deconvolution_kernel=kernel, convolution_mode="fft", verbose=False) + exact = _relative_error(montage.obj, reference) + montage.reconstruct( + deconvolution_kernel=kernel, + convolution_mode="stencil", + stencil_radius=5, + verbose=False, + ) + truncated = _relative_error(montage.obj, reference) + + assert exact < 0.01 * truncated + + def test_pad_boundary_stays_linear(self, dataset4d): + """A Fourier convolution wraps; `"pad"` doubles the canvas so that it does not. + + Compared against a stencil wide enough that its own truncation is the larger error. + """ + _, montage = self._pair(dataset4d, boundary="pad") + montage.reconstruct( + deconvolution_kernel="ssb", + convolution_mode="stencil", + stencil_radius=14, + pad_px=4, + verbose=False, + ) + stencil = montage.obj.copy() + montage.reconstruct( + deconvolution_kernel="ssb", convolution_mode="fft", pad_px=4, verbose=False + ) + + assert montage.obj.shape == stencil.shape + assert _relative_error(montage.obj, stencil) < 0.05 + + @pytest.mark.parametrize("interpolation", ["nearest", "bilinear"]) + @pytest.mark.parametrize("boundary", ["wrap", "pad"]) + def test_conv2d_stencil_matches_the_scatter(self, dataset4d, boundary, interpolation): + """`splat_and_convolve` is a reorganization of `scatter_add_convolve`, not a change. + + The subtlety it has to reproduce is that the scatter tests the boundary at the + *deposit* position, so a point outside the canvas still contributes inward through + the kernel -- which is why the splat happens on a canvas grown by the radius. + """ + torch.manual_seed(0) + radius = 3 + span = torch.arange(-radius, radius + 1) + offsets = torch.stack(torch.meshgrid(span, span, indexing="ij"), dim=-1).reshape(-1, 2) + weights = torch.randn(4, offsets.shape[0], dtype=torch.complex64) + values = torch.randn(4, 30) + # deliberately spill outside the canvas on both sides + coords = torch.rand(4, 30, 2) * torch.tensor([28.0, 24.0]) - torch.tensor([4.0, 4.0]) + shape = (20, 16) + + scattered = torch.zeros(shape[0] * shape[1], dtype=torch.complex64) + scatter_add_convolve( + values, + coords, + shape, + offsets, + weights, + boundary=boundary, + interpolation=interpolation, + out=scattered, + ) + convolved = splat_and_convolve( + values, + coords, + shape, + weights, + radius, + boundary=boundary, + interpolation=interpolation, + ).sum(0) + + assert _relative_error(to_numpy(convolved), to_numpy(scattered.view(shape))) < 1e-5 + + def test_splat_stack_matches_a_shared_canvas(self): + """Summing the per-image canvases must give what the shared-canvas splat gives.""" + torch.manual_seed(1) + values = torch.randn(6, 25) + coords = torch.rand(6, 25, 2) * 12.0 + shape = (12, 12) + + stack = splat_stack(values, coords, shape, boundary="wrap", interpolation="bilinear") + _, shared, _ = scatter_add_splat( + values, coords, shape, boundary="wrap", interpolation="bilinear" + ) + + assert _relative_error(to_numpy(stack.sum(0)), to_numpy(shared.view(shape))) < 1e-5 + + @pytest.mark.parametrize("stencil_radius, expected", [("auto", "fft"), (6, "stencil")]) + def test_auto_reads_the_stencil_radius(self, stencil_radius, expected): + """Naming a radius is a request to truncate; leaving it 'auto' takes the exact route.""" + resolve = DirectPtychographyMontage._resolve_convolution_mode + assert resolve("auto", "ssb", stencil_radius) == expected + + def test_parallax_ignores_the_mode(self): + resolve = DirectPtychographyMontage._resolve_convolution_mode + assert resolve("fft", "prlx", "auto") == "splat" + + def test_unknown_mode_raises(self, dataset4d): + _, montage = self._pair(dataset4d) + with pytest.raises(ValueError, match="must be 'auto', 'fft' or 'stencil'"): + montage.reconstruct( + deconvolution_kernel="ssb", convolution_mode="direct", verbose=False + ) + + def test_fft_mode_leaves_no_stencil_info(self, dataset4d): + _, montage = self._pair(dataset4d) + montage.reconstruct(deconvolution_kernel="ssb", convolution_mode="fft", verbose=False) + + assert montage._stencil_info is None + + +class TestSharedCanvas: + """`obj_origin` / `obj_fov` pin the canvas to a window of the *specimen*. + + Without them the canvas follows each acquisition's own position bounding box, so two + scans of the same region reconstruct onto canvases of different shapes, offset from one + another by a fraction of the drift you are trying to measure. + """ + + @staticmethod + def _montage(dataset4d, defocus, position_offset=(0.0, 0.0), **kwargs): + dataset3d = Dataset3d.from_array( + dataset4d.array.reshape(-1, N, N), + name="synthetic 3D stack", + units=("index", "A^-1", "A^-1"), + sampling=(1, RECIPROCAL_SAMPLING, RECIPROCAL_SAMPLING), + ) + n_scan = dataset4d.shape[0] + positions = scan_positions_px()[: n_scan * n_scan] * SCAN_SAMPLING + positions = positions + np.asarray(position_offset, dtype=np.float64) + return DirectPtychographyMontage.from_dataset3d( + dataset3d, + positions, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + boundary="pad", + **_common_kwargs(defocus), + **kwargs, + ) + + def test_scan_origin_recovers_the_input_positions(self, dataset4d): + """`scan_origin + positions_px * scan_sampling` must be what was passed in.""" + offset = (37.0, -12.5) + montage = self._montage(dataset4d, integer_shift_defocus(1), position_offset=offset) + + n_scan = dataset4d.shape[0] + expected = scan_positions_px()[: n_scan * n_scan] * SCAN_SAMPLING + np.asarray(offset) + recovered = np.asarray(montage.scan_origin) + to_numpy(montage.positions_px) * np.asarray( + montage.scan_sampling + ) + + assert recovered == pytest.approx(expected, abs=1e-4) + + def test_pinned_window_is_reported_back(self, dataset4d): + montage = self._montage(dataset4d, integer_shift_defocus(1)) + origin = (-40.0, -30.0) + fov = (SCAN_SAMPLING * 20, SCAN_SAMPLING * 24) + + montage.reconstruct( + deconvolution_kernel="prlx", obj_origin=origin, obj_fov=fov, verbose=False + ) + + assert montage.obj.shape == (20, 24) + assert montage.obj_origin == pytest.approx(origin, abs=1e-4) + assert montage._obj_fov == pytest.approx(fov, rel=1e-6) + + def test_offset_acquisitions_land_on_the_same_canvas(self, dataset4d): + """The headline: the same specimen scanned in shifted coordinates must agree. + + The two montages differ only in a rigid offset applied to *both* the positions and + the requested window, so the reconstructions have to be the same image -- which is + exactly what makes a cross-correlation between frames measure drift and nothing else. + """ + defocus = integer_shift_defocus(1) + # deliberately not a whole number of scan pixels, which is the case that used to + # leave the two canvases a fraction of a pixel apart + offset = np.array([2.3 * SCAN_SAMPLING, -1.7 * SCAN_SAMPLING]) + origin = np.array([-3.0 * SCAN_SAMPLING, -3.0 * SCAN_SAMPLING]) + fov = (SCAN_SAMPLING * 24, SCAN_SAMPLING * 24) + + plain = self._montage(dataset4d, defocus) + moved = self._montage(dataset4d, defocus, position_offset=offset) + + plain.reconstruct( + deconvolution_kernel="prlx", obj_origin=tuple(origin), obj_fov=fov, verbose=False + ) + moved.reconstruct( + deconvolution_kernel="prlx", + obj_origin=tuple(origin + offset), + obj_fov=fov, + verbose=False, + ) + + assert moved.obj.shape == plain.obj.shape + assert moved.obj_origin == pytest.approx(tuple(origin + offset), abs=1e-4) + assert _relative_error(moved.obj, plain.obj) < 1e-5 + + def test_unpinned_canvas_is_reproduced_by_pinning_it(self, dataset4d): + """Reading the window back out and passing it in must be a no-op.""" + montage = self._montage(dataset4d, integer_shift_defocus(1)) + montage.reconstruct(deconvolution_kernel="prlx", pad_px=4, verbose=False) + automatic = montage.obj.copy() + + montage.reconstruct( + deconvolution_kernel="prlx", + obj_origin=montage.obj_origin, + obj_fov=montage._obj_fov, + verbose=False, + ) + + assert montage.obj.shape == automatic.shape + assert _relative_error(montage.obj, automatic) < 1e-6 + + def test_obj_fov_and_pad_px_are_mutually_exclusive(self, dataset4d): + montage = self._montage(dataset4d, integer_shift_defocus(1)) + with pytest.raises(ValueError, match="both size the canvas"): + montage.reconstruct( + deconvolution_kernel="prlx", pad_px=4, obj_fov=(100.0, 100.0), verbose=False + ) + + def test_wrap_boundary_accepts_a_pinned_canvas(self, dataset4d): + """Pinning the canvas is orthogonal to the boundary rule. + + ``"wrap"`` defaults to the scan grid, but that is a default rather than a + restriction: given a window, it wraps into that window instead. Needed whenever the + canvas has to be a particular size for reasons other than the scan -- matching an + empirical probe's reciprocal grid, say. + """ + montage = self._montage(dataset4d, integer_shift_defocus(1)) + fov = (SCAN_SAMPLING * 20, SCAN_SAMPLING * 24) + + montage.reconstruct( + deconvolution_kernel="prlx", + boundary="wrap", + obj_origin=(0.0, 0.0), + obj_fov=fov, + verbose=False, + ) + + assert montage.obj.shape == (20, 24) + assert montage.obj_origin == pytest.approx((0.0, 0.0), abs=1e-4) + assert np.isfinite(montage.obj).all() + + def test_a_wrapped_canvas_smaller_than_the_scan_warns(self, dataset4d): + """Positions outside a wrapped canvas fold a second copy of the specimen over it. + + The failure looks like a real reconstruction with a ghost in it, which is easy to + blame on something else -- so it has to say so. + """ + montage = self._montage(dataset4d, integer_shift_defocus(1)) + span = np.ptp(to_numpy(montage.positions_px), axis=0) * SCAN_SAMPLING + + with pytest.warns(UserWarning, match="wrap around"): + montage.reconstruct( + deconvolution_kernel="prlx", + boundary="wrap", + obj_origin=tuple(np.asarray(montage.scan_origin)), + obj_fov=(span[0] / 2, span[1]), + verbose=False, + ) + + @pytest.mark.parametrize("boundary", ["wrap", "pad"]) + def test_a_canvas_covering_the_scan_is_silent(self, dataset4d, boundary, recwarn): + montage = self._montage(dataset4d, integer_shift_defocus(1)) + span = np.ptp(to_numpy(montage.positions_px), axis=0) * SCAN_SAMPLING + + montage.reconstruct( + deconvolution_kernel="prlx", + boundary=boundary, + obj_origin=tuple(np.asarray(montage.scan_origin) - 4 * SCAN_SAMPLING), + obj_fov=tuple(span + 8 * SCAN_SAMPLING), + verbose=False, + ) + + assert not [w for w in recwarn if "wrap around" in str(w.message)] + + def test_wrap_defaults_to_the_scan_grid(self, dataset4d): + """The default is unchanged: no window given, the canvas is the scan.""" + montage = self._montage(dataset4d, integer_shift_defocus(1)) + montage.reconstruct(deconvolution_kernel="prlx", boundary="wrap", verbose=False) + + assert montage.obj.shape == tuple(montage.scan_gpts) + + def test_obj_fov_must_be_positive(self, dataset4d): + montage = self._montage(dataset4d, integer_shift_defocus(1)) + with pytest.raises(ValueError, match="positive"): + montage.reconstruct(deconvolution_kernel="prlx", obj_fov=(0.0, 10.0), verbose=False) + + def test_scan_origin_survives_a_round_trip(self, dataset4d, tmp_path): + montage = self._montage(dataset4d, integer_shift_defocus(1), position_offset=(11.0, -3.0)) + montage.reconstruct(deconvolution_kernel="prlx", verbose=False) + + path = str(tmp_path / "montage.zip") + montage.save(path, mode="o") + restored = load(path) + + assert restored.scan_origin == pytest.approx(montage.scan_origin) + assert restored.obj_origin == pytest.approx(montage.obj_origin) + + def test_gridded_construction_has_a_zero_scan_origin(self, dataset4d): + """A raster acquisition's grid *defines* the coordinates, so its origin is zero.""" + _, montage = _build_pair(dataset4d, integer_shift_defocus(1)) + assert montage.scan_origin == (0.0, 0.0) + + +class TestFrameDrift: + """`estimate_frame_drift` on interleaved frames with a seeded, known displacement.""" + + #: a window well inside the object, so every frame's canvas is fully supported and the + #: correlation is not measuring which frame happened to reach nearer the edge + ORIGIN = (8 * CTF_SAMPLING, 8 * CTF_SAMPLING) + FOV = (48 * CTF_SAMPLING, 48 * CTF_SAMPLING) + DRIFT = np.array([[0.0, 0.0], [1.6, -0.8], [3.2, -1.6], [4.8, -2.4]]) * CTF_SAMPLING + + def _frames(self, ctf_scene, drift=None): + complex_obj, _, probe, _, _ = ctf_scene + drift = self.DRIFT if drift is None else drift + datasets, positions = ctf_interleaved_frames(complex_obj, probe, drift) + + montages = [] + for dataset, pos in zip(datasets, positions): + montage = DirectPtychographyMontage.from_dataset3d( + dataset, + pos, + scan_sampling=(CTF_SCAN_SAMPLING, CTF_SCAN_SAMPLING), + boundary="pad", + **ctf_kwargs(), + ) + montage.reconstruct( + deconvolution_kernel="prlx", + obj_origin=self.ORIGIN, + obj_fov=self.FOV, + verbose=False, + ) + montages.append(montage) + return montages, positions + + def test_recovers_a_seeded_drift(self, ctf_scene): + montages, _ = self._frames(ctf_scene) + measured = estimate_frame_drift(montages, verbose=False) + + # the estimate is referred to the mean over frames, so compare mean-centered + expected = self.DRIFT - self.DRIFT.mean(axis=0) + assert measured == pytest.approx(expected, abs=0.35 * CTF_SCAN_SAMPLING) + + def test_drift_sums_to_zero(self, ctf_scene): + montages, _ = self._frames(ctf_scene) + measured = estimate_frame_drift(montages, verbose=False) + assert measured.mean(axis=0) == pytest.approx((0.0, 0.0), abs=1e-9) + + def test_correcting_the_positions_sharpens_the_combined_reconstruction(self, ctf_scene): + """The point of the exercise: `positions - drift` must beat `positions`. + + Also pins the sign, which no amount of reasoning about correlation conventions + substitutes for. + """ + complex_obj, _, probe, ctf_full, _ = ctf_scene + datasets, positions = ctf_interleaved_frames(complex_obj, probe, self.DRIFT) + montages, _ = self._frames(ctf_scene) + drift = estimate_frame_drift(montages, verbose=False) + + combined = Dataset3d.from_array( + np.concatenate([d.array for d in datasets]), + name="combined frames", + units=datasets[0].units, + sampling=datasets[0].sampling, + ) + + def reconstruct(position_list): + montage = DirectPtychographyMontage.from_dataset3d( + combined, + np.concatenate(position_list), + scan_sampling=(CTF_SCAN_SAMPLING, CTF_SCAN_SAMPLING), + boundary="pad", + **ctf_kwargs(), + ) + return montage.reconstruct( + deconvolution_kernel="prlx", + obj_origin=self.ORIGIN, + obj_fov=self.FOV, + verbose=False, + ) + + uncorrected = reconstruct(positions) + corrected = reconstruct([p - d for p, d in zip(positions, drift)]) + + # drift smears the montage, so undoing it has to raise the contrast + assert corrected.obj.std() > uncorrected.obj.std() + # and the wrong sign must make it worse, not better + wrong_sign = reconstruct([p + d for p, d in zip(positions, drift)]) + assert corrected.obj.std() > wrong_sign.obj.std() + + def test_zero_drift_is_recovered_as_zero(self, ctf_scene): + no_drift = np.zeros((3, 2)) + montages, _ = self._frames(ctf_scene, drift=no_drift) + measured = estimate_frame_drift(montages, verbose=False) + assert np.abs(measured).max() < 0.35 * CTF_SCAN_SAMPLING + + def test_rejects_frames_on_different_canvases(self, ctf_scene): + montages, _ = self._frames(ctf_scene) + montages[1].reconstruct( + deconvolution_kernel="prlx", + obj_origin=self.ORIGIN, + obj_fov=(self.FOV[0] + 4 * CTF_SAMPLING, self.FOV[1]), + verbose=False, + ) + with pytest.raises(ValueError, match="share a canvas"): + estimate_frame_drift(montages, verbose=False) + + def test_rejects_frames_at_different_origins(self, ctf_scene): + """The subtle failure: same shape, wrong place, silently measuring the offset.""" + montages, _ = self._frames(ctf_scene) + montages[1].reconstruct( + deconvolution_kernel="prlx", + obj_origin=(self.ORIGIN[0] + 3 * CTF_SAMPLING, self.ORIGIN[1]), + obj_fov=self.FOV, + verbose=False, + ) + with pytest.raises(ValueError, match="not a canvas origin"): + estimate_frame_drift(montages, verbose=False) + + def test_rejects_unreconstructed_frames(self, ctf_scene): + montages, _ = self._frames(ctf_scene) + montages[1]._reset_reconstruction() + with pytest.raises(ValueError, match="not been reconstructed"): + estimate_frame_drift(montages, verbose=False) + + def test_needs_at_least_two_frames(self, ctf_scene): + montages, _ = self._frames(ctf_scene) + with pytest.raises(ValueError, match="at least two"): + estimate_frame_drift(montages[:1], verbose=False) + + def test_works_on_gridded_reconstructions(self, dataset4d): + """Nothing about it is montage-specific -- two raster reconstructions also align.""" + defocus = integer_shift_defocus(1) + first, second = ( + DirectPtychography.from_dataset4d( + dataset4d, edge_blend_pixels=0, **_common_kwargs(defocus) + ) + for _ in range(2) + ) + first.reconstruct(deconvolution_kernel="prlx", verbose=False) + second.reconstruct(deconvolution_kernel="prlx", verbose=False) + + drift = estimate_frame_drift([first, second], verbose=False) + assert np.abs(drift).max() < 1e-6 + + +@pytest.mark.skipif(not ACCELERATORS, reason="no accelerator available") +@pytest.mark.parametrize("device", ACCELERATORS) +class TestAccelerators: + """Both classes must run, and agree with CPU, on whatever accelerator is present. + + MPS has no float64 at all, so every positional tensor -- coordinates, shifts, canvas + origins -- has to take its dtype from the device rather than hardcode one. + """ + + @staticmethod + def _kwargs(dataset4d): + return dict(_common_kwargs(integer_shift_defocus(1)), edge_blend_pixels=0) + + @staticmethod + def _ungridded(dataset4d): + """A `Dataset3d` and positions, with the origin *fitted* rather than forced. + + Forcing the origin skips the centre-of-mass fit entirely, which is where the + device mismatch lived, so these deliberately let it run. + """ + dataset3d = Dataset3d.from_array( + np.asarray(dataset4d.array).reshape(-1, N, N), + name="ungridded", + sampling=(1.0, dataset4d.sampling[-2], dataset4d.sampling[-1]), + units=("index", "A^-1", "A^-1"), + ) + return dataset3d, scan_positions_px() * SCAN_SAMPLING + + @pytest.mark.parametrize("boundary", ["wrap", "pad"]) + def test_montage_matches_cpu(self, dataset4d, device, boundary): + def run(where): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, device=where, boundary=boundary, **self._kwargs(dataset4d) + ) + return montage.reconstruct(verbose=False).obj + + on_device, on_cpu = run(device), run("cpu") + + assert on_device.shape == on_cpu.shape + assert np.abs(on_device - on_cpu).max() / np.abs(on_cpu).max() < 1e-4 + + def test_montage_upsampled_matches_cpu(self, dataset4d, device): + def run(where): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, device=where, **self._kwargs(dataset4d) + ) + return montage.reconstruct( + upsampling_factor=2, interpolation="bilinear", verbose=False + ).obj + + assert np.abs(run(device) - run("cpu")).max() / np.abs(run("cpu")).max() < 1e-4 + + def test_defocus_gradient_matches_cpu(self, dataset4d, device): + def run(where): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, + device=where, + boundary="pad", + defocus_gradient=(30.0, -10.0), + **self._kwargs(dataset4d), + ) + return montage.reconstruct(verbose=False).obj + + on_device, on_cpu = run(device), run("cpu") + + assert on_device.shape == on_cpu.shape + assert np.abs(on_device - on_cpu).max() / np.abs(on_cpu).max() < 1e-4 + + def test_real_space_kernel_runs(self, dataset4d, device): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, device=device, **self._kwargs(dataset4d) + ) + obj = montage.reconstruct(deconvolution_kernel="ssb", stencil_radius=4, verbose=False).obj + + assert np.isfinite(obj).all() + + def test_variance_loss_and_search_run(self, dataset4d, device): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, device=device, **self._kwargs(dataset4d) + ) + montage.reconstruct(verbose=False) + + assert np.isfinite(float(montage.variance_loss())) + + @pytest.mark.parametrize( + "cls", [DirectPtychography, DirectPtychographyMontage], ids=["fourier", "montage"] + ) + def test_ungridded_fits_the_origin_on_device(self, dataset4d, device, cls): + """Regression: probe positions arrive as numpy and met the measured origin on MPS.""" + dataset3d, positions = self._ungridded(dataset4d) + + reconstruction = cls.from_dataset3d( + dataset3d, + positions, + energy=PROBE_ENERGY, + semiangle_cutoff=SEMIANGLE_CUTOFF, + rotation_angle=0.0, + scan_sampling=(SCAN_SAMPLING, SCAN_SAMPLING), + device=device, + verbose=False, + ) + obj = reconstruction.reconstruct(deconvolution_kernel="prlx", verbose=False).obj + + assert np.isfinite(obj).all() + + +@pytest.mark.skipif(not ACCELERATORS, reason="no accelerator available") +@pytest.mark.parametrize("device", ACCELERATORS) +def test_padded_canvas_shape_is_device_independent(dataset4d, device): + """Float noise must not push a canvas bound across an integer on one device only.""" + + def shape(where): + montage = DirectPtychographyMontage.from_dataset4d( + dataset4d, + device=where, + boundary="pad", + **dict(_common_kwargs(integer_shift_defocus(1)), edge_blend_pixels=0), + ) + return montage.reconstruct(verbose=False).obj.shape + + assert shape(device) == shape("cpu") diff --git a/tests/diffractive_imaging/test_object_tensor_decomp.py b/tests/diffractive_imaging/test_object_tensor_decomp.py index 9343da75..17d87055 100644 --- a/tests/diffractive_imaging/test_object_tensor_decomp.py +++ b/tests/diffractive_imaging/test_object_tensor_decomp.py @@ -417,7 +417,7 @@ def test_potential_identity_default_and_positivity_penalty(self): 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 + 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) diff --git a/tests/utils/test_imaging_utils.py b/tests/utils/test_imaging_utils.py new file mode 100644 index 00000000..b8688de2 --- /dev/null +++ b/tests/utils/test_imaging_utils.py @@ -0,0 +1,99 @@ +import numpy as np +import pytest +from scipy.ndimage import fourier_shift, gaussian_filter + +from quantem.core.utils.imaging_utils import ( + cross_correlation_shift, + cross_correlation_shift_torch, + dft_upsample, +) + +# shifts spanning integer, half-integer, sub-pixel and large displacements +SHIFTS = [(0.0, 0.0), (3.0, -2.0), (1.25, 0.75), (-4.5, 2.5), (0.37, -0.11), (-7.3, 5.9)] + + +@pytest.fixture(scope="module") +def image(): + rng = np.random.default_rng(0) + return gaussian_filter(rng.normal(size=(64, 64)), 2) + + +def _shifted(image, shift): + return np.real(np.fft.ifft2(fourier_shift(np.fft.fft2(image), shift))) + + +class TestDftUpsample: + def test_agrees_with_ifft2_at_unit_upsampling(self, image): + """At `up=1` the fine grid lands on integer pixels, where `ifft2` is the answer.""" + spectrum = np.fft.fft2(image) + du = int(np.ceil(1.5 * 1)) + row, col = 5, 7 + + local = dft_upsample(spectrum, 1, (row, col)) + expected = np.real(np.fft.ifft2(spectrum))[ + row - du : row + du + 1, col - du : col + du + 1 + ] + + # `dft_upsample` omits the 1/(M*N) normalization `ifft2` applies + assert np.allclose(local, expected * spectrum.size, atol=1e-8) + + @pytest.mark.parametrize("up", [2, 4, 8]) + def test_peak_lands_on_the_center_tap(self, image, up): + """Centering the fine grid on the true peak must put the maximum at its middle. + + The middle is `ceil(1.5 * up)`, since the grid spans `arange(-du, du+1) / up`. + """ + cc = np.fft.fft2(image) * np.conj(np.fft.fft2(image)) + local = dft_upsample(cc, up, (0.0, 0.0)) + + du = int(np.ceil(1.5 * up)) + assert local.shape == (2 * du + 1, 2 * du + 1) + assert np.unravel_index(np.argmax(local), local.shape) == (du, du) + + +class TestCrossCorrelationShift: + @pytest.mark.parametrize("shift", SHIFTS) + @pytest.mark.parametrize("up", [2, 4, 8, 16]) + def test_recovers_a_known_shift(self, image, shift, up): + """The returned shift realigns `im` onto `im_ref`, so it is the negated input.""" + measured = np.asarray(cross_correlation_shift(image, _shifted(image, shift), up)) + assert measured == pytest.approx(-np.asarray(shift), abs=0.02) + + def test_upsampling_beats_no_upsampling(self, image): + """Sub-pixel refinement has to be an improvement on the integer peak.""" + shift = (1.25, 0.75) + shifted = _shifted(image, shift) + + coarse = np.asarray(cross_correlation_shift(image, shifted, upsample_factor=1)) + fine = np.asarray(cross_correlation_shift(image, shifted, upsample_factor=8)) + + want = -np.asarray(shift) + assert np.abs(fine - want).max() < np.abs(coarse - want).max() + + def test_identical_images_do_not_shift(self, image): + """Guards the off-by-`ceil(1.5*up)` that used to bias every result by half a pixel.""" + for up in (2, 4, 8, 16): + measured = np.asarray(cross_correlation_shift(image, image, upsample_factor=up)) + assert measured == pytest.approx((0.0, 0.0), abs=1e-6) + + @pytest.mark.parametrize("shift", SHIFTS) + def test_matches_the_torch_implementation(self, image, shift): + """The two paths are independent ports; they must not disagree.""" + import torch + + shifted = _shifted(image, shift) + numpy_shift = np.asarray(cross_correlation_shift(image, shifted, upsample_factor=8)) + torch_shift = cross_correlation_shift_torch( + torch.as_tensor(image), torch.as_tensor(shifted), upsample_factor=8 + ).numpy() + + assert numpy_shift == pytest.approx(torch_shift, abs=0.02) + + def test_shifted_image_is_realigned(self, image): + """`return_shifted_image` must undo the displacement it just measured.""" + shifted = _shifted(image, (2.5, -1.5)) + _, realigned = cross_correlation_shift( + image, shifted, upsample_factor=8, return_shifted_image=True + ) + + assert np.corrcoef(realigned.ravel(), image.ravel())[0, 1] > 0.999