From 26f2ff7d18d637bbad3da6c644bd463b16f34074 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Mon, 24 Aug 2026 21:25:07 -0700 Subject: [PATCH 01/11] refactor: split drift implementation into a package --- src/quantem/imaging/{drift.py => drift/__init__.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/quantem/imaging/{drift.py => drift/__init__.py} (100%) diff --git a/src/quantem/imaging/drift.py b/src/quantem/imaging/drift/__init__.py similarity index 100% rename from src/quantem/imaging/drift.py rename to src/quantem/imaging/drift/__init__.py From 9179261dd19e5254dbb3c7380936e63fb4829233 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Mon, 24 Aug 2026 21:42:01 -0700 Subject: [PATCH 02/11] refactor: separate drift package responsibilities --- src/quantem/imaging/drift/__init__.py | 1838 +------------------- src/quantem/imaging/drift/apply.py | 399 +++++ src/quantem/imaging/drift/core/__init__.py | 1 + src/quantem/imaging/drift/core/affine.py | 378 ++++ src/quantem/imaging/drift/core/nonrigid.py | 449 +++++ src/quantem/imaging/drift/core/warping.py | 196 +++ src/quantem/imaging/drift/correction.py | 186 ++ src/quantem/imaging/drift/plot.py | 111 ++ src/quantem/imaging/drift/preparation.py | 203 +++ 9 files changed, 1931 insertions(+), 1830 deletions(-) create mode 100644 src/quantem/imaging/drift/apply.py create mode 100644 src/quantem/imaging/drift/core/__init__.py create mode 100644 src/quantem/imaging/drift/core/affine.py create mode 100644 src/quantem/imaging/drift/core/nonrigid.py create mode 100644 src/quantem/imaging/drift/core/warping.py create mode 100644 src/quantem/imaging/drift/correction.py create mode 100644 src/quantem/imaging/drift/plot.py create mode 100644 src/quantem/imaging/drift/preparation.py diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index 415789d8..a19bef65 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -1,1832 +1,10 @@ -import warnings -from collections.abc import Sequence -from typing import Self - -import matplotlib.pyplot as plt -import numpy as np -import torch -from torch.fft import fftfreq -from numpy.typing import NDArray -from scipy.interpolate import interp1d -from scipy.ndimage import distance_transform_edt, gaussian_filter -from scipy.optimize import minimize -from tqdm import tqdm - -from quantem.core.config import validate_device - -from quantem.core.datastructures.dataset2d import Dataset2d -from quantem.core.datastructures.dataset3d import Dataset3d -from quantem.core.io.serialize import AutoSerialize -from quantem.core.utils.compound_validators import ( - validate_list_of_dataset2d, - validate_pad_value, -) -from quantem.core.utils.imaging_utils import ( - bilinear_kde, - cross_correlation_shift, - fourier_cropping, -) -from quantem.imaging.drift_utils import ( - bilinear_kde_batch, - cross_corr_batch, - gaussian_smooth_1d, - transform_coordinates_single_knot, - translate_align, +from quantem.imaging.drift.correction import DriftCorrection as DriftCorrection +from quantem.imaging.drift.core.warping import ( + DriftInterpolator as DriftInterpolator, + _bounded_sine_sigmoid_torch as _bounded_sine_sigmoid_torch, + _fourier_crop_torch as _fourier_crop_torch, + bounded_sine_sigmoid as bounded_sine_sigmoid, ) -from quantem.core.utils.validators import ensure_valid_array -from quantem.core.visualization import show_2d - - -class DriftCorrection(AutoSerialize): - """ - DriftCorrection provides translation, affine, and non-rigid drift correction for - sequential 2D images using scan direction metadata and flexible spatial interpolation. - - This class supports input data as numpy arrays, Dataset2d, or Dataset3d instances, - with various padding strategies and configurable spline interpolation of scanline - trajectories via Bézier knot control. - - Features - -------- - - Load data from arrays or files - - Apply initial scanline resampling using Bézier curves - - Align images using translation, affine, or non-rigid optimization - - Visualize intermediate and final results with optional knot overlays - - Serialize state with `.save()` and restore with `.load()` - - Parameters (via `from_data` or `from_file`) - ------------------------------------------- - images : list of 2D arrays, Dataset2d, Dataset3d, or file names, or a 3D numpy array - The image stack to correct for drift. - scan_direction_degrees : list of float - The scan direction angle (in degrees) for each image, measured relative to vertical. - pad_fraction : float, default 0.25 - Fraction of padding to add around each image during interpolation. - pad_value : str, float, or list of float, default 'median' - How to pad outside the image area during warping. Can be: - - One of: 'median', 'mean', 'min', 'max' - - A float quantile value (e.g., 0.25) - - A list of per-image float values - number_knots : int, default 1 - Number of knots to use for Bézier interpolation of scanline trajectories. - We strongly recommend using `number_knots = 1` unless the fast scan direction is - expected to vary within the image. - - Example - ------- - Instantiate the DriftCorrection class, run preprocessing and alignment, and save/load results: - - >>> drift = DriftCorrection.from_data( - ... images=[ - ... image0, # 2D numpy array or Dataset2d - ... image1, - ... ], - ... scan_direction_degrees=[0, 90], - ... ).preprocess( - ... pad_fraction=0.25, - ... pad_value='median', - ... number_knots=1, - ... ) - - >>> drift.align_affine() - >>> drift.align_nonrigid() - >>> drift.plot_merged_images() - >>> image_corr = drift.generate_corrected_image() - - >>> drift.save("drift_result.zip") - >>> drift_reloaded = quantem.io.load("drift_result.zip") - - >>> image_corr.save("image_corrected.zip") - >>> image_corr_reloaded = quantem.io.load("image_corrected.zip") - - Notes - ----- - - Use `align_translation()` for rigid shifts, `align_affine()` for scan-shear or uniform drift, - and `align_nonrigid()` for flexible per-row or per-image correction. - - The class stores resampled images in `self.images_warped` and the control knots in `self.knots`. - - Visualization is supported through `plot_merged_images()` and `plot_transformed_images()`. - - Performance - ----------- - ``align_affine`` uses PyTorch to run all heavy operations on GPU - (works on CUDA, MPS, and CPU). The key optimizations are: - - - **Batched grid search**: all ~97 candidate drift vectors are warped - and scored in parallel, instead of one-at-a-time in a Python loop. - - **Batched bilinear KDE** (``drift_utils.bilinear_kde_batch``): - scatter-based image warping via ``scatter_add_`` with int32 indices. - - **Batched FFT cross-correlation** (``drift_utils.cross_corr_batch``): - sub-pixel alignment using DFT upsampling across all candidates at once. - - **Zero CPU round-trips**: coordinate transforms, Gaussian smoothing, - translation alignment, and error computation all stay on GPU until - the final sync. - - This gives ~300× speedup over the original NumPy implementation - (e.g. 436 s → 1.5 s on 2048×2048 image pairs). - - Memory is automatically chunked when the full batch doesn't fit. - Approximate memory per candidate at common sizes: - - ========== =========== ================ - Input size Canvas size Mem / candidate - ========== =========== ================ - 1024×1024 1280×1280 85 MB - 2048×2048 2560×2560 341 MB - 4096×4096 5120×5120 1.36 GB - ========== =========== ================ - """ - - _token = object() - - def __init__( - self, - images: list[Dataset2d], - scan_direction_degrees: NDArray, - _token: object | None = None, - ): - if _token is not self._token: - raise RuntimeError( - "Use DriftCorrection.from_data() or .from_file() to instantiate this class." - ) - - self.images = images - self.scan_direction_degrees = ensure_valid_array(scan_direction_degrees, ndim=1) - - device, _ = validate_device(None) - self._device = device - self._dtype = torch.float32 - - @classmethod - def from_file( - cls, - file_paths: Sequence[str], - scan_direction_degrees: Sequence[float] | NDArray, - file_type: str | None = None, - ) -> Self: - image_list = [Dataset2d.from_file(fp, file_type=file_type) for fp in file_paths] - return cls.from_data( - image_list, - scan_direction_degrees, - ) - - @classmethod - def from_data( - cls, - images: list[Dataset2d] | list[NDArray] | Dataset3d | NDArray, - scan_direction_degrees: list[float] | NDArray, - ) -> Self: - validated_images = validate_list_of_dataset2d(images) - - return cls( - images=validated_images, - scan_direction_degrees=scan_direction_degrees, - _token=cls._token, - ) - - def preprocess( - self, - pad_fraction: float = 0.25, - pad_value: float | str | list[float] = "median", - kde_sigma: float = 0.5, - number_knots: int = 1, - show_merged: bool = False, - show_images: bool = False, - show_knots: bool = True, - **kwargs, - ): - """Prepare images for drift correction by building the scanline model. - - Computes scan direction vectors, initializes Bezier knots that map - each scanline onto a padded canvas, and generates the initial warped - images. This must be called before any alignment step. - - Without preprocessing, there is no spatial model connecting the raw - images to the shared canvas - alignment methods would have no - coordinates to optimize. - - Parameters - ---------- - pad_fraction : float - Fraction of the image size to add as padding around the canvas. - Larger values give more room for drift but use more memory. - ``pad_fraction=0.25`` adds 25% on each side. - pad_value : float, str, or list[float] - Fill value for pixels outside the image footprint. Can be - ``'median'``, ``'mean'``, ``'min'``, ``'max'``, a quantile - (e.g. ``0.25``), or a per-image list of floats. - kde_sigma : float - Gaussian smoothing sigma (in pixels) applied after bilinear - scatter. Smooths the warped images to reduce scatter noise. - number_knots : int - Number of Bezier knots per scanline. Use ``1`` (recommended) - for linear drift correction. Higher values allow per-scanline - curvature but are slower and rarely needed. - show_merged : bool - Display the merged (averaged) warped images after preprocessing. - show_images : bool - Display each individual warped image after preprocessing. - show_knots : bool - Overlay knot positions on displayed images. - **kwargs - Additional keyword arguments passed to plotting functions. - - Returns - ------- - Self - For method chaining: ``drift.preprocess().align_affine()``. - - Examples - -------- - >>> drift = DriftCorrection.from_data( - ... images=[im0, im1], scan_direction_degrees=[0, 90]) - >>> drift.preprocess(pad_fraction=0.25, kde_sigma=0.5, number_knots=1) - """ - self.pad_fraction = float(pad_fraction) - self.pad_value = validate_pad_value(pad_value, self.images) - self.kde_sigma = float(kde_sigma) - self.number_knots = int(number_knots) - self.scan_direction = np.deg2rad(self.scan_direction_degrees) - self.scan_fast = np.stack( - [np.sin(-self.scan_direction), np.cos(-self.scan_direction)], axis=1) - self.scan_slow = np.stack( - [np.cos(-self.scan_direction), -np.sin(-self.scan_direction)], axis=1) - self.shape = ( - len(self.images), - int(np.round(self.images[0].shape[0] * (1 + self.pad_fraction) / 2) * 2), - int(np.round(self.images[1].shape[1] * (1 + self.pad_fraction) / 2) * 2), - ) - # Initialize knots - each image's scanlines mapped to the padded canvas - self.knots = [] - for img_idx in range(self.shape[0]): - shape = self.images[img_idx].shape - v_slow = np.linspace(-(shape[0] - 1) / 2, (shape[0] - 1) / 2, shape[0]) - u_fast = np.linspace(-(shape[1] - 1) / 2, (shape[1] - 1) / 2, self.number_knots) - row_knots = ((self.shape[1] - 1) / 2 - + u_fast[None, :] * self.scan_fast[img_idx, 0] - + v_slow[:, None] * self.scan_slow[img_idx, 0]) - col_knots = ((self.shape[2] - 1) / 2 - + u_fast[None, :] * self.scan_fast[img_idx, 1] - + v_slow[:, None] * self.scan_slow[img_idx, 1]) - self.knots.append(np.stack([row_knots, col_knots], axis=0)) - self.interpolator = [ - DriftInterpolator( - input_shape=self.images[i].shape, - output_shape=self.shape[1:], - scan_fast=self.scan_fast[i], - scan_slow=self.scan_slow[i], - pad_value=self.pad_value[i], - kde_sigma=self.kde_sigma, - ) - for i in range(self.shape[0]) - ] - # Cache source data on GPU and generate initial warped images - device = self._device - dtype = self._dtype - self.images_t = [ - torch.tensor(self.images[i].array, dtype=dtype, device=device) - for i in range(self.shape[0]) - ] - self.scan_fast_t = [ - torch.tensor(self.scan_fast[i], dtype=dtype, device=device) - for i in range(self.shape[0]) - ] - self.images_warped = Dataset3d.from_shape(self.shape) - self.weights_warped = Dataset3d.from_shape(self.shape) - canvas_shape = (self.shape[1], self.shape[2]) - warped_t = torch.zeros(self.shape[0], *canvas_shape, dtype=dtype, device=device) - for img_idx in range(self.shape[0]): - knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) - warped, weights = bilinear_kde_batch( - row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, - self.kde_sigma, self.pad_value[img_idx]) - warped_t[img_idx] = warped[0] - self.images_warped.array[img_idx] = warped[0].cpu().numpy() - self.weights_warped.array[img_idx] = weights[0].cpu().numpy() - self.calculate_error(0, _warped_t=warped_t) - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images(show_knots=show_knots, title="Merged: initial", **kwargs) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: initial" for i in range(self.shape[0])], - **kwargs, - ) - return self - - def align_translation( - self, - upsample_factor: int = 8, - min_image_shift: float | None = None, - max_image_shift: float = 32, - show_merged: bool = True, - show_images: bool = False, - show_knots: bool = True, - **kwargs, - ): - """ - Solve for the translation between all images in DriftCorrection.images_warped - """ - dxy = np.zeros((self.shape[0], 2)) - F_ref = np.fft.fft2(self.images_warped.array[0]) - for ind in range(1, self.shape[0]): - shifts, image_shift = cross_correlation_shift( - F_ref, - np.fft.fft2(self.images_warped.array[ind]), - upsample_factor=upsample_factor, - max_shift=max_image_shift, - fft_input=True, - fft_output=True, - return_shifted_image=True, - ) - dxy[ind, :] = shifts - F_ref = F_ref * ind / (ind + 1) + image_shift / (ind + 1) - dxy -= np.mean(dxy, axis=0) - if min_image_shift is not None: - if np.linalg.norm(dxy[ind]) < min_image_shift: - dxy[ind] = 0.0 - for ind in range(self.shape[0]): - self.knots[ind][0] += dxy[ind, 0] - self.knots[ind][1] += dxy[ind, 1] - for ind in range(self.shape[0]): - self.images_warped.array[ind], self.weights_warped.array[ind] = self.interpolator[ - ind - ].warp_image( - self.images[ind].array, - self.knots[ind], - ) - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images(show_knots=show_knots, title="Merged: translation", **kwargs) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: translation" for i in range(self.shape[0])], - **kwargs, - ) - return self - - # Affine alignment - def align_affine( - self, - step: float = 0.01, - num_tests: int = 9, - refine: bool = True, - upsample_factor: int = 8, - max_image_shift: float | None = 32, - chunk_size: int | None = None, - show_merged: bool = True, - show_images: bool = False, - show_knots: bool = True, - verbose: bool = False, - **kwargs, - ): - """Correct affine drift between scan pairs using a batched grid search. - - Builds a grid of candidate linear-drift vectors, warps both images - for each candidate, and picks the one with the lowest cross-correlation - cost. An optional refinement pass subdivides the winning cell for - sub-step accuracy. Without affine correction, per-scanline drift - causes shear distortion that translation alignment alone cannot fix. - - Parameters - ---------- - step : float - Search resolution in pixels per scan line. The grid search - tests drift rates from ``-step * num_tests/2`` to - ``+step * num_tests/2`` px/line. For example, ``step=0.02`` - with ``num_tests=11`` searches drifts from -0.10 to +0.10 - px/line. Smaller values detect subtler drift but test more - candidates. - num_tests : int - Number of drift rates to test along each axis. Must be odd - so the grid is centered on zero drift. Total candidates - ≈ ``π/4 * num_tests²``: ``num_tests=5`` → 21, - ``num_tests=9`` → 61, ``num_tests=11`` → 97. - refine : bool - If True, run a second pass at ``step / (num_tests - 1)`` - resolution, centered on the coarse winner. - upsample_factor : int - Sub-pixel precision for measuring the translational shift - between warped image pairs. 8 means 1/8-pixel precision. - Higher values are more accurate but slower. - max_image_shift : float or None - Maximum allowed translational shift in pixels. Cross-correlation - peaks beyond this radius are masked to reject spurious matches - from noise or periodic artifacts. Set to None to allow any shift. - chunk_size : int or None - Number of candidates per pass. If None, all candidates at once. - Set to a smaller value if you run out of memory. - show_merged : bool - Display the merged (averaged) image after alignment. - show_images : bool - Display each individual warped image after alignment. - show_knots : bool - Overlay knot positions on the displayed images. - verbose : bool - If True, print the top 5 candidate drift vectors with their - cost and direction after the grid search. Useful for - diagnosing ambiguous alignments or verifying the winning - candidate has a clear margin over runner-ups. - **kwargs - Additional keyword arguments passed to the plotting functions. - - Returns - ------- - DriftCorrection - Self, for method chaining. - - Examples - -------- - >>> drift = DriftCorrection.from_data( - ... images=[im0, im1], scan_direction_degrees=[0, 90]) - >>> drift.preprocess().align_affine(step=0.02, num_tests=11) - """ - if self.shape[0] < 2: - raise ValueError( - f"align_affine requires at least 2 images (got {self.shape[0]}). " - f"Provide image pairs with different scan directions." - ) - if num_tests % 2 == 0: - raise ValueError( - f"num_tests must be odd (got {num_tests}). Try {num_tests + 1}." - ) - # Build candidate grid with circular mask (~21% fewer than square) - grid_axis = np.arange(-(num_tests - 1) / 2, (num_tests + 1) / 2) - row_grid, col_grid = np.meshgrid(grid_axis, grid_axis, indexing="ij") - circular_mask = row_grid**2 + col_grid**2 <= (num_tests / 2) ** 2 - drift_vectors = np.vstack((row_grid[circular_mask], col_grid[circular_mask])).T * step - - def _print_top_candidates(label, candidates, costs_tensor): - costs_np = costs_tensor.cpu().numpy() - ranked = np.argsort(costs_np) - best_cost = costs_np[ranked[0]] - print(f" {label} - top 5 candidates:") - for rank in range(min(5, len(ranked))): - idx = ranked[rank] - drift_row, drift_col = candidates[idx] - magnitude = np.sqrt(drift_row**2 + drift_col**2) - gap = (costs_np[idx] - best_cost) / best_cost * 100 if rank > 0 else 0 - print(f" drift=({drift_row:+.4f}, {drift_col:+.4f}) px/line " - f"({magnitude:.4f} magnitude), cost={costs_np[idx]:.4f}" - f"{f' (+{gap:.1f}%)' if rank > 0 else ' (best)'}") - - def _apply_drift(drift_vec): - for img_idx in range(self.shape[0]): - scanline_offset = np.arange(self.knots[img_idx].shape[1]) - (self.knots[img_idx].shape[1] - 1) / 2 - self.knots[img_idx][0] += drift_vec[0] * scanline_offset[:, None] - self.knots[img_idx][1] += drift_vec[1] * scanline_offset[:, None] - - def _search_and_apply(candidates, label): - best_idx, costs = self._affine_grid_search_batch(candidates, upsample_factor, max_image_shift, chunk_size) - _apply_drift(candidates[best_idx]) - if verbose: - _print_top_candidates(label, candidates, costs) - warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor) - self.calculate_error(1, _warped_t=warped_t) - return candidates[best_idx] - - drift_total = _search_and_apply(drift_vectors, "Coarse search") - if refine: - drift_fine = drift_vectors / (num_tests - 1) - drift_total = drift_total + _search_and_apply(drift_fine, "Refine search") - if verbose: - num_rows = self.images[0].shape[0] - drift_rate = np.sqrt(drift_total[0] ** 2 + drift_total[1] ** 2) - total_shift = drift_rate * num_rows - angle_deg = np.degrees(np.arctan2(drift_total[1], drift_total[0])) - print(f"align_affine: step={step}, num_tests={num_tests} " - f"({len(drift_vectors)} candidates), refine={refine}, " - f"max_image_shift={max_image_shift}") - msg = (f"Drift: ({drift_total[0]:+.4f}, {drift_total[1]:+.4f}) px/line, " - f"{drift_rate:.4f} magnitude, {angle_deg:.1f} deg, " - f"{total_shift:.1f} px total over {num_rows} lines") - if self.images[0].sampling is not None: - px_size = self.images[0].sampling[0] - unit = self.images[0].units[0] if self.images[0].units else "px" - msg += f" = {total_shift * px_size:.2f} {unit}" - print(msg) - err = self.error_track - print(f"Error: {err[0, 1]:.2f} -> {err[-1, 1]:.2f} " - f"({(err[0, 1] - err[-1, 1]) / err[0, 1] * 100:+.1f}%)") - - # Plots - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images( - show_knots=show_knots, - title="Merged: affine", - **kwargs, - ) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: affine" for i in range(self.shape[0])], - **kwargs, - ) - - return self - - @torch.inference_mode() - def _affine_grid_search_batch(self, drift_vectors, upsample_factor, max_image_shift, chunk_size=None): - """Evaluate all candidate drift vectors in parallel. - - Warps both images for each candidate using ``bilinear_kde_batch`` - and scores alignment quality via ``cross_corr_batch``. Without - batching, each candidate would be a separate Python iteration - this - is the key operation that enables the 300x speedup. - - Parameters - ---------- - drift_vectors : ndarray, shape (N, 2) - Candidate drift vectors to test, columns are (row, col). - upsample_factor : int - Subpixel cross-correlation upsampling factor. - max_image_shift : float or None - Maximum allowed shift for cross-correlation peak search. - chunk_size : int or None - Number of candidates per pass. If None, all at once. - - Returns - ------- - tuple[int, torch.Tensor] - Index of the best candidate in ``drift_vectors``, and the full - cost tensor of shape ``(N,)`` for all candidates (used by - verbose mode to rank runner-ups). - """ - device = self._device - dtype = self._dtype - num_candidates = drift_vectors.shape[0] - drift_vectors_t = torch.tensor(drift_vectors, dtype=dtype, device=device) - canvas_shape = (self.shape[1], self.shape[2]) - # Base coordinates shared across all candidates - base_data = [] - for img_idx in range(2): - knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) - row_base, col_base = transform_coordinates_single_knot( - knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) - num_rows = self.knots[img_idx].shape[1] - scanline_offset = (torch.arange(num_rows, dtype=dtype, device=device) - - (num_rows - 1) / 2) - base_data.append((self.images_t[img_idx], row_base, col_base, scanline_offset)) - # Precompute shift mask and frequency grids (shared across chunks) - shift_mask = None - if max_image_shift is not None: - canvas_rows, canvas_cols = canvas_shape - freq_row = fftfreq(canvas_rows, 1.0 / canvas_rows, device=device, dtype=dtype) - freq_col = fftfreq(canvas_cols, 1.0 / canvas_cols, device=device, dtype=dtype) - shift_mask = freq_row[:, None] ** 2 + freq_col[None, :] ** 2 >= max_image_shift ** 2 - freq_grids = ( - fftfreq(canvas_shape[0], device=device, dtype=dtype)[:, None], - fftfreq(canvas_shape[1], device=device, dtype=dtype)[None, :], - ) - if chunk_size is None: - chunk_size = self._auto_chunk_size(num_candidates, canvas_shape, dtype, device) - on_cuda = torch.device(device).type == "cuda" - chunked = on_cuda and chunk_size < num_candidates - all_costs = [] - chunk_start = 0 - chunk_idx = 0 - while chunk_start < num_candidates: - chunk_end = min(chunk_start + chunk_size, num_candidates) - drift_chunk = drift_vectors_t[chunk_start:chunk_end] - if chunk_idx == 0 and chunked: - torch.cuda.reset_peak_memory_stats(device) - warped_pair = [] - for img_idx in range(2): - image_t, row_base, col_base, scanline_offset = base_data[img_idx] - row_candidates = row_base[None] + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] - col_candidates = col_base[None] + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] - warped, _ = bilinear_kde_batch( - row_candidates, col_candidates, image_t, - canvas_shape, self.kde_sigma, - self.pad_value[img_idx]) - warped_pair.append(warped) - all_costs.append(cross_corr_batch( - warped_pair[0], warped_pair[1], - upsample_factor, - max_shift_mask=shift_mask, - freq_grids=freq_grids)) - # After chunk 0, replace the conservative static estimate with the - # actual measured per-candidate cost and print one summary line so - # the user can see how the chunking adapted to their GPU state. - if chunk_idx == 0 and chunked: - per_candidate_actual = torch.cuda.max_memory_allocated(device) / chunk_size - free_bytes, total_bytes = torch.cuda.mem_get_info(device) - tuned_chunk_size = max(1, int(free_bytes * 0.5 / per_candidate_actual)) - tuned_chunk_size = min(tuned_chunk_size, num_candidates) - if tuned_chunk_size > chunk_size: - chunk_size = tuned_chunk_size - num_chunks_final = 1 + (num_candidates - chunk_end + chunk_size - 1) // chunk_size - print( - f" affine grid: {num_candidates} cand × {canvas_shape[0]}×{canvas_shape[1]}, " - f"{per_candidate_actual / 1e9:.2f} GB/cand → {chunk_size}/chunk × {num_chunks_final} passes " - f"({free_bytes / 1e9:.0f}/{total_bytes / 1e9:.0f} GB free)" - ) - chunk_start = chunk_end - chunk_idx += 1 - all_costs = torch.cat(all_costs) - return torch.argmin(all_costs).item(), all_costs - - @staticmethod - def _auto_chunk_size(num_candidates, canvas_shape, dtype, device): - """Pick a candidate-batch size that fits in current free GPU memory. - - Empirical per-candidate peak (measured at 4096×4096): bilinear KDE - scatter buffers, gaussian smoothing temporaries, then cross-correlation - FFT pairs (complex64) - together about ``32 × canvas_pixels`` - ``× dtype_bytes`` at peak. We sample free memory at call time, divide - by that estimate with a 0.4 safety factor, and cap the result at - ``num_candidates`` (no point splitting if it all fits). - On CPU we just process all candidates at once - no VRAM constraint. - """ - device = torch.device(device) - if device.type != "cuda": - return num_candidates - bytes_per_element = torch.finfo(dtype).bits // 8 - per_candidate_bytes = canvas_shape[0] * canvas_shape[1] * bytes_per_element * 32 - free_bytes, _ = torch.cuda.mem_get_info(device) - chunk_size = max(1, int(free_bytes * 0.4 / per_candidate_bytes)) - return min(chunk_size, num_candidates) - - @torch.inference_mode() - def _warp_and_translate_torch( - self, - max_image_shift: float | None, - upsample_factor: int = 8, - knots_batch: torch.Tensor | None = None, - solve_translation: bool = True, - ) -> torch.Tensor: - """Regenerate warped images and solve translation on GPU. - - Three phases: warp → solve translation → re-warp. When ``knots_batch`` - is provided, reads/writes a single batched torch tensor (zero numpy - crossings). Without it, falls back to ``self.knots`` (numpy) for - compatibility with ``align_affine``. - - Set ``solve_translation=False`` to only warp and sync without - re-solving translation - used after the nonrigid loop to populate - ``self.images_warped`` from final knots. - - Parameters - ---------- - max_image_shift : float or None - Maximum allowed translational shift in pixels. - upsample_factor : int - Sub-pixel precision for cross-correlation (1/N pixel). - knots_batch : torch.Tensor or None - If provided, batched ``(N, 2, num_rows)`` torch tensor on GPU. - Translation shifts are applied in-place. Skips numpy sync. - solve_translation : bool - If False, skip translation alignment (Phase 2+3). Only warp - once using current knots and sync to CPU. - - Returns - ------- - torch.Tensor - Warped images on GPU, shape ``(num_images, H, W)``. - """ - device = self._device - dtype = self._dtype - num_images = self.shape[0] - canvas_shape = (self.shape[1], self.shape[2]) - - def _warp_all(warped_t, weights_t): - """Warp all images onto the canvas using current knots.""" - for img_idx in range(num_images): - if knots_batch is not None: - # transform_coordinates_single_knot expects (2, N, 1) - knots_img = knots_batch[img_idx].detach()[:, :, None] - else: - knots_img = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_img, self.scan_fast_t[img_idx], self.images[img_idx].shape) - warped, weights = bilinear_kde_batch( - row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, - self.kde_sigma, self.pad_value[img_idx]) - warped_t[img_idx] = warped[0] - weights_t[img_idx] = weights[0] - - warped_t = torch.zeros(num_images, *canvas_shape, dtype=dtype, device=device) - weights_t = torch.zeros_like(warped_t) - _warp_all(warped_t, weights_t) - if not solve_translation: - self.images_warped.array[:] = warped_t.cpu().numpy() - self.weights_warped.array[:] = weights_t.cpu().numpy() - return warped_t - # Solve translation shifts and apply to knots - shifts_t = translate_align(warped_t, upsample_factor, max_image_shift) - if knots_batch is not None: - knots_batch[:, 0] += shifts_t[:, 0:1] - knots_batch[:, 1] += shifts_t[:, 1:2] - else: - shifts_np = shifts_t.cpu().numpy() - for img_idx in range(num_images): - self.knots[img_idx][0] += shifts_np[img_idx, 0] - self.knots[img_idx][1] += shifts_np[img_idx, 1] - # Re-warp with corrected knots - _warp_all(warped_t, weights_t) - if knots_batch is None: - self.images_warped.array[:] = warped_t.cpu().numpy() - self.weights_warped.array[:] = weights_t.cpu().numpy() - return warped_t - - def align_nonrigid( - self, - backend: str = "pytorch", - optimizer_name: str = "adam", - num_iterations: int = 8, - regularization_sigma_px: float = 16.0, - regularization_update_step_size: float | None = 0.8, - regularization_poly_order: int = 1, - max_image_shift: float | None = 32.0, - adam_steps: int = 30, - lr: float | None = None, - lbfgs_max_iter: int = 20, - max_optimize_iterations: int = 10, - regularization_max_image_shift_px: float | None = None, - solve_individual_rows: bool = True, - show_merged: bool = True, - show_images: bool = False, - show_knots: bool = True, - **kwargs, - ): - """ - Non-rigid drift correction using PyTorch (default) or SciPy backend. - - Parameters - ---------- - backend : str, default "pytorch" - Optimization backend. - - "pytorch": GPU-accelerated batched optimization. Single-knot only. - - "scipy": CPU L-BFGS row-by-row. Use when you need multi-knot - mode (``number_knots > 1``), which the pytorch path does not - yet support. - optimizer_name : str, default "adam" - PyTorch optimizer (ignored if backend="scipy"). - - **"adam"** - first-order momentum optimizer. Default. Best when: - - You want the fastest possible runtime, especially at image - sizes ≤512 px where the per-step grid_sample is small and - Adam's tight inner loop wins on launch overhead. - - You're confident ``max_image_shift`` reflects the true drift - bound (Adam's auto-lr derives from it; if it's too small, - Adam silently under-converges). - - You want bit-reproducible results across runs (LBFGS line - search has subtle non-determinism from Wolfe condition checks). - - **Provisional override guidance** (validated on one real-data - pair - Bob's gold-nanoparticle HAADF on a spectra background - - and the synthetic chevron test; needs broader testing on - diverse datasets before being treated as authoritative). If - you override ``lr`` manually, the rough formula is - ``expected_drift_px / (num_iterations * adam_steps)``. - Indicative starting values for the default ``num_iterations=8, - adam_steps=30`` (240 total steps): - * ~5 px drift (synthetic chevron, small drift): ``lr≈0.02`` - * ~50-100 px drift (gold-nanoparticle HAADF, real STEM): ``lr≈0.5`` - * larger / unknown drift: prefer ``optimizer_name="lbfgs"`` - which auto-scales via line search and doesn't need this - per-dataset tuning. - - **"lbfgs"** - quasi-Newton optimizer with strong-Wolfe line search. - Best when: - - The image is ≥512 px and you don't mind paying Python closure - overhead for fewer total steps (typically 2-3× faster than - Adam at 2048+ px because it converges in ~30 steps not 240). - - You're unsure about the drift magnitude or don't want to think - about ``lr`` tuning - LBFGS line search auto-scales the step - without any hand-tuning. - - You want quality over speed. - - **Failure modes to avoid:** - - **Don't normalize inputs to [0, 1] when using LBFGS** - - strong-Wolfe's curvature condition needs absolute gradient - magnitude above a threshold; with normalized intensities the - gradient is ~1e-4 and the line search returns step=0, - producing zero correction silently. Adam is unaffected. - - **Don't set ``max_image_shift`` smaller than your actual drift - if using Adam with default ``lr=None``** - Adam's auto-derived - lr scales with max_image_shift, so a too-small bound silently - clamps how much drift Adam can recover. LBFGS is unaffected. - - If unsure, start with Adam (the default) for ≤1024 px images and - switch to LBFGS for ≥2048 px or for unknown-drift exploratory work. - - Shared Parameters - ----------------- - num_iterations : int, default 8 - Number of outer iterations for alternating optimization. - regularization_sigma_px : float, default 16.0 - Gaussian smoothing sigma for knot regularization. - regularization_update_step_size : float, default 0.8 - Step size for knot updates (0-1, lower = more conservative). - regularization_poly_order : int, default 1 - Polynomial order for trend removal in knot regularization - (used by both pytorch and scipy backends). - max_image_shift : float, default 32.0 - Maximum shift for translation alignment between iterations. - - PyTorch Parameters (ignored if backend="scipy") - ----------------------------------------------- - adam_steps : int, default 30 - Number of Adam optimization steps per outer iteration. - lr : float or None, default None - Learning rate for Adam. When None (default), auto-derived as - ``max_image_shift / (num_iterations * adam_steps * 4)``. - - **Why auto-derive?** Adam's ``m/sqrt(v)`` update self-normalizes - the gradient, so each step moves a knot by ~``lr`` pixels - regardless of image intensity scale. The total movement budget - is ``lr × num_iterations × adam_steps`` and is hard-bounded: - Adam cannot find drift larger than that budget no matter how - many iterations you give it. This means ``lr`` must be matched - to the EXPECTED DRIFT MAGNITUDE IN PIXELS, not to gradient - magnitude - the same default value that works on small synthetic - drift will silently under-converge on real data with larger drift. - - The auto-derived formula reserves half the total step budget for - search (covering up to ``max_image_shift / 2`` of nonlinear drift) - and the other half for refinement near the minimum. - - Override with an explicit float when you know the actual drift - magnitude - e.g. ``lr=2.0`` for very-large-drift in-situ data, - or ``lr=0.005`` for atomic-resolution stable samples. - lbfgs_max_iter : int, default 20 - Maximum LBFGS iterations per outer iteration (line search probes - within each iter happen automatically). Only used when - optimizer="lbfgs". - - SciPy Parameters (ignored if backend="pytorch") - ----------------------------------------------- - max_optimize_iterations : int, default 10 - Maximum L-BFGS iterations per row. - regularization_max_image_shift_px : float, optional - Maximum allowed shift per iteration. - solve_individual_rows : bool, default True - If True, optimize each row independently. - - Display Parameters - ------------------ - show_merged : bool, default True - Show merged image after alignment. - show_images : bool, default False - Show individual aligned images. - show_knots : bool, default True - Overlay knot positions on visualizations. - - Notes - ----- - With backend="pytorch", ``self.images_warped`` is left STALE after - the loop and refreshed lazily on first access via plot methods or - ``calculate_error()``. Code that reads ``self.images_warped.array`` - directly should call ``self._ensure_warped_images()`` first, or - use ``generate_corrected_image()`` which builds its own warps from - ``self.knots``. - """ - if not hasattr(self, "knots"): - raise RuntimeError( - "No knots found. Call .preprocess() before running alignment." - ) - if backend == "pytorch": - device = self._device - dtype = self._dtype - num_images = self.shape[0] - canvas_shape = (self.shape[1], self.shape[2]) - if any(self.knots[i].shape[2] != 1 for i in range(num_images)): - raise NotImplementedError( - "PyTorch backend only supports single knot. " - "Use backend='scipy' for multiple knots.") - knots_batch = torch.tensor( - np.stack([self.knots[i][:, :, 0] for i in range(num_images)]), - dtype=dtype, device=device, requires_grad=True) - num_rows_knot = knots_batch.shape[2] - target_batch = torch.stack(self.images_t) - # Build u tensors once and reuse - same scan-position vector projects - # onto row and col offsets via the per-image scan_fast components. - u_t = [ - torch.as_tensor(self.interpolator[i].u, dtype=dtype, device=device) - for i in range(num_images) - ] - row_scan_offsets = torch.stack([ - u_t[i] * (self.interpolator[i].scan_fast[0] * (self.images[i].shape[0] - 1)) - for i in range(num_images) - ]) - col_scan_offsets = torch.stack([ - u_t[i] * (self.interpolator[i].scan_fast[1] * (self.images[i].shape[1] - 1)) - for i in range(num_images) - ]) - row_scale = 2.0 / (canvas_shape[0] - 1) - col_scale = 2.0 / (canvas_shape[1] - 1) - if optimizer_name == "adam": - # Auto-derive lr so the total movement budget covers a quarter - # of max_image_shift. The safety factor of 4 (not 2) prevents - # over-shooting at small image sizes where actual drift is well - # below max_image_shift; at large sizes the same factor still - # converges because the loss surface is smoother. See the `lr` - # parameter docstring for the full rationale. - adam_lr = lr if lr is not None else max_image_shift / (num_iterations * adam_steps * 4) - optimizer = torch.optim.Adam([knots_batch], lr=adam_lr, fused=True) - elif optimizer_name == "lbfgs": - optimizer = torch.optim.LBFGS( - [knots_batch], lr=1.0, max_iter=lbfgs_max_iter, - line_search_fn="strong_wolfe") - else: - raise ValueError(f"optimizer_name must be 'adam' or 'lbfgs', got {optimizer_name!r}") - if regularization_sigma_px is not None and regularization_sigma_px > 0: - x_knot = torch.arange(num_rows_knot, dtype=dtype, device=device) - x_norm = (x_knot - x_knot.mean()) / x_knot.std() - vander = torch.stack([x_norm ** p for p in range(regularization_poly_order + 1)], dim=1) - warped_t = self._warp_and_translate_torch( - max_image_shift, upsample_factor=8, knots_batch=knots_batch) - error_buffer = [] - for _ in tqdm(range(num_iterations), desc=f"Solving nonrigid drift ({optimizer_name})"): - # Build the reference under no_grad: arithmetic on warped_t (an - # inference tensor) would otherwise return an autograd-tracked - # leaf, and the optimizer would build a graph through it. - with torch.no_grad(): - warped_sum = warped_t.sum(0) - ref_batch = (warped_sum[None] - warped_t) / (num_images - 1) - knots_prev = knots_batch.detach().clone() - # Regularization alters the loss surface between outer iters, so - # stale momentum / curvature history would push knots the wrong way. - optimizer.state.clear() - if optimizer_name == "adam": - self._optimize_knots_adam( - ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, adam_steps) - else: - self._optimize_knots_lbfgs( - ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer) - self._regularize_knots( - knots_batch, knots_prev, vander, - regularization_max_image_shift_px, - regularization_sigma_px, - regularization_update_step_size) - warped_t = self._warp_and_translate_torch( - max_image_shift, upsample_factor=8, knots_batch=knots_batch) - # Per-iter error stays on GPU; sync once after the loop - images_mean = warped_t.mean(dim=0) - error_buffer.append(torch.mean(torch.abs(warped_t - images_mean[None]), dim=(1, 2))) - # Sync knots back to numpy; leave images_warped lazy so callers - # that never plot avoid the GPU→CPU transfer of the warped stack. - knots_final = knots_batch.detach().cpu().numpy() - for img_idx in range(num_images): - self.knots[img_idx][:, :, 0] = knots_final[img_idx] - self._images_warped_stale = True - self._max_image_shift_cached = max_image_shift - if error_buffer: - # Build all error rows in one DtoH transfer + one vstack, instead of - # the quadratic vstack-per-iteration pattern used by calculate_error. - errors_np = torch.stack(error_buffer).cpu().numpy() # (num_iterations, num_images) - mode_col = np.full((len(errors_np), 1), 2.0) - mean_col = errors_np.mean(axis=1, keepdims=True) - new_rows = np.hstack((mode_col, mean_col, errors_np)) - if not hasattr(self, "error_track"): - self.error_track = new_rows - else: - self.error_track = np.vstack((self.error_track, new_rows)) - else: - for _ in tqdm(range(num_iterations), desc="Solving nonrigid drift (scipy)"): - for ind in range(self.shape[0]): - image_ref = np.delete(self.images_warped.array, ind, axis=0).mean(axis=0) - knots_updated = self._optimize_knots_scipy( - ind, image_ref, self.knots[ind], - max_optimize_iterations=max_optimize_iterations, - solve_individual_rows=solve_individual_rows) - if regularization_max_image_shift_px is not None: - knots_shift = knots_updated - self.knots[ind] - knots_dist = np.sqrt(np.sum(knots_shift**2, axis=0)) - sub = knots_dist > regularization_max_image_shift_px - knots_updated[0][sub] = (self.knots[ind][0][sub] - + knots_shift[0][sub] * regularization_max_image_shift_px / knots_dist[sub]) - knots_updated[1][sub] = (self.knots[ind][1][sub] - + knots_shift[1][sub] * regularization_max_image_shift_px / knots_dist[sub]) - if regularization_sigma_px is not None and regularization_sigma_px > 0: - knots_smoothed = knots_updated.copy() - for dim in range(2): - x = np.arange(knots_updated.shape[1]) - for knot_ind in range(knots_updated.shape[2]): - y = knots_updated[dim, :, knot_ind] - coefs = np.polyfit(x, y, deg=regularization_poly_order) - trend = np.polyval(coefs, x) - residual = y - trend - residual_smooth = gaussian_filter(residual, sigma=regularization_sigma_px) - knots_smoothed[dim, :, knot_ind] = residual_smooth + trend - knots_updated = knots_smoothed - if regularization_update_step_size is not None: - knots_updated = (self.knots[ind] - + (knots_updated - self.knots[ind]) * regularization_update_step_size) - self.knots[ind] = knots_updated - warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor=8) - self.calculate_error(2, _warped_t=warped_t) - - if show_merged: - self.plot_merged_images( - show_knots=show_knots, - title="Merged: non-rigid", - **kwargs, - ) - - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: non-rigid" for i in range(self.shape[0])], - **kwargs, - ) - - return self - - def _optimize_knots_adam( - self, ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, adam_steps, - ): - """Run ``adam_steps`` of Adam on a batched knot tensor against ``_compiled_loss_fn``.""" - ref_t = ref_batch[:, None] - for _ in range(adam_steps): - optimizer.zero_grad() - loss = self._compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale) - loss.backward() - optimizer.step() - - @staticmethod - @torch.compile(mode="reduce-overhead", dynamic=False) - def _compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - ): - """Fused forward pass: knot offsets → grid → grid_sample → MSE loss. - - The MSE is averaged over both the batch (N images) and the spatial - dims, so each image's gradient is scaled by 1/N relative to a - per-image solve. Adam's adaptive step size absorbs the constant - rescale; LBFGS line search rescales itself. - """ - grid_row = (knots_batch[:, 0, :, None] + row_scan_offsets[:, None, :]) * row_scale - 1.0 - grid_col = (knots_batch[:, 1, :, None] + col_scan_offsets[:, None, :]) * col_scale - 1.0 - grid = torch.stack([grid_col, grid_row], dim=-1) - warped = torch.nn.functional.grid_sample( - ref_t, grid, mode='bilinear', align_corners=True, padding_mode='border')[:, 0] - return ((warped - target_batch) ** 2).mean() - - def _optimize_knots_lbfgs( - self, ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, - ): - """Run one LBFGS outer step (line search re-evaluates the closure several times).""" - ref_t = ref_batch[:, None] - def closure(): - optimizer.zero_grad() - loss = self._compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale) - loss.backward() - return loss - optimizer.step(closure) - - def _regularize_knots( - self, knots_batch, knots_prev, vander, - max_shift_px, sigma_px, step_size, - ): - """Apply per-iteration knot regularization (in-place on ``knots_batch``). - - Three independent stages, each gated by its parameter being non-None: - 1. Per-knot shift cap: clamp ``|new - prev|`` to ``max_shift_px`` - so the optimizer can't move any knot too far in one outer iter. - 2. Polynomial detrend + Gaussian smooth: keep low-order trends, - smooth the residual along the scan-line dimension. Removes - high-frequency optimizer wobble while preserving the drift signal. - 3. Step-size blend: ``new = prev + step_size · (new - prev)``, - under-relaxes the update for stability across outer iterations. - """ - num_images, _, num_rows_knot = knots_batch.shape - with torch.no_grad(): - if max_shift_px is not None: - shift = knots_batch - knots_prev - dist = torch.norm(shift, dim=1, keepdim=True) - scale_factor = torch.clamp(max_shift_px / dist.clamp(min=1e-8), max=1.0) - knots_batch.copy_(knots_prev + shift * scale_factor) - if sigma_px is not None and sigma_px > 0: - # Detrend + smooth all (N*2, num_rows) knots in one batched lstsq + smooth - knots_flat = knots_batch.reshape(-1, num_rows_knot).T # (num_rows, N*2) - coefs, _, _, _ = torch.linalg.lstsq(vander, knots_flat) - trend = (vander @ coefs).T # (N*2, num_rows) - residual = knots_batch.reshape(-1, num_rows_knot) - trend - smoothed = gaussian_smooth_1d(residual, sigma_px) - knots_batch.copy_((smoothed + trend).reshape(num_images, 2, num_rows_knot)) - if step_size is not None: - knots_batch.copy_(knots_prev + (knots_batch - knots_prev) * step_size) - - def _optimize_knots_scipy( - self, idx: int, image_ref: np.ndarray, knots_init: np.ndarray, - max_optimize_iterations: int = 10, solve_individual_rows: bool = True, - ) -> np.ndarray: - """SciPy L-BFGS optimization for one image.""" - shape_knots = knots_init.shape - options = {"maxiter": max_optimize_iterations} if max_optimize_iterations else {} - if solve_individual_rows: - knots_updated = np.zeros_like(knots_init) - for row_ind in range(knots_init.shape[1]): - x0 = knots_init[:, row_ind, :].ravel() - def cost_function(x): - knots_row = x.reshape(shape_knots[0], shape_knots[2]) - xa, ya = self.interpolator[idx].transform_rows(knots_row) - xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) - yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) - dx, dy = xa - xf, ya - yf - warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) - + image_ref[xf + 1, yf] * dx * (1 - dy) - + image_ref[xf, yf + 1] * (1 - dx) * dy - + image_ref[xf + 1, yf + 1] * dx * dy) - return np.sum((warped - self.images[idx].array[row_ind, :]) ** 2) - result = minimize(cost_function, x0, method="L-BFGS-B", options=options) - knots_updated[:, row_ind, :] = result.x.reshape((2, -1)) - else: - x0 = knots_init.ravel() - def cost_function(x): - knots = x.reshape(shape_knots) - xa, ya = self.interpolator[idx].transform_coordinates(knots) - xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) - yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) - dx, dy = xa - xf, ya - yf - warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) - + image_ref[xf + 1, yf] * dx * (1 - dy) - + image_ref[xf, yf + 1] * (1 - dx) * dy - + image_ref[xf + 1, yf + 1] * dx * dy) - return np.sum((warped - self.images[idx].array) ** 2) - result = minimize(cost_function, x0, method="L-BFGS-B", options=options) - knots_updated = result.x.reshape(shape_knots) - return knots_updated - - @torch.inference_mode() - def generate_corrected( - self, - upsample_factor: int = 2, - output_original_shape: bool = True, - strip_padding: bool = False, - mask_output: bool = True, - mask_edge_blend: float = 8.0, - fourier_filter: bool = True, - filter_midpoint: float = 0.5, - kde_sigma: float | None = 0.5, - weight_thresh: float = 0.1, - show_merged: bool = True, - **kwargs, - ): - """Generate the final drift-corrected image on GPU using torch. - - Parameters - ---------- - upsample_factor : int, default 2 - Factor to upsample the output image for enhanced interpolation accuracy. - output_original_shape : bool, default True - If True, crop the output image back to the original input dimensions. - strip_padding : bool, default False - If True (and output_original_shape is True), also strip the scan padding - to return only the original scan-area pixels. - mask_output : bool, default True - If True, mask the output using the probe position weights. - mask_edge_blend : float, default 8.0 - Pixels over which the mask edge is blended. - fourier_filter : bool, default True - Whether to apply Fourier-based directional filtering to merge corrected images. - filter_midpoint : float, default 0.5 - Midpoint for the sigmoid-based Fourier weighting filter. - kde_sigma : float or None, default 0.5 - Standard deviation for kernel density estimation. Uses object's kde_sigma if None. - weight_thresh : float, default 0.1 - Threshold for masking outputs. - show_merged : bool, default True - Whether to display the final corrected image. - **kwargs - Additional keyword arguments passed to the plotting function. - - Returns - ------- - image_corr : Dataset2d - The final drift-corrected output image. - """ - if not hasattr(self, "knots"): - raise RuntimeError( - "No knots found. Call .preprocess() before generating the corrected image." - ) - - device = self._device - dtype = self._dtype - up_h = round(self.shape[1] * upsample_factor) - up_w = round(self.shape[2] * upsample_factor) - canvas_shape = (up_h, up_w) - - if kde_sigma is None: - kde_sigma = self.kde_sigma - - stack_corr = torch.zeros(self.shape[0], up_h, up_w, dtype=dtype, device=device) - weight_corr = torch.zeros_like(stack_corr) - for img_idx in range(self.shape[0]): - knots_t = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_t, - self.scan_fast_t[img_idx], - self.images[img_idx].shape, - ) - warped, weights = bilinear_kde_batch( - row_t[None] * upsample_factor, - col_t[None] * upsample_factor, - self.images_t[img_idx], - canvas_shape, - kde_sigma * upsample_factor, - self.pad_value[img_idx], - ) - stack_corr[img_idx] = warped[0] - weight_corr[img_idx] = weights[0] - - if fourier_filter: - freq_row = torch.fft.fftfreq(up_h, dtype=dtype, device=device)[:, None] - freq_col = torch.fft.fftfreq(up_w, dtype=dtype, device=device)[None, :] - freq_angle = torch.atan2(freq_col, freq_row) - stack_fft = torch.fft.fft2(stack_corr) - weights = torch.zeros_like(stack_corr) - for img_idx in range(self.shape[0]): - weights[img_idx] = torch.abs( - torch.remainder( - (freq_angle - self.scan_direction[img_idx]) / np.pi + 0.5, - 1.0, - ) - 0.5 - ) / 0.5 - weights[img_idx, 0, 0] = 1.0 - weights[img_idx] = _bounded_sine_sigmoid_torch( - weights[img_idx], - midpoint=filter_midpoint, - ) - stack_fft[img_idx] *= weights[img_idx] - weights_sum = weights.sum(0) - fft_sum = stack_fft.sum(0) - image_corr_fft = torch.where( - weights_sum > 0.0, - fft_sum / weights_sum.clamp(min=1e-8), - torch.zeros_like(fft_sum), - ) - else: - image_corr_fft = torch.fft.fft2(stack_corr.mean(0)) - - if mask_output: - weight_np = weight_corr.cpu().numpy() - mask_edge = np.prod(weight_np >= (weight_thresh / upsample_factor**2), axis=0) - mask_edge[:, 0] = False - mask_edge[:, -1] = False - mask_edge[0, :] = False - mask_edge[-1, :] = False - mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend - mask_np = ( - np.cos( - (np.pi / 2) - * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) - ) - ** 2 - ) - mask_t = torch.as_tensor(mask_np, dtype=dtype, device=device) - pad_value_mean = float(np.mean(self.pad_value)) - image_corr_fft = torch.fft.fft2( - torch.fft.ifft2(image_corr_fft).real * mask_t + pad_value_mean * (1 - mask_t) - ) - - if output_original_shape: - image_corr_fft = _fourier_crop_torch(image_corr_fft, self.shape[-2:]) / upsample_factor**2 - - corr_np = torch.fft.ifft2(image_corr_fft).real.cpu().numpy() - if strip_padding and output_original_shape: - scan_h, scan_w = self.images[0].shape[:2] - canvas_h, canvas_w = corr_np.shape[:2] - pad_h = (canvas_h - scan_h) // 2 - pad_w = (canvas_w - scan_w) // 2 - corr_np = corr_np[pad_h:pad_h + scan_h, pad_w:pad_w + scan_w] - - image_corr = Dataset2d.from_array( - corr_np, - name="drift corrected image", - origin=self.images[0].origin, - sampling=self.images[0].sampling, - units=self.images[0].units, - ) - if show_merged: - show_2d(image_corr.array, **kwargs) - plt.show() - return image_corr - - def generate_corrected_image( - self, - upsample_factor: int = 2, - output_original_shape: bool = True, - mask_output: bool = True, - mask_edge_blend: float = 8.0, - fourier_filter: bool = True, - filter_midpoint: float = 0.5, - kde_sigma: float = 0.5, - weight_thresh=0.1, - show_image: bool = True, - **kwargs, - ): - """ - Generate the final drift-corrected image after aligning a stack of input images. - - Parameters - ---------- - upsample_factor : int, default 2 - Factor to upsample the output image for enhanced interpolation accuracy. - output_original_shape : bool, default True - If True, crop the output image back to the original input dimensions after processing. - mask_output : bool, default True - If true, mask the output using the probe position weights - mask_edge_blend : float, default 8.0 - Value in pixels to blend from the edge of the mask (where we have data) - fourier_filter : bool, default True - Whether to apply Fourier-based directional filtering to merge corrected images. - filter_midpoint : float, default 0.5 - Midpoint for the sigmoid-based Fourier weighting filter, determining transition smoothness. - Setting this to a low value close to 0 will include more signal but also more slow scan artifacts. - If using 2 images at 0 and 90 degrees scan angles, any value >0.75 will be unstable. - Only use larger values (close to 1.0) if multiple images covering many scan angles are used. - kde_sigma : float, default 0.5 - Standard deviation for kernel density estimation used during image interpolation. Defaults - to the object's stored kde_sigma if set to None. - weight_thresh: float, default 0.1 - This value sets the threshold for masking the outputs. - For very large jitter artifacts this value can be lowered. - show_image : bool, default True - Whether to display the final corrected image after processing. - **kwargs : dict - Additional keyword arguments passed to the plotting function when displaying the image. - - Returns - ------- - image_corr : Dataset2d - The final drift-corrected output image encapsulated in a Dataset2d object. - - Notes - ----- - - The function applies per-frame warping using knot-based interpolation and optionally - performs directional Fourier filtering to blend multiple warped images. - - The Fourier filter suppresses directional artifacts by weighting image contributions based - on their scan angles, utilizing a bounded sine sigmoid for smooth transition. - - Upsampling enhances interpolation precision but may increase computational cost. - """ - - # init - stack_corr = np.zeros( - ( - self.shape[0], - np.round(self.shape[1] * upsample_factor).astype("int"), - np.round(self.shape[2] * upsample_factor).astype("int"), - ) - ) - weight_corr = np.zeros( - ( - self.shape[0], - np.round(self.shape[1] * upsample_factor).astype("int"), - np.round(self.shape[2] * upsample_factor).astype("int"), - ) - ) - - if kde_sigma is None: - kde_sigma = self.kde_sigma - - # Update images - for ind in range(self.shape[0]): - stack_corr[ind], weight_corr[ind] = self.interpolator[ind].warp_image( - self.images[ind].array, - self.knots[ind], - kde_sigma=kde_sigma, - upsample_factor=upsample_factor, - ) - - if fourier_filter: - # Apply fourier filtering - kx = np.fft.fftfreq(stack_corr.shape[1])[:, None] - ky = np.fft.fftfreq(stack_corr.shape[2])[None, :] - kt = np.arctan2(ky, kx) - - stack_fft = np.fft.fft2(stack_corr) - weights = np.zeros_like(stack_corr) - - for ind in range(stack_corr.shape[0]): - # Calculate weights as a function of angle - weights[ind] = np.abs( - np.mod((kt - self.scan_direction[ind]) / np.pi + 0.5, 1.0) - 0.5 - ) / (1 / 2) - weights[ind][0, 0] = 1.0 - - # Apply sigmoid to weighting function - weights[ind] = bounded_sine_sigmoid( - weights[ind], - midpoint=filter_midpoint, - ) - - # Weight the fourier transformed images - stack_fft[ind] *= weights[ind] - - weights_sum = np.sum(weights, axis=0) - image_corr_fft = np.zeros_like(weights_sum, dtype=complex) - np.divide( - np.sum(stack_fft, axis=0), - weights_sum, - where=weights_sum > 0.0, - out=image_corr_fft, - ) - - else: - image_corr_fft = np.fft.fft2(np.mean(stack_corr, axis=0)) - - if mask_output: - # Note that we compute 2 boolean masks to round off the corners of image blending - - # calculate mask from product of individual image masks - # scale weights by upsample factor to normalize to mean value of 1.0 - mask_edge = np.prod(weight_corr >= (weight_thresh / upsample_factor**2), axis=0) - # Set outermost pixels to False to define the boundary for edge blending - mask_edge[:, 0] = False - mask_edge[:, -1] = False - mask_edge[0, :] = False - mask_edge[-1, :] = False - # Find inner boundary mask - mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend - # compute mask using edge blending value - mask = ( - np.cos( - (np.pi / 2) - * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) - ) - ** 2 - ) - # Mean pad value - pad_value_mean = np.mean([ind.pad_value for ind in self.interpolator]) - # apply mask - image_corr_fft = np.fft.fft2( - np.fft.ifft2(image_corr_fft) * mask + pad_value_mean * (1 - mask) - ) - - if output_original_shape: - image_corr_fft = fourier_cropping(image_corr_fft, self.shape[-2:]) / upsample_factor**2 - - # TODO - adjust origin / sampling if output sampling is different from input - # i.e. if output_original_shape is False, and upsample_factor > 1 - image_corr = Dataset2d.from_array( - np.real(np.fft.ifft2(image_corr_fft)), - name="drift corrected image", - origin=self.images[0].origin, - sampling=self.images[0].sampling, - units=self.images[0].units, - ) - - if show_image: - fig, ax = show_2d(image_corr.array, **kwargs) - # Force a render whether we're drawing into a provided Axes or a fresh Figure - ax_to_draw = kwargs.get("ax", ax) - try: - ax_to_draw.figure.canvas.draw_idle() - # If we're not drawing into a caller-provided Axes, also pop the window - if "ax" not in kwargs: - plt.show() - except Exception: - # Fallback: if backend is odd, try a blocking show - plt.show() - return image_corr - - def calculate_error( - self, - mode: int, - _warped_t: torch.Tensor | None = None, - ): - """Compute per-image MAE against the mean and append to error history. - - Measures how well the warped images agree by computing the mean - absolute difference of each image from the stack mean. Without - error tracking, there is no way to verify that alignment steps - are actually improving the result. - - Parameters - ---------- - mode : int - Stage identifier (0=preprocess, 1=affine, 2=nonrigid). - _warped_t : torch.Tensor or None - If provided, compute error from this tensor directly, - avoiding a GPU-to-CPU round-trip. - """ - if _warped_t is not None: - images_mean = _warped_t.mean(dim=0) - sig_diff = torch.mean( - torch.abs(_warped_t - images_mean[None]), dim=(1, 2) - ).cpu().numpy() - else: - # Lazy refresh: align_nonrigid defers the warped→numpy sync until - # someone reads it, so calculate_error must trigger the refresh. - self._ensure_warped_images() - images_mean = np.mean(self.images_warped.array, axis=0) - sig_diff = np.mean( - np.abs(self.images_warped.array - images_mean[None, :, :]), axis=(1, 2) - ) - - # Error vector - error_current = np.hstack((mode, np.mean(sig_diff), sig_diff)) - - # Initialize or append to error tracking array - if not hasattr(self, "error_track"): - self.error_track = error_current[None, :] # initialize with first row - else: - self.error_track = np.vstack((self.error_track, error_current)) - - def plot_transformed_images(self, show_knots: bool = True, **kwargs): - self._ensure_warped_images() - fig, ax = show_2d( - list(self.images_warped.array), - **kwargs, - ) - if show_knots: - for a0 in range(self.shape[0]): - x = self.knots[a0][0] - y = self.knots[a0][1] - ax[a0].plot( - y, - x, - color="r", - ) - - def plot_convergence( - self, - figsize=(8, 3), - **kwargs, - ): - """ - Plot the convergence of the drift correction. - """ - sub = np.abs(self.error_track[:, 0] - 2) < 0.1 - error = self.error_track[:, 1] - it = np.arange(error.shape[0]) - - from matplotlib.ticker import FormatStrFormatter, MaxNLocator - - fig, ax = plt.subplots(1, 2, figsize=figsize) - color = (1, 0, 0) # red - - # Plot Affine - if np.any(~sub): - ax[0].plot( - it[~sub], - 100 * error[~sub], - marker="o", - color=color, - linestyle="-", - label="Affine", - **kwargs, - ) - ax[0].set_xlabel("Affine Iterations") - ax[0].set_ylabel("Mean Error [%]") - ax[0].xaxis.set_major_locator(MaxNLocator(integer=True)) - ax[0].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) - else: - ax[0].axis("off") - - # Plot Non-Rigid - if np.any(sub): - first_true = np.argmax(sub) - if first_true > 0: - sub[first_true - 1] = True - - ax[1].plot( - it[sub], - 100 * error[sub], - marker="o", - color=color, - linestyle="-", - label="Non-Rigid", - **kwargs, - ) - ax[1].set_xlabel("Non-Rigid Iterations") - ax[1].xaxis.set_major_locator(MaxNLocator(integer=True)) - ax[1].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) - else: - ax[1].axis("off") - - plt.tight_layout() - - return self - - def _ensure_warped_images(self): - """Lazily populate images_warped from current knots if marked stale.""" - if getattr(self, "_images_warped_stale", False): - self._warp_and_translate_torch( - self._max_image_shift_cached, upsample_factor=8, - solve_translation=False) - self._images_warped_stale = False - - def plot_merged_images(self, show_knots: bool = True, **kwargs): - """ - Plot the current transformed images, with knot overlays. - """ - self._ensure_warped_images() - fig, ax = show_2d( - self.images_warped.array.mean(0), - **kwargs, - ) - if show_knots: - for a0 in range(self.shape[0]): - x = self.knots[a0][0] - y = self.knots[a0][1] - ax.plot( - y, - x, - ) - - -class DriftInterpolator: - def __init__( - self, - input_shape, - output_shape, - scan_fast, - scan_slow, - pad_value, - kde_sigma, - ): - self.input_shape = input_shape - self.output_shape = output_shape - self.scan_fast = scan_fast - self.scan_slow = scan_slow - self.pad_value = pad_value - self.kde_sigma = kde_sigma - - self.rows_input = np.arange(input_shape[0]) - self.cols_input = np.arange(input_shape[1]) - self.u = np.linspace(0, 1, input_shape[1]) - - def transform_rows( - self, - knots_row: NDArray, - ): - num_knots = knots_row.shape[-1] - basis = np.linspace(0, 1, num_knots) - - if num_knots == 1: - xa = knots_row[0] + self.u[None, :] * self.scan_fast[0] * (self.input_shape[0] - 1) - ya = knots_row[1] + self.u[None, :] * self.scan_fast[1] * (self.input_shape[1] - 1) - elif num_knots == 2: - xa = interp1d(basis, knots_row[0], kind="linear", assume_sorted=True)(self.u) - ya = interp1d(basis, knots_row[1], kind="linear", assume_sorted=True)(self.u) - else: - kind = "quadratic" if num_knots == 3 else "cubic" - xa = interp1d( - basis, - knots_row[0], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - ya = interp1d( - basis, - knots_row[1], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - - return xa, ya - - def transform_coordinates( - self, - knots: NDArray, - ): - num_knots = knots.shape[-1] - - if num_knots == 1: - # vectorized version for speed - xa, ya = self.transform_rows(knots) - else: - xa = np.zeros(self.input_shape) - ya = np.zeros(self.input_shape) - for i in range(self.input_shape[0]): - xa[i], ya[i] = self.transform_rows(knots[:, i]) - - return xa, ya - - def warp_image( - self, - image: NDArray, - knots: NDArray, # shape: (2, rows, num_knots) - kde_sigma=None, - output_shape=None, - pad_value=None, - upsample_factor=None, - ) -> NDArray: - xa, ya = self.transform_coordinates( - knots, - ) - - if kde_sigma is None: - kde_sigma = self.kde_sigma - - if output_shape is None: - output_shape = self.output_shape - - if pad_value is None: - pad_value = self.pad_value - - if upsample_factor is None: - upsample_factor = 1.0 - - image_interp, weight_interp = bilinear_kde( - xa=xa * upsample_factor, # rows - ya=ya * upsample_factor, # cols - values=image, - output_shape=np.round(np.array(output_shape) * upsample_factor).astype("int"), - kde_sigma=kde_sigma * upsample_factor, - pad_value=pad_value, - return_pix_count=True, - ) - - return image_interp, weight_interp - - -def bounded_sine_sigmoid(x, midpoint=0.5, width=1.0): - """ - Piecewise bounded sigmoid: zero, raised sine squared, one. - - Parameters - ---------- - x : array-like, shape (...,) - Input values in [0, 1]. - midpoint : float - Center of the sigmoid transition. - width : float - Width of the sigmoid (range over which it ramps from 0 to 1). - Returns - ------- - y : array-like - Output in [0, 1], same shape as x. - """ - x = np.asarray(x) - # Truncate width if midpoint too close to edge - left_max = midpoint - width / 2 - right_min = midpoint + width / 2 - if left_max < 0: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * midpoint}.", - RuntimeWarning, - ) - width = 2 * midpoint - - if right_min > 1: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * (1 - midpoint)}.", - RuntimeWarning, - ) - width = 2 * (1 - midpoint) - # Recalculate edges - left = midpoint - width / 2 - right = midpoint + width / 2 - - y = np.zeros_like(x, dtype=float) - in_band = (x >= left) & (x <= right) - # Map [left, right] to [0, pi/2] - t = (x[in_band] - left) / width # goes from 0 to 1 - y[in_band] = np.sin(t * np.pi / 2) ** 2 - y[x > right] = 1.0 - return y - - -def _bounded_sine_sigmoid_torch( - x: torch.Tensor, - midpoint: float = 0.5, - width: float = 1.0, -) -> torch.Tensor: - width = min(width, 2 * midpoint, 2 * (1 - midpoint)) - left = midpoint - width / 2 - right = midpoint + width / 2 - t = ((x - left) / width).clamp(0.0, 1.0) - return torch.where(x > right, torch.ones_like(x), torch.sin(t * (np.pi / 2)) ** 2) - -def _fourier_crop_torch( - fft_array: torch.Tensor, - crop_shape: tuple[int, int], -) -> torch.Tensor: - crop_h, crop_w = crop_shape - h1 = crop_h // 2 - h2 = crop_h - h1 - w1 = crop_w // 2 - w2 = crop_w - w1 - result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) - result[:h1, :w1] = fft_array[:h1, :w1] - result[:h1, -w2:] = fft_array[:h1, -w2:] - result[-h2:, :w1] = fft_array[-h2:, :w1] - result[-h2:, -w2:] = fft_array[-h2:, -w2:] - return result +DriftCorrection.__module__ = __name__ +DriftInterpolator.__module__ = __name__ diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py new file mode 100644 index 00000000..921223ee --- /dev/null +++ b/src/quantem/imaging/drift/apply.py @@ -0,0 +1,399 @@ + +import matplotlib.pyplot as plt +import numpy as np +import torch +from scipy.ndimage import distance_transform_edt + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.utils.imaging_utils import ( + fourier_cropping, +) +from quantem.core.visualization import show_2d +from quantem.imaging.drift.core.warping import ( + _bounded_sine_sigmoid_torch, + _fourier_crop_torch, + bounded_sine_sigmoid, +) +from quantem.imaging.drift_utils import ( + bilinear_kde_batch, + transform_coordinates_single_knot, +) + + +@torch.inference_mode() +def generate_corrected( + self, + upsample_factor: int = 2, + output_original_shape: bool = True, + strip_padding: bool = False, + mask_output: bool = True, + mask_edge_blend: float = 8.0, + fourier_filter: bool = True, + filter_midpoint: float = 0.5, + kde_sigma: float | None = 0.5, + weight_thresh: float = 0.1, + show_merged: bool = True, + **kwargs, +): + """Generate the final drift-corrected image on GPU using torch. + + Parameters + ---------- + upsample_factor : int, default 2 + Factor to upsample the output image for enhanced interpolation accuracy. + output_original_shape : bool, default True + If True, crop the output image back to the original input dimensions. + strip_padding : bool, default False + If True (and output_original_shape is True), also strip the scan padding + to return only the original scan-area pixels. + mask_output : bool, default True + If True, mask the output using the probe position weights. + mask_edge_blend : float, default 8.0 + Pixels over which the mask edge is blended. + fourier_filter : bool, default True + Whether to apply Fourier-based directional filtering to merge corrected images. + filter_midpoint : float, default 0.5 + Midpoint for the sigmoid-based Fourier weighting filter. + kde_sigma : float or None, default 0.5 + Standard deviation for kernel density estimation. Uses object's kde_sigma if None. + weight_thresh : float, default 0.1 + Threshold for masking outputs. + show_merged : bool, default True + Whether to display the final corrected image. + **kwargs + Additional keyword arguments passed to the plotting function. + + Returns + ------- + image_corr : Dataset2d + The final drift-corrected output image. + """ + if not hasattr(self, "knots"): + raise RuntimeError( + "No knots found. Call .preprocess() before generating the corrected image." + ) + + device = self._device + dtype = self._dtype + up_h = round(self.shape[1] * upsample_factor) + up_w = round(self.shape[2] * upsample_factor) + canvas_shape = (up_h, up_w) + + if kde_sigma is None: + kde_sigma = self.kde_sigma + + stack_corr = torch.zeros(self.shape[0], up_h, up_w, dtype=dtype, device=device) + weight_corr = torch.zeros_like(stack_corr) + for img_idx in range(self.shape[0]): + knots_t = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) + row_t, col_t = transform_coordinates_single_knot( + knots_t, + self.scan_fast_t[img_idx], + self.images[img_idx].shape, + ) + warped, weights = bilinear_kde_batch( + row_t[None] * upsample_factor, + col_t[None] * upsample_factor, + self.images_t[img_idx], + canvas_shape, + kde_sigma * upsample_factor, + self.pad_value[img_idx], + ) + stack_corr[img_idx] = warped[0] + weight_corr[img_idx] = weights[0] + + if fourier_filter: + freq_row = torch.fft.fftfreq(up_h, dtype=dtype, device=device)[:, None] + freq_col = torch.fft.fftfreq(up_w, dtype=dtype, device=device)[None, :] + freq_angle = torch.atan2(freq_col, freq_row) + stack_fft = torch.fft.fft2(stack_corr) + weights = torch.zeros_like(stack_corr) + for img_idx in range(self.shape[0]): + weights[img_idx] = torch.abs( + torch.remainder( + (freq_angle - self.scan_direction[img_idx]) / np.pi + 0.5, + 1.0, + ) - 0.5 + ) / 0.5 + weights[img_idx, 0, 0] = 1.0 + weights[img_idx] = _bounded_sine_sigmoid_torch( + weights[img_idx], + midpoint=filter_midpoint, + ) + stack_fft[img_idx] *= weights[img_idx] + weights_sum = weights.sum(0) + fft_sum = stack_fft.sum(0) + image_corr_fft = torch.where( + weights_sum > 0.0, + fft_sum / weights_sum.clamp(min=1e-8), + torch.zeros_like(fft_sum), + ) + else: + image_corr_fft = torch.fft.fft2(stack_corr.mean(0)) + + if mask_output: + weight_np = weight_corr.cpu().numpy() + mask_edge = np.prod(weight_np >= (weight_thresh / upsample_factor**2), axis=0) + mask_edge[:, 0] = False + mask_edge[:, -1] = False + mask_edge[0, :] = False + mask_edge[-1, :] = False + mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend + mask_np = ( + np.cos( + (np.pi / 2) + * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) + ) + ** 2 + ) + mask_t = torch.as_tensor(mask_np, dtype=dtype, device=device) + pad_value_mean = float(np.mean(self.pad_value)) + image_corr_fft = torch.fft.fft2( + torch.fft.ifft2(image_corr_fft).real * mask_t + pad_value_mean * (1 - mask_t) + ) + + if output_original_shape: + image_corr_fft = _fourier_crop_torch(image_corr_fft, self.shape[-2:]) / upsample_factor**2 + + corr_np = torch.fft.ifft2(image_corr_fft).real.cpu().numpy() + if strip_padding and output_original_shape: + scan_h, scan_w = self.images[0].shape[:2] + canvas_h, canvas_w = corr_np.shape[:2] + pad_h = (canvas_h - scan_h) // 2 + pad_w = (canvas_w - scan_w) // 2 + corr_np = corr_np[pad_h:pad_h + scan_h, pad_w:pad_w + scan_w] + + image_corr = Dataset2d.from_array( + corr_np, + name="drift corrected image", + origin=self.images[0].origin, + sampling=self.images[0].sampling, + units=self.images[0].units, + ) + if show_merged: + show_2d(image_corr.array, **kwargs) + plt.show() + return image_corr + + +def generate_corrected_image( + self, + upsample_factor: int = 2, + output_original_shape: bool = True, + mask_output: bool = True, + mask_edge_blend: float = 8.0, + fourier_filter: bool = True, + filter_midpoint: float = 0.5, + kde_sigma: float = 0.5, + weight_thresh=0.1, + show_image: bool = True, + **kwargs, +): + """ + Generate the final drift-corrected image after aligning a stack of input images. + + Parameters + ---------- + upsample_factor : int, default 2 + Factor to upsample the output image for enhanced interpolation accuracy. + output_original_shape : bool, default True + If True, crop the output image back to the original input dimensions after processing. + mask_output : bool, default True + If true, mask the output using the probe position weights + mask_edge_blend : float, default 8.0 + Value in pixels to blend from the edge of the mask (where we have data) + fourier_filter : bool, default True + Whether to apply Fourier-based directional filtering to merge corrected images. + filter_midpoint : float, default 0.5 + Midpoint for the sigmoid-based Fourier weighting filter, determining transition smoothness. + Setting this to a low value close to 0 will include more signal but also more slow scan artifacts. + If using 2 images at 0 and 90 degrees scan angles, any value >0.75 will be unstable. + Only use larger values (close to 1.0) if multiple images covering many scan angles are used. + kde_sigma : float, default 0.5 + Standard deviation for kernel density estimation used during image interpolation. Defaults + to the object's stored kde_sigma if set to None. + weight_thresh: float, default 0.1 + This value sets the threshold for masking the outputs. + For very large jitter artifacts this value can be lowered. + show_image : bool, default True + Whether to display the final corrected image after processing. + **kwargs : dict + Additional keyword arguments passed to the plotting function when displaying the image. + + Returns + ------- + image_corr : Dataset2d + The final drift-corrected output image encapsulated in a Dataset2d object. + + Notes + ----- + - The function applies per-frame warping using knot-based interpolation and optionally + performs directional Fourier filtering to blend multiple warped images. + - The Fourier filter suppresses directional artifacts by weighting image contributions based + on their scan angles, utilizing a bounded sine sigmoid for smooth transition. + - Upsampling enhances interpolation precision but may increase computational cost. + """ + + # init + stack_corr = np.zeros( + ( + self.shape[0], + np.round(self.shape[1] * upsample_factor).astype("int"), + np.round(self.shape[2] * upsample_factor).astype("int"), + ) + ) + weight_corr = np.zeros( + ( + self.shape[0], + np.round(self.shape[1] * upsample_factor).astype("int"), + np.round(self.shape[2] * upsample_factor).astype("int"), + ) + ) + + if kde_sigma is None: + kde_sigma = self.kde_sigma + + # Update images + for ind in range(self.shape[0]): + stack_corr[ind], weight_corr[ind] = self.interpolator[ind].warp_image( + self.images[ind].array, + self.knots[ind], + kde_sigma=kde_sigma, + upsample_factor=upsample_factor, + ) + + if fourier_filter: + # Apply fourier filtering + kx = np.fft.fftfreq(stack_corr.shape[1])[:, None] + ky = np.fft.fftfreq(stack_corr.shape[2])[None, :] + kt = np.arctan2(ky, kx) + + stack_fft = np.fft.fft2(stack_corr) + weights = np.zeros_like(stack_corr) + + for ind in range(stack_corr.shape[0]): + # Calculate weights as a function of angle + weights[ind] = np.abs( + np.mod((kt - self.scan_direction[ind]) / np.pi + 0.5, 1.0) - 0.5 + ) / (1 / 2) + weights[ind][0, 0] = 1.0 + + # Apply sigmoid to weighting function + weights[ind] = bounded_sine_sigmoid( + weights[ind], + midpoint=filter_midpoint, + ) + + # Weight the fourier transformed images + stack_fft[ind] *= weights[ind] + + weights_sum = np.sum(weights, axis=0) + image_corr_fft = np.zeros_like(weights_sum, dtype=complex) + np.divide( + np.sum(stack_fft, axis=0), + weights_sum, + where=weights_sum > 0.0, + out=image_corr_fft, + ) + + else: + image_corr_fft = np.fft.fft2(np.mean(stack_corr, axis=0)) + + if mask_output: + # Note that we compute 2 boolean masks to round off the corners of image blending + + # calculate mask from product of individual image masks + # scale weights by upsample factor to normalize to mean value of 1.0 + mask_edge = np.prod(weight_corr >= (weight_thresh / upsample_factor**2), axis=0) + # Set outermost pixels to False to define the boundary for edge blending + mask_edge[:, 0] = False + mask_edge[:, -1] = False + mask_edge[0, :] = False + mask_edge[-1, :] = False + # Find inner boundary mask + mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend + # compute mask using edge blending value + mask = ( + np.cos( + (np.pi / 2) + * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) + ) + ** 2 + ) + # Mean pad value + pad_value_mean = np.mean([ind.pad_value for ind in self.interpolator]) + # apply mask + image_corr_fft = np.fft.fft2( + np.fft.ifft2(image_corr_fft) * mask + pad_value_mean * (1 - mask) + ) + + if output_original_shape: + image_corr_fft = fourier_cropping(image_corr_fft, self.shape[-2:]) / upsample_factor**2 + + # TODO - adjust origin / sampling if output sampling is different from input + # i.e. if output_original_shape is False, and upsample_factor > 1 + image_corr = Dataset2d.from_array( + np.real(np.fft.ifft2(image_corr_fft)), + name="drift corrected image", + origin=self.images[0].origin, + sampling=self.images[0].sampling, + units=self.images[0].units, + ) + + if show_image: + fig, ax = show_2d(image_corr.array, **kwargs) + # Force a render whether we're drawing into a provided Axes or a fresh Figure + ax_to_draw = kwargs.get("ax", ax) + try: + ax_to_draw.figure.canvas.draw_idle() + # If we're not drawing into a caller-provided Axes, also pop the window + if "ax" not in kwargs: + plt.show() + except Exception: + # Fallback: if backend is odd, try a blocking show + plt.show() + return image_corr + + +def calculate_error( + self, + mode: int, + _warped_t: torch.Tensor | None = None, +): + """Compute per-image MAE against the mean and append to error history. + + Measures how well the warped images agree by computing the mean + absolute difference of each image from the stack mean. Without + error tracking, there is no way to verify that alignment steps + are actually improving the result. + + Parameters + ---------- + mode : int + Stage identifier (0=preprocess, 1=affine, 2=nonrigid). + _warped_t : torch.Tensor or None + If provided, compute error from this tensor directly, + avoiding a GPU-to-CPU round-trip. + """ + if _warped_t is not None: + images_mean = _warped_t.mean(dim=0) + sig_diff = torch.mean( + torch.abs(_warped_t - images_mean[None]), dim=(1, 2) + ).cpu().numpy() + else: + # Lazy refresh: align_nonrigid defers the warped→numpy sync until + # someone reads it, so calculate_error must trigger the refresh. + self._ensure_warped_images() + images_mean = np.mean(self.images_warped.array, axis=0) + sig_diff = np.mean( + np.abs(self.images_warped.array - images_mean[None, :, :]), axis=(1, 2) + ) + + # Error vector + error_current = np.hstack((mode, np.mean(sig_diff), sig_diff)) + + # Initialize or append to error tracking array + if not hasattr(self, "error_track"): + self.error_track = error_current[None, :] # initialize with first row + else: + self.error_track = np.vstack((self.error_track, error_current)) diff --git a/src/quantem/imaging/drift/core/__init__.py b/src/quantem/imaging/drift/core/__init__.py new file mode 100644 index 00000000..b9253d00 --- /dev/null +++ b/src/quantem/imaging/drift/core/__init__.py @@ -0,0 +1 @@ +"""Numerical stages for drift correction.""" diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py new file mode 100644 index 00000000..b226246f --- /dev/null +++ b/src/quantem/imaging/drift/core/affine.py @@ -0,0 +1,378 @@ + +import numpy as np +import torch +from torch.fft import fftfreq + +from quantem.imaging.drift_utils import ( + bilinear_kde_batch, + cross_corr_batch, + transform_coordinates_single_knot, + translate_align, +) + + +def align_affine( + self, + step: float = 0.01, + num_tests: int = 9, + refine: bool = True, + upsample_factor: int = 8, + max_image_shift: float | None = 32, + chunk_size: int | None = None, + show_merged: bool = True, + show_images: bool = False, + show_knots: bool = True, + verbose: bool = False, + **kwargs, +): + """Correct affine drift between scan pairs using a batched grid search. + + Builds a grid of candidate linear-drift vectors, warps both images + for each candidate, and picks the one with the lowest cross-correlation + cost. An optional refinement pass subdivides the winning cell for + sub-step accuracy. Without affine correction, per-scanline drift + causes shear distortion that translation alignment alone cannot fix. + + Parameters + ---------- + step : float + Search resolution in pixels per scan line. The grid search + tests drift rates from ``-step * num_tests/2`` to + ``+step * num_tests/2`` px/line. For example, ``step=0.02`` + with ``num_tests=11`` searches drifts from -0.10 to +0.10 + px/line. Smaller values detect subtler drift but test more + candidates. + num_tests : int + Number of drift rates to test along each axis. Must be odd + so the grid is centered on zero drift. Total candidates + ≈ ``π/4 * num_tests²``: ``num_tests=5`` → 21, + ``num_tests=9`` → 61, ``num_tests=11`` → 97. + refine : bool + If True, run a second pass at ``step / (num_tests - 1)`` + resolution, centered on the coarse winner. + upsample_factor : int + Sub-pixel precision for measuring the translational shift + between warped image pairs. 8 means 1/8-pixel precision. + Higher values are more accurate but slower. + max_image_shift : float or None + Maximum allowed translational shift in pixels. Cross-correlation + peaks beyond this radius are masked to reject spurious matches + from noise or periodic artifacts. Set to None to allow any shift. + chunk_size : int or None + Number of candidates per pass. If None, all candidates at once. + Set to a smaller value if you run out of memory. + show_merged : bool + Display the merged (averaged) image after alignment. + show_images : bool + Display each individual warped image after alignment. + show_knots : bool + Overlay knot positions on the displayed images. + verbose : bool + If True, print the top 5 candidate drift vectors with their + cost and direction after the grid search. Useful for + diagnosing ambiguous alignments or verifying the winning + candidate has a clear margin over runner-ups. + **kwargs + Additional keyword arguments passed to the plotting functions. + + Returns + ------- + DriftCorrection + Self, for method chaining. + + Examples + -------- + >>> drift = DriftCorrection.from_data( + ... images=[im0, im1], scan_direction_degrees=[0, 90]) + >>> drift.preprocess().align_affine(step=0.02, num_tests=11) + """ + if self.shape[0] < 2: + raise ValueError( + f"align_affine requires at least 2 images (got {self.shape[0]}). " + f"Provide image pairs with different scan directions." + ) + if num_tests % 2 == 0: + raise ValueError( + f"num_tests must be odd (got {num_tests}). Try {num_tests + 1}." + ) + # Build candidate grid with circular mask (~21% fewer than square) + grid_axis = np.arange(-(num_tests - 1) / 2, (num_tests + 1) / 2) + row_grid, col_grid = np.meshgrid(grid_axis, grid_axis, indexing="ij") + circular_mask = row_grid**2 + col_grid**2 <= (num_tests / 2) ** 2 + drift_vectors = np.vstack((row_grid[circular_mask], col_grid[circular_mask])).T * step + + def _print_top_candidates(label, candidates, costs_tensor): + costs_np = costs_tensor.cpu().numpy() + ranked = np.argsort(costs_np) + best_cost = costs_np[ranked[0]] + print(f" {label} - top 5 candidates:") + for rank in range(min(5, len(ranked))): + idx = ranked[rank] + drift_row, drift_col = candidates[idx] + magnitude = np.sqrt(drift_row**2 + drift_col**2) + gap = (costs_np[idx] - best_cost) / best_cost * 100 if rank > 0 else 0 + print(f" drift=({drift_row:+.4f}, {drift_col:+.4f}) px/line " + f"({magnitude:.4f} magnitude), cost={costs_np[idx]:.4f}" + f"{f' (+{gap:.1f}%)' if rank > 0 else ' (best)'}") + + def _apply_drift(drift_vec): + for img_idx in range(self.shape[0]): + scanline_offset = np.arange(self.knots[img_idx].shape[1]) - (self.knots[img_idx].shape[1] - 1) / 2 + self.knots[img_idx][0] += drift_vec[0] * scanline_offset[:, None] + self.knots[img_idx][1] += drift_vec[1] * scanline_offset[:, None] + + def _search_and_apply(candidates, label): + best_idx, costs = self._affine_grid_search_batch(candidates, upsample_factor, max_image_shift, chunk_size) + _apply_drift(candidates[best_idx]) + if verbose: + _print_top_candidates(label, candidates, costs) + warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor) + self.calculate_error(1, _warped_t=warped_t) + return candidates[best_idx] + + drift_total = _search_and_apply(drift_vectors, "Coarse search") + if refine: + drift_fine = drift_vectors / (num_tests - 1) + drift_total = drift_total + _search_and_apply(drift_fine, "Refine search") + if verbose: + num_rows = self.images[0].shape[0] + drift_rate = np.sqrt(drift_total[0] ** 2 + drift_total[1] ** 2) + total_shift = drift_rate * num_rows + angle_deg = np.degrees(np.arctan2(drift_total[1], drift_total[0])) + print(f"align_affine: step={step}, num_tests={num_tests} " + f"({len(drift_vectors)} candidates), refine={refine}, " + f"max_image_shift={max_image_shift}") + msg = (f"Drift: ({drift_total[0]:+.4f}, {drift_total[1]:+.4f}) px/line, " + f"{drift_rate:.4f} magnitude, {angle_deg:.1f} deg, " + f"{total_shift:.1f} px total over {num_rows} lines") + if self.images[0].sampling is not None: + px_size = self.images[0].sampling[0] + unit = self.images[0].units[0] if self.images[0].units else "px" + msg += f" = {total_shift * px_size:.2f} {unit}" + print(msg) + err = self.error_track + print(f"Error: {err[0, 1]:.2f} -> {err[-1, 1]:.2f} " + f"({(err[0, 1] - err[-1, 1]) / err[0, 1] * 100:+.1f}%)") + + # Plots + kwargs.pop("title", None) + if show_merged: + self.plot_merged_images( + show_knots=show_knots, + title="Merged: affine", + **kwargs, + ) + if show_images: + self.plot_transformed_images( + show_knots=show_knots, + title=[f"Image {i}: affine" for i in range(self.shape[0])], + **kwargs, + ) + + return self + + +@torch.inference_mode() +def _affine_grid_search_batch(self, drift_vectors, upsample_factor, max_image_shift, chunk_size=None): + """Evaluate all candidate drift vectors in parallel. + + Warps both images for each candidate using ``bilinear_kde_batch`` + and scores alignment quality via ``cross_corr_batch``. Without + batching, each candidate would be a separate Python iteration - this + is the key operation that enables the 300x speedup. + + Parameters + ---------- + drift_vectors : ndarray, shape (N, 2) + Candidate drift vectors to test, columns are (row, col). + upsample_factor : int + Subpixel cross-correlation upsampling factor. + max_image_shift : float or None + Maximum allowed shift for cross-correlation peak search. + chunk_size : int or None + Number of candidates per pass. If None, all at once. + + Returns + ------- + tuple[int, torch.Tensor] + Index of the best candidate in ``drift_vectors``, and the full + cost tensor of shape ``(N,)`` for all candidates (used by + verbose mode to rank runner-ups). + """ + device = self._device + dtype = self._dtype + num_candidates = drift_vectors.shape[0] + drift_vectors_t = torch.tensor(drift_vectors, dtype=dtype, device=device) + canvas_shape = (self.shape[1], self.shape[2]) + # Base coordinates shared across all candidates + base_data = [] + for img_idx in range(2): + knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) + row_base, col_base = transform_coordinates_single_knot( + knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) + num_rows = self.knots[img_idx].shape[1] + scanline_offset = (torch.arange(num_rows, dtype=dtype, device=device) + - (num_rows - 1) / 2) + base_data.append((self.images_t[img_idx], row_base, col_base, scanline_offset)) + # Precompute shift mask and frequency grids (shared across chunks) + shift_mask = None + if max_image_shift is not None: + canvas_rows, canvas_cols = canvas_shape + freq_row = fftfreq(canvas_rows, 1.0 / canvas_rows, device=device, dtype=dtype) + freq_col = fftfreq(canvas_cols, 1.0 / canvas_cols, device=device, dtype=dtype) + shift_mask = freq_row[:, None] ** 2 + freq_col[None, :] ** 2 >= max_image_shift ** 2 + freq_grids = ( + fftfreq(canvas_shape[0], device=device, dtype=dtype)[:, None], + fftfreq(canvas_shape[1], device=device, dtype=dtype)[None, :], + ) + if chunk_size is None: + chunk_size = self._auto_chunk_size(num_candidates, canvas_shape, dtype, device) + on_cuda = torch.device(device).type == "cuda" + chunked = on_cuda and chunk_size < num_candidates + all_costs = [] + chunk_start = 0 + chunk_idx = 0 + while chunk_start < num_candidates: + chunk_end = min(chunk_start + chunk_size, num_candidates) + drift_chunk = drift_vectors_t[chunk_start:chunk_end] + if chunk_idx == 0 and chunked: + torch.cuda.reset_peak_memory_stats(device) + warped_pair = [] + for img_idx in range(2): + image_t, row_base, col_base, scanline_offset = base_data[img_idx] + row_candidates = row_base[None] + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] + col_candidates = col_base[None] + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] + warped, _ = bilinear_kde_batch( + row_candidates, col_candidates, image_t, + canvas_shape, self.kde_sigma, + self.pad_value[img_idx]) + warped_pair.append(warped) + all_costs.append(cross_corr_batch( + warped_pair[0], warped_pair[1], + upsample_factor, + max_shift_mask=shift_mask, + freq_grids=freq_grids)) + # After chunk 0, replace the conservative static estimate with the + # actual measured per-candidate cost and print one summary line so + # the user can see how the chunking adapted to their GPU state. + if chunk_idx == 0 and chunked: + per_candidate_actual = torch.cuda.max_memory_allocated(device) / chunk_size + free_bytes, total_bytes = torch.cuda.mem_get_info(device) + tuned_chunk_size = max(1, int(free_bytes * 0.5 / per_candidate_actual)) + tuned_chunk_size = min(tuned_chunk_size, num_candidates) + if tuned_chunk_size > chunk_size: + chunk_size = tuned_chunk_size + num_chunks_final = 1 + (num_candidates - chunk_end + chunk_size - 1) // chunk_size + print( + f" affine grid: {num_candidates} cand × {canvas_shape[0]}×{canvas_shape[1]}, " + f"{per_candidate_actual / 1e9:.2f} GB/cand → {chunk_size}/chunk × {num_chunks_final} passes " + f"({free_bytes / 1e9:.0f}/{total_bytes / 1e9:.0f} GB free)" + ) + chunk_start = chunk_end + chunk_idx += 1 + all_costs = torch.cat(all_costs) + return torch.argmin(all_costs).item(), all_costs + + +def _auto_chunk_size(num_candidates, canvas_shape, dtype, device): + """Pick a candidate-batch size that fits in current free GPU memory. + + Empirical per-candidate peak (measured at 4096×4096): bilinear KDE + scatter buffers, gaussian smoothing temporaries, then cross-correlation + FFT pairs (complex64) - together about ``32 × canvas_pixels`` + ``× dtype_bytes`` at peak. We sample free memory at call time, divide + by that estimate with a 0.4 safety factor, and cap the result at + ``num_candidates`` (no point splitting if it all fits). + On CPU we just process all candidates at once - no VRAM constraint. + """ + device = torch.device(device) + if device.type != "cuda": + return num_candidates + bytes_per_element = torch.finfo(dtype).bits // 8 + per_candidate_bytes = canvas_shape[0] * canvas_shape[1] * bytes_per_element * 32 + free_bytes, _ = torch.cuda.mem_get_info(device) + chunk_size = max(1, int(free_bytes * 0.4 / per_candidate_bytes)) + return min(chunk_size, num_candidates) + + +@torch.inference_mode() +def _warp_and_translate_torch( + self, + max_image_shift: float | None, + upsample_factor: int = 8, + knots_batch: torch.Tensor | None = None, + solve_translation: bool = True, +) -> torch.Tensor: + """Regenerate warped images and solve translation on GPU. + + Three phases: warp → solve translation → re-warp. When ``knots_batch`` + is provided, reads/writes a single batched torch tensor (zero numpy + crossings). Without it, falls back to ``self.knots`` (numpy) for + compatibility with ``align_affine``. + + Set ``solve_translation=False`` to only warp and sync without + re-solving translation - used after the nonrigid loop to populate + ``self.images_warped`` from final knots. + + Parameters + ---------- + max_image_shift : float or None + Maximum allowed translational shift in pixels. + upsample_factor : int + Sub-pixel precision for cross-correlation (1/N pixel). + knots_batch : torch.Tensor or None + If provided, batched ``(N, 2, num_rows)`` torch tensor on GPU. + Translation shifts are applied in-place. Skips numpy sync. + solve_translation : bool + If False, skip translation alignment (Phase 2+3). Only warp + once using current knots and sync to CPU. + + Returns + ------- + torch.Tensor + Warped images on GPU, shape ``(num_images, H, W)``. + """ + device = self._device + dtype = self._dtype + num_images = self.shape[0] + canvas_shape = (self.shape[1], self.shape[2]) + + def _warp_all(warped_t, weights_t): + """Warp all images onto the canvas using current knots.""" + for img_idx in range(num_images): + if knots_batch is not None: + # transform_coordinates_single_knot expects (2, N, 1) + knots_img = knots_batch[img_idx].detach()[:, :, None] + else: + knots_img = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) + row_t, col_t = transform_coordinates_single_knot( + knots_img, self.scan_fast_t[img_idx], self.images[img_idx].shape) + warped, weights = bilinear_kde_batch( + row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, + self.kde_sigma, self.pad_value[img_idx]) + warped_t[img_idx] = warped[0] + weights_t[img_idx] = weights[0] + + warped_t = torch.zeros(num_images, *canvas_shape, dtype=dtype, device=device) + weights_t = torch.zeros_like(warped_t) + _warp_all(warped_t, weights_t) + if not solve_translation: + self.images_warped.array[:] = warped_t.cpu().numpy() + self.weights_warped.array[:] = weights_t.cpu().numpy() + return warped_t + # Solve translation shifts and apply to knots + shifts_t = translate_align(warped_t, upsample_factor, max_image_shift) + if knots_batch is not None: + knots_batch[:, 0] += shifts_t[:, 0:1] + knots_batch[:, 1] += shifts_t[:, 1:2] + else: + shifts_np = shifts_t.cpu().numpy() + for img_idx in range(num_images): + self.knots[img_idx][0] += shifts_np[img_idx, 0] + self.knots[img_idx][1] += shifts_np[img_idx, 1] + # Re-warp with corrected knots + _warp_all(warped_t, weights_t) + if knots_batch is None: + self.images_warped.array[:] = warped_t.cpu().numpy() + self.weights_warped.array[:] = weights_t.cpu().numpy() + return warped_t diff --git a/src/quantem/imaging/drift/core/nonrigid.py b/src/quantem/imaging/drift/core/nonrigid.py new file mode 100644 index 00000000..cd701fed --- /dev/null +++ b/src/quantem/imaging/drift/core/nonrigid.py @@ -0,0 +1,449 @@ + +import numpy as np +import torch +from scipy.ndimage import gaussian_filter +from scipy.optimize import minimize +from tqdm import tqdm + +from quantem.imaging.drift_utils import ( + gaussian_smooth_1d, +) + + +def align_nonrigid( + self, + backend: str = "pytorch", + optimizer_name: str = "adam", + num_iterations: int = 8, + regularization_sigma_px: float = 16.0, + regularization_update_step_size: float | None = 0.8, + regularization_poly_order: int = 1, + max_image_shift: float | None = 32.0, + adam_steps: int = 30, + lr: float | None = None, + lbfgs_max_iter: int = 20, + max_optimize_iterations: int = 10, + regularization_max_image_shift_px: float | None = None, + solve_individual_rows: bool = True, + show_merged: bool = True, + show_images: bool = False, + show_knots: bool = True, + **kwargs, +): + """ + Non-rigid drift correction using PyTorch (default) or SciPy backend. + + Parameters + ---------- + backend : str, default "pytorch" + Optimization backend. + - "pytorch": GPU-accelerated batched optimization. Single-knot only. + - "scipy": CPU L-BFGS row-by-row. Use when you need multi-knot + mode (``number_knots > 1``), which the pytorch path does not + yet support. + optimizer_name : str, default "adam" + PyTorch optimizer (ignored if backend="scipy"). + + **"adam"** - first-order momentum optimizer. Default. Best when: + - You want the fastest possible runtime, especially at image + sizes ≤512 px where the per-step grid_sample is small and + Adam's tight inner loop wins on launch overhead. + - You're confident ``max_image_shift`` reflects the true drift + bound (Adam's auto-lr derives from it; if it's too small, + Adam silently under-converges). + - You want bit-reproducible results across runs (LBFGS line + search has subtle non-determinism from Wolfe condition checks). + + **Provisional override guidance** (validated on one real-data + pair - Bob's gold-nanoparticle HAADF on a spectra background - + and the synthetic chevron test; needs broader testing on + diverse datasets before being treated as authoritative). If + you override ``lr`` manually, the rough formula is + ``expected_drift_px / (num_iterations * adam_steps)``. + Indicative starting values for the default ``num_iterations=8, + adam_steps=30`` (240 total steps): + * ~5 px drift (synthetic chevron, small drift): ``lr≈0.02`` + * ~50-100 px drift (gold-nanoparticle HAADF, real STEM): ``lr≈0.5`` + * larger / unknown drift: prefer ``optimizer_name="lbfgs"`` + which auto-scales via line search and doesn't need this + per-dataset tuning. + + **"lbfgs"** - quasi-Newton optimizer with strong-Wolfe line search. + Best when: + - The image is ≥512 px and you don't mind paying Python closure + overhead for fewer total steps (typically 2-3× faster than + Adam at 2048+ px because it converges in ~30 steps not 240). + - You're unsure about the drift magnitude or don't want to think + about ``lr`` tuning - LBFGS line search auto-scales the step + without any hand-tuning. + - You want quality over speed. + + **Failure modes to avoid:** + - **Don't normalize inputs to [0, 1] when using LBFGS** - + strong-Wolfe's curvature condition needs absolute gradient + magnitude above a threshold; with normalized intensities the + gradient is ~1e-4 and the line search returns step=0, + producing zero correction silently. Adam is unaffected. + - **Don't set ``max_image_shift`` smaller than your actual drift + if using Adam with default ``lr=None``** - Adam's auto-derived + lr scales with max_image_shift, so a too-small bound silently + clamps how much drift Adam can recover. LBFGS is unaffected. + + If unsure, start with Adam (the default) for ≤1024 px images and + switch to LBFGS for ≥2048 px or for unknown-drift exploratory work. + + Shared Parameters + ----------------- + num_iterations : int, default 8 + Number of outer iterations for alternating optimization. + regularization_sigma_px : float, default 16.0 + Gaussian smoothing sigma for knot regularization. + regularization_update_step_size : float, default 0.8 + Step size for knot updates (0-1, lower = more conservative). + regularization_poly_order : int, default 1 + Polynomial order for trend removal in knot regularization + (used by both pytorch and scipy backends). + max_image_shift : float, default 32.0 + Maximum shift for translation alignment between iterations. + + PyTorch Parameters (ignored if backend="scipy") + ----------------------------------------------- + adam_steps : int, default 30 + Number of Adam optimization steps per outer iteration. + lr : float or None, default None + Learning rate for Adam. When None (default), auto-derived as + ``max_image_shift / (num_iterations * adam_steps * 4)``. + + **Why auto-derive?** Adam's ``m/sqrt(v)`` update self-normalizes + the gradient, so each step moves a knot by ~``lr`` pixels + regardless of image intensity scale. The total movement budget + is ``lr × num_iterations × adam_steps`` and is hard-bounded: + Adam cannot find drift larger than that budget no matter how + many iterations you give it. This means ``lr`` must be matched + to the EXPECTED DRIFT MAGNITUDE IN PIXELS, not to gradient + magnitude - the same default value that works on small synthetic + drift will silently under-converge on real data with larger drift. + + The auto-derived formula reserves half the total step budget for + search (covering up to ``max_image_shift / 2`` of nonlinear drift) + and the other half for refinement near the minimum. + + Override with an explicit float when you know the actual drift + magnitude - e.g. ``lr=2.0`` for very-large-drift in-situ data, + or ``lr=0.005`` for atomic-resolution stable samples. + lbfgs_max_iter : int, default 20 + Maximum LBFGS iterations per outer iteration (line search probes + within each iter happen automatically). Only used when + optimizer="lbfgs". + + SciPy Parameters (ignored if backend="pytorch") + ----------------------------------------------- + max_optimize_iterations : int, default 10 + Maximum L-BFGS iterations per row. + regularization_max_image_shift_px : float, optional + Maximum allowed shift per iteration. + solve_individual_rows : bool, default True + If True, optimize each row independently. + + Display Parameters + ------------------ + show_merged : bool, default True + Show merged image after alignment. + show_images : bool, default False + Show individual aligned images. + show_knots : bool, default True + Overlay knot positions on visualizations. + + Notes + ----- + With backend="pytorch", ``self.images_warped`` is left STALE after + the loop and refreshed lazily on first access via plot methods or + ``calculate_error()``. Code that reads ``self.images_warped.array`` + directly should call ``self._ensure_warped_images()`` first, or + use ``generate_corrected_image()`` which builds its own warps from + ``self.knots``. + """ + if not hasattr(self, "knots"): + raise RuntimeError( + "No knots found. Call .preprocess() before running alignment." + ) + if backend == "pytorch": + device = self._device + dtype = self._dtype + num_images = self.shape[0] + canvas_shape = (self.shape[1], self.shape[2]) + if any(self.knots[i].shape[2] != 1 for i in range(num_images)): + raise NotImplementedError( + "PyTorch backend only supports single knot. " + "Use backend='scipy' for multiple knots.") + knots_batch = torch.tensor( + np.stack([self.knots[i][:, :, 0] for i in range(num_images)]), + dtype=dtype, device=device, requires_grad=True) + num_rows_knot = knots_batch.shape[2] + target_batch = torch.stack(self.images_t) + # Build u tensors once and reuse - same scan-position vector projects + # onto row and col offsets via the per-image scan_fast components. + u_t = [ + torch.as_tensor(self.interpolator[i].u, dtype=dtype, device=device) + for i in range(num_images) + ] + row_scan_offsets = torch.stack([ + u_t[i] * (self.interpolator[i].scan_fast[0] * (self.images[i].shape[0] - 1)) + for i in range(num_images) + ]) + col_scan_offsets = torch.stack([ + u_t[i] * (self.interpolator[i].scan_fast[1] * (self.images[i].shape[1] - 1)) + for i in range(num_images) + ]) + row_scale = 2.0 / (canvas_shape[0] - 1) + col_scale = 2.0 / (canvas_shape[1] - 1) + if optimizer_name == "adam": + # Auto-derive lr so the total movement budget covers a quarter + # of max_image_shift. The safety factor of 4 (not 2) prevents + # over-shooting at small image sizes where actual drift is well + # below max_image_shift; at large sizes the same factor still + # converges because the loss surface is smoother. See the `lr` + # parameter docstring for the full rationale. + adam_lr = lr if lr is not None else max_image_shift / (num_iterations * adam_steps * 4) + optimizer = torch.optim.Adam([knots_batch], lr=adam_lr, fused=True) + elif optimizer_name == "lbfgs": + optimizer = torch.optim.LBFGS( + [knots_batch], lr=1.0, max_iter=lbfgs_max_iter, + line_search_fn="strong_wolfe") + else: + raise ValueError(f"optimizer_name must be 'adam' or 'lbfgs', got {optimizer_name!r}") + if regularization_sigma_px is not None and regularization_sigma_px > 0: + x_knot = torch.arange(num_rows_knot, dtype=dtype, device=device) + x_norm = (x_knot - x_knot.mean()) / x_knot.std() + vander = torch.stack([x_norm ** p for p in range(regularization_poly_order + 1)], dim=1) + warped_t = self._warp_and_translate_torch( + max_image_shift, upsample_factor=8, knots_batch=knots_batch) + error_buffer = [] + for _ in tqdm(range(num_iterations), desc=f"Solving nonrigid drift ({optimizer_name})"): + # Build the reference under no_grad: arithmetic on warped_t (an + # inference tensor) would otherwise return an autograd-tracked + # leaf, and the optimizer would build a graph through it. + with torch.no_grad(): + warped_sum = warped_t.sum(0) + ref_batch = (warped_sum[None] - warped_t) / (num_images - 1) + knots_prev = knots_batch.detach().clone() + # Regularization alters the loss surface between outer iters, so + # stale momentum / curvature history would push knots the wrong way. + optimizer.state.clear() + if optimizer_name == "adam": + self._optimize_knots_adam( + ref_batch, target_batch, knots_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale, + optimizer, adam_steps) + else: + self._optimize_knots_lbfgs( + ref_batch, target_batch, knots_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale, + optimizer) + self._regularize_knots( + knots_batch, knots_prev, vander, + regularization_max_image_shift_px, + regularization_sigma_px, + regularization_update_step_size) + warped_t = self._warp_and_translate_torch( + max_image_shift, upsample_factor=8, knots_batch=knots_batch) + # Per-iter error stays on GPU; sync once after the loop + images_mean = warped_t.mean(dim=0) + error_buffer.append(torch.mean(torch.abs(warped_t - images_mean[None]), dim=(1, 2))) + # Sync knots back to numpy; leave images_warped lazy so callers + # that never plot avoid the GPU→CPU transfer of the warped stack. + knots_final = knots_batch.detach().cpu().numpy() + for img_idx in range(num_images): + self.knots[img_idx][:, :, 0] = knots_final[img_idx] + self._images_warped_stale = True + self._max_image_shift_cached = max_image_shift + if error_buffer: + # Build all error rows in one DtoH transfer + one vstack, instead of + # the quadratic vstack-per-iteration pattern used by calculate_error. + errors_np = torch.stack(error_buffer).cpu().numpy() # (num_iterations, num_images) + mode_col = np.full((len(errors_np), 1), 2.0) + mean_col = errors_np.mean(axis=1, keepdims=True) + new_rows = np.hstack((mode_col, mean_col, errors_np)) + if not hasattr(self, "error_track"): + self.error_track = new_rows + else: + self.error_track = np.vstack((self.error_track, new_rows)) + else: + for _ in tqdm(range(num_iterations), desc="Solving nonrigid drift (scipy)"): + for ind in range(self.shape[0]): + image_ref = np.delete(self.images_warped.array, ind, axis=0).mean(axis=0) + knots_updated = self._optimize_knots_scipy( + ind, image_ref, self.knots[ind], + max_optimize_iterations=max_optimize_iterations, + solve_individual_rows=solve_individual_rows) + if regularization_max_image_shift_px is not None: + knots_shift = knots_updated - self.knots[ind] + knots_dist = np.sqrt(np.sum(knots_shift**2, axis=0)) + sub = knots_dist > regularization_max_image_shift_px + knots_updated[0][sub] = (self.knots[ind][0][sub] + + knots_shift[0][sub] * regularization_max_image_shift_px / knots_dist[sub]) + knots_updated[1][sub] = (self.knots[ind][1][sub] + + knots_shift[1][sub] * regularization_max_image_shift_px / knots_dist[sub]) + if regularization_sigma_px is not None and regularization_sigma_px > 0: + knots_smoothed = knots_updated.copy() + for dim in range(2): + x = np.arange(knots_updated.shape[1]) + for knot_ind in range(knots_updated.shape[2]): + y = knots_updated[dim, :, knot_ind] + coefs = np.polyfit(x, y, deg=regularization_poly_order) + trend = np.polyval(coefs, x) + residual = y - trend + residual_smooth = gaussian_filter(residual, sigma=regularization_sigma_px) + knots_smoothed[dim, :, knot_ind] = residual_smooth + trend + knots_updated = knots_smoothed + if regularization_update_step_size is not None: + knots_updated = (self.knots[ind] + + (knots_updated - self.knots[ind]) * regularization_update_step_size) + self.knots[ind] = knots_updated + warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor=8) + self.calculate_error(2, _warped_t=warped_t) + + if show_merged: + self.plot_merged_images( + show_knots=show_knots, + title="Merged: non-rigid", + **kwargs, + ) + + if show_images: + self.plot_transformed_images( + show_knots=show_knots, + title=[f"Image {i}: non-rigid" for i in range(self.shape[0])], + **kwargs, + ) + + return self + + +def _optimize_knots_adam( + self, ref_batch, target_batch, knots_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale, + optimizer, adam_steps, +): + """Run ``adam_steps`` of Adam on a batched knot tensor against ``_compiled_loss_fn``.""" + ref_t = ref_batch[:, None] + for _ in range(adam_steps): + optimizer.zero_grad() + loss = self._compiled_loss_fn( + knots_batch, ref_t, target_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale) + loss.backward() + optimizer.step() + + +@torch.compile(mode="reduce-overhead", dynamic=False) +def _compiled_loss_fn( + knots_batch, ref_t, target_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale, +): + """Fused forward pass: knot offsets → grid → grid_sample → MSE loss. + + The MSE is averaged over both the batch (N images) and the spatial + dims, so each image's gradient is scaled by 1/N relative to a + per-image solve. Adam's adaptive step size absorbs the constant + rescale; LBFGS line search rescales itself. + """ + grid_row = (knots_batch[:, 0, :, None] + row_scan_offsets[:, None, :]) * row_scale - 1.0 + grid_col = (knots_batch[:, 1, :, None] + col_scan_offsets[:, None, :]) * col_scale - 1.0 + grid = torch.stack([grid_col, grid_row], dim=-1) + warped = torch.nn.functional.grid_sample( + ref_t, grid, mode='bilinear', align_corners=True, padding_mode='border')[:, 0] + return ((warped - target_batch) ** 2).mean() + + +def _optimize_knots_lbfgs( + self, ref_batch, target_batch, knots_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale, + optimizer, +): + """Run one LBFGS outer step (line search re-evaluates the closure several times).""" + ref_t = ref_batch[:, None] + def closure(): + optimizer.zero_grad() + loss = self._compiled_loss_fn( + knots_batch, ref_t, target_batch, + row_scan_offsets, col_scan_offsets, row_scale, col_scale) + loss.backward() + return loss + optimizer.step(closure) + + +def _regularize_knots( + self, knots_batch, knots_prev, vander, + max_shift_px, sigma_px, step_size, +): + """Apply per-iteration knot regularization (in-place on ``knots_batch``). + + Three independent stages, each gated by its parameter being non-None: + 1. Per-knot shift cap: clamp ``|new - prev|`` to ``max_shift_px`` + so the optimizer can't move any knot too far in one outer iter. + 2. Polynomial detrend + Gaussian smooth: keep low-order trends, + smooth the residual along the scan-line dimension. Removes + high-frequency optimizer wobble while preserving the drift signal. + 3. Step-size blend: ``new = prev + step_size · (new - prev)``, + under-relaxes the update for stability across outer iterations. + """ + num_images, _, num_rows_knot = knots_batch.shape + with torch.no_grad(): + if max_shift_px is not None: + shift = knots_batch - knots_prev + dist = torch.norm(shift, dim=1, keepdim=True) + scale_factor = torch.clamp(max_shift_px / dist.clamp(min=1e-8), max=1.0) + knots_batch.copy_(knots_prev + shift * scale_factor) + if sigma_px is not None and sigma_px > 0: + # Detrend + smooth all (N*2, num_rows) knots in one batched lstsq + smooth + knots_flat = knots_batch.reshape(-1, num_rows_knot).T # (num_rows, N*2) + coefs, _, _, _ = torch.linalg.lstsq(vander, knots_flat) + trend = (vander @ coefs).T # (N*2, num_rows) + residual = knots_batch.reshape(-1, num_rows_knot) - trend + smoothed = gaussian_smooth_1d(residual, sigma_px) + knots_batch.copy_((smoothed + trend).reshape(num_images, 2, num_rows_knot)) + if step_size is not None: + knots_batch.copy_(knots_prev + (knots_batch - knots_prev) * step_size) + + +def _optimize_knots_scipy( + self, idx: int, image_ref: np.ndarray, knots_init: np.ndarray, + max_optimize_iterations: int = 10, solve_individual_rows: bool = True, +) -> np.ndarray: + """SciPy L-BFGS optimization for one image.""" + shape_knots = knots_init.shape + options = {"maxiter": max_optimize_iterations} if max_optimize_iterations else {} + if solve_individual_rows: + knots_updated = np.zeros_like(knots_init) + for row_ind in range(knots_init.shape[1]): + x0 = knots_init[:, row_ind, :].ravel() + def cost_function(x): + knots_row = x.reshape(shape_knots[0], shape_knots[2]) + xa, ya = self.interpolator[idx].transform_rows(knots_row) + xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) + yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) + dx, dy = xa - xf, ya - yf + warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) + + image_ref[xf + 1, yf] * dx * (1 - dy) + + image_ref[xf, yf + 1] * (1 - dx) * dy + + image_ref[xf + 1, yf + 1] * dx * dy) + return np.sum((warped - self.images[idx].array[row_ind, :]) ** 2) + result = minimize(cost_function, x0, method="L-BFGS-B", options=options) + knots_updated[:, row_ind, :] = result.x.reshape((2, -1)) + else: + x0 = knots_init.ravel() + def cost_function(x): + knots = x.reshape(shape_knots) + xa, ya = self.interpolator[idx].transform_coordinates(knots) + xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) + yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) + dx, dy = xa - xf, ya - yf + warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) + + image_ref[xf + 1, yf] * dx * (1 - dy) + + image_ref[xf, yf + 1] * (1 - dx) * dy + + image_ref[xf + 1, yf + 1] * dx * dy) + return np.sum((warped - self.images[idx].array) ** 2) + result = minimize(cost_function, x0, method="L-BFGS-B", options=options) + knots_updated = result.x.reshape(shape_knots) + return knots_updated diff --git a/src/quantem/imaging/drift/core/warping.py b/src/quantem/imaging/drift/core/warping.py new file mode 100644 index 00000000..5d8017cf --- /dev/null +++ b/src/quantem/imaging/drift/core/warping.py @@ -0,0 +1,196 @@ +import warnings + +import numpy as np +import torch +from numpy.typing import NDArray +from scipy.interpolate import interp1d + +from quantem.core.utils.imaging_utils import ( + bilinear_kde, +) + + +class DriftInterpolator: + def __init__( + self, + input_shape, + output_shape, + scan_fast, + scan_slow, + pad_value, + kde_sigma, + ): + self.input_shape = input_shape + self.output_shape = output_shape + self.scan_fast = scan_fast + self.scan_slow = scan_slow + self.pad_value = pad_value + self.kde_sigma = kde_sigma + + self.rows_input = np.arange(input_shape[0]) + self.cols_input = np.arange(input_shape[1]) + self.u = np.linspace(0, 1, input_shape[1]) + + def transform_rows( + self, + knots_row: NDArray, + ): + num_knots = knots_row.shape[-1] + basis = np.linspace(0, 1, num_knots) + + if num_knots == 1: + xa = knots_row[0] + self.u[None, :] * self.scan_fast[0] * (self.input_shape[0] - 1) + ya = knots_row[1] + self.u[None, :] * self.scan_fast[1] * (self.input_shape[1] - 1) + elif num_knots == 2: + xa = interp1d(basis, knots_row[0], kind="linear", assume_sorted=True)(self.u) + ya = interp1d(basis, knots_row[1], kind="linear", assume_sorted=True)(self.u) + else: + kind = "quadratic" if num_knots == 3 else "cubic" + xa = interp1d( + basis, + knots_row[0], + kind=kind, + fill_value="extrapolate", + assume_sorted=True, + )(self.u) + ya = interp1d( + basis, + knots_row[1], + kind=kind, + fill_value="extrapolate", + assume_sorted=True, + )(self.u) + + return xa, ya + + def transform_coordinates( + self, + knots: NDArray, + ): + num_knots = knots.shape[-1] + + if num_knots == 1: + # vectorized version for speed + xa, ya = self.transform_rows(knots) + else: + xa = np.zeros(self.input_shape) + ya = np.zeros(self.input_shape) + for i in range(self.input_shape[0]): + xa[i], ya[i] = self.transform_rows(knots[:, i]) + + return xa, ya + + def warp_image( + self, + image: NDArray, + knots: NDArray, # shape: (2, rows, num_knots) + kde_sigma=None, + output_shape=None, + pad_value=None, + upsample_factor=None, + ) -> NDArray: + xa, ya = self.transform_coordinates( + knots, + ) + + if kde_sigma is None: + kde_sigma = self.kde_sigma + + if output_shape is None: + output_shape = self.output_shape + + if pad_value is None: + pad_value = self.pad_value + + if upsample_factor is None: + upsample_factor = 1.0 + + image_interp, weight_interp = bilinear_kde( + xa=xa * upsample_factor, # rows + ya=ya * upsample_factor, # cols + values=image, + output_shape=np.round(np.array(output_shape) * upsample_factor).astype("int"), + kde_sigma=kde_sigma * upsample_factor, + pad_value=pad_value, + return_pix_count=True, + ) + + return image_interp, weight_interp + + +def bounded_sine_sigmoid(x, midpoint=0.5, width=1.0): + """ + Piecewise bounded sigmoid: zero, raised sine squared, one. + + Parameters + ---------- + x : array-like, shape (...,) + Input values in [0, 1]. + midpoint : float + Center of the sigmoid transition. + width : float + Width of the sigmoid (range over which it ramps from 0 to 1). + Returns + ------- + y : array-like + Output in [0, 1], same shape as x. + """ + x = np.asarray(x) + # Truncate width if midpoint too close to edge + left_max = midpoint - width / 2 + right_min = midpoint + width / 2 + if left_max < 0: + warnings.warn( + f"width={width} is too large for midpoint={midpoint}, " + f"clamping width to {2 * midpoint}.", + RuntimeWarning, + ) + width = 2 * midpoint + + if right_min > 1: + warnings.warn( + f"width={width} is too large for midpoint={midpoint}, " + f"clamping width to {2 * (1 - midpoint)}.", + RuntimeWarning, + ) + width = 2 * (1 - midpoint) + # Recalculate edges + left = midpoint - width / 2 + right = midpoint + width / 2 + + y = np.zeros_like(x, dtype=float) + in_band = (x >= left) & (x <= right) + # Map [left, right] to [0, pi/2] + t = (x[in_band] - left) / width # goes from 0 to 1 + y[in_band] = np.sin(t * np.pi / 2) ** 2 + y[x > right] = 1.0 + return y + + +def _bounded_sine_sigmoid_torch( + x: torch.Tensor, + midpoint: float = 0.5, + width: float = 1.0, +) -> torch.Tensor: + width = min(width, 2 * midpoint, 2 * (1 - midpoint)) + left = midpoint - width / 2 + right = midpoint + width / 2 + t = ((x - left) / width).clamp(0.0, 1.0) + return torch.where(x > right, torch.ones_like(x), torch.sin(t * (np.pi / 2)) ** 2) + + +def _fourier_crop_torch( + fft_array: torch.Tensor, + crop_shape: tuple[int, int], +) -> torch.Tensor: + crop_h, crop_w = crop_shape + h1 = crop_h // 2 + h2 = crop_h - h1 + w1 = crop_w // 2 + w2 = crop_w - w1 + result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) + result[:h1, :w1] = fft_array[:h1, :w1] + result[:h1, -w2:] = fft_array[:h1, -w2:] + result[-h2:, :w1] = fft_array[-h2:, :w1] + result[-h2:, -w2:] = fft_array[-h2:, -w2:] + return result diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py new file mode 100644 index 00000000..c75f1cda --- /dev/null +++ b/src/quantem/imaging/drift/correction.py @@ -0,0 +1,186 @@ +from collections.abc import Sequence +from typing import Self + +import torch +from numpy.typing import NDArray + +from quantem.core.config import validate_device +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset3d import Dataset3d +from quantem.core.io.serialize import AutoSerialize +from quantem.core.utils.compound_validators import ( + validate_list_of_dataset2d, +) +from quantem.core.utils.validators import ensure_valid_array + +from . import apply as drift_apply +from . import plot as drift_plot +from . import preparation +from .core import affine, nonrigid + + +class DriftCorrection(AutoSerialize): + """ + DriftCorrection provides translation, affine, and non-rigid drift correction for + sequential 2D images using scan direction metadata and flexible spatial interpolation. + + This class supports input data as numpy arrays, Dataset2d, or Dataset3d instances, + with various padding strategies and configurable spline interpolation of scanline + trajectories via Bézier knot control. + + Features + -------- + - Load data from arrays or files + - Apply initial scanline resampling using Bézier curves + - Align images using translation, affine, or non-rigid optimization + - Visualize intermediate and final results with optional knot overlays + - Serialize state with `.save()` and restore with `.load()` + + Parameters (via `from_data` or `from_file`) + ------------------------------------------- + images : list of 2D arrays, Dataset2d, Dataset3d, or file names, or a 3D numpy array + The image stack to correct for drift. + scan_direction_degrees : list of float + The scan direction angle (in degrees) for each image, measured relative to vertical. + pad_fraction : float, default 0.25 + Fraction of padding to add around each image during interpolation. + pad_value : str, float, or list of float, default 'median' + How to pad outside the image area during warping. Can be: + - One of: 'median', 'mean', 'min', 'max' + - A float quantile value (e.g., 0.25) + - A list of per-image float values + number_knots : int, default 1 + Number of knots to use for Bézier interpolation of scanline trajectories. + We strongly recommend using `number_knots = 1` unless the fast scan direction is + expected to vary within the image. + + Example + ------- + Instantiate the DriftCorrection class, run preprocessing and alignment, and save/load results: + + >>> drift = DriftCorrection.from_data( + ... images=[ + ... image0, # 2D numpy array or Dataset2d + ... image1, + ... ], + ... scan_direction_degrees=[0, 90], + ... ).preprocess( + ... pad_fraction=0.25, + ... pad_value='median', + ... number_knots=1, + ... ) + + >>> drift.align_affine() + >>> drift.align_nonrigid() + >>> drift.plot_merged_images() + >>> image_corr = drift.generate_corrected_image() + + >>> drift.save("drift_result.zip") + >>> drift_reloaded = quantem.io.load("drift_result.zip") + + >>> image_corr.save("image_corrected.zip") + >>> image_corr_reloaded = quantem.io.load("image_corrected.zip") + + Notes + ----- + - Use `align_translation()` for rigid shifts, `align_affine()` for scan-shear or uniform drift, + and `align_nonrigid()` for flexible per-row or per-image correction. + - The class stores resampled images in `self.images_warped` and the control knots in `self.knots`. + - Visualization is supported through `plot_merged_images()` and `plot_transformed_images()`. + + Performance + ----------- + ``align_affine`` uses PyTorch to run all heavy operations on GPU + (works on CUDA, MPS, and CPU). The key optimizations are: + + - **Batched grid search**: all ~97 candidate drift vectors are warped + and scored in parallel, instead of one-at-a-time in a Python loop. + - **Batched bilinear KDE** (``drift_utils.bilinear_kde_batch``): + scatter-based image warping via ``scatter_add_`` with int32 indices. + - **Batched FFT cross-correlation** (``drift_utils.cross_corr_batch``): + sub-pixel alignment using DFT upsampling across all candidates at once. + - **Zero CPU round-trips**: coordinate transforms, Gaussian smoothing, + translation alignment, and error computation all stay on GPU until + the final sync. + + This gives ~300× speedup over the original NumPy implementation + (e.g. 436 s → 1.5 s on 2048×2048 image pairs). + + Memory is automatically chunked when the full batch doesn't fit. + Approximate memory per candidate at common sizes: + + ========== =========== ================ + Input size Canvas size Mem / candidate + ========== =========== ================ + 1024×1024 1280×1280 85 MB + 2048×2048 2560×2560 341 MB + 4096×4096 5120×5120 1.36 GB + ========== =========== ================ + """ + + _token = object() + + def __init__( + self, + images: list[Dataset2d], + scan_direction_degrees: NDArray, + _token: object | None = None, + ): + if _token is not self._token: + raise RuntimeError( + "Use DriftCorrection.from_data() or .from_file() to instantiate this class." + ) + + self.images = images + self.scan_direction_degrees = ensure_valid_array(scan_direction_degrees, ndim=1) + + device, _ = validate_device(None) + self._device = device + self._dtype = torch.float32 + + @classmethod + def from_file( + cls, + file_paths: Sequence[str], + scan_direction_degrees: Sequence[float] | NDArray, + file_type: str | None = None, + ) -> Self: + image_list = [Dataset2d.from_file(fp, file_type=file_type) for fp in file_paths] + return cls.from_data( + image_list, + scan_direction_degrees, + ) + + @classmethod + def from_data( + cls, + images: list[Dataset2d] | list[NDArray] | Dataset3d | NDArray, + scan_direction_degrees: list[float] | NDArray, + ) -> Self: + validated_images = validate_list_of_dataset2d(images) + + return cls( + images=validated_images, + scan_direction_degrees=scan_direction_degrees, + _token=cls._token, + ) + + preprocess = preparation.preprocess + align_translation = preparation.align_translation + align_affine = affine.align_affine + _affine_grid_search_batch = affine._affine_grid_search_batch + _auto_chunk_size = staticmethod(affine._auto_chunk_size) + _warp_and_translate_torch = affine._warp_and_translate_torch + align_nonrigid = nonrigid.align_nonrigid + _optimize_knots_adam = nonrigid._optimize_knots_adam + _compiled_loss_fn = staticmethod(nonrigid._compiled_loss_fn) + _optimize_knots_lbfgs = nonrigid._optimize_knots_lbfgs + _regularize_knots = nonrigid._regularize_knots + _optimize_knots_scipy = nonrigid._optimize_knots_scipy + generate_corrected = drift_apply.generate_corrected + generate_corrected_image = drift_apply.generate_corrected_image + calculate_error = drift_apply.calculate_error + plot_transformed_images = drift_plot.plot_transformed_images + plot_convergence = drift_plot.plot_convergence + _ensure_warped_images = drift_plot._ensure_warped_images + plot_merged_images = drift_plot.plot_merged_images diff --git a/src/quantem/imaging/drift/plot.py b/src/quantem/imaging/drift/plot.py new file mode 100644 index 00000000..ee587567 --- /dev/null +++ b/src/quantem/imaging/drift/plot.py @@ -0,0 +1,111 @@ + +import matplotlib.pyplot as plt +import numpy as np + +from quantem.core.visualization import show_2d + + +def plot_transformed_images(self, show_knots: bool = True, **kwargs): + self._ensure_warped_images() + fig, ax = show_2d( + list(self.images_warped.array), + **kwargs, + ) + if show_knots: + for a0 in range(self.shape[0]): + x = self.knots[a0][0] + y = self.knots[a0][1] + ax[a0].plot( + y, + x, + color="r", + ) + + +def plot_convergence( + self, + figsize=(8, 3), + **kwargs, +): + """ + Plot the convergence of the drift correction. + """ + sub = np.abs(self.error_track[:, 0] - 2) < 0.1 + error = self.error_track[:, 1] + it = np.arange(error.shape[0]) + + from matplotlib.ticker import FormatStrFormatter, MaxNLocator + + fig, ax = plt.subplots(1, 2, figsize=figsize) + color = (1, 0, 0) # red + + # Plot Affine + if np.any(~sub): + ax[0].plot( + it[~sub], + 100 * error[~sub], + marker="o", + color=color, + linestyle="-", + label="Affine", + **kwargs, + ) + ax[0].set_xlabel("Affine Iterations") + ax[0].set_ylabel("Mean Error [%]") + ax[0].xaxis.set_major_locator(MaxNLocator(integer=True)) + ax[0].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) + else: + ax[0].axis("off") + + # Plot Non-Rigid + if np.any(sub): + first_true = np.argmax(sub) + if first_true > 0: + sub[first_true - 1] = True + + ax[1].plot( + it[sub], + 100 * error[sub], + marker="o", + color=color, + linestyle="-", + label="Non-Rigid", + **kwargs, + ) + ax[1].set_xlabel("Non-Rigid Iterations") + ax[1].xaxis.set_major_locator(MaxNLocator(integer=True)) + ax[1].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) + else: + ax[1].axis("off") + + plt.tight_layout() + + return self + + +def _ensure_warped_images(self): + """Lazily populate images_warped from current knots if marked stale.""" + if getattr(self, "_images_warped_stale", False): + self._warp_and_translate_torch( + self._max_image_shift_cached, upsample_factor=8, + solve_translation=False) + self._images_warped_stale = False + + +def plot_merged_images(self, show_knots: bool = True, **kwargs): + """ + Plot the current transformed images, with knot overlays. + """ + self._ensure_warped_images() + fig, ax = show_2d( + self.images_warped.array.mean(0), + **kwargs, + ) + if show_knots: + for a0 in range(self.shape[0]): + x = self.knots[a0][0] + y = self.knots[a0][1] + ax.plot( + y, + x, + ) diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py new file mode 100644 index 00000000..445b6654 --- /dev/null +++ b/src/quantem/imaging/drift/preparation.py @@ -0,0 +1,203 @@ + +import numpy as np +import torch + +from quantem.core.datastructures.dataset3d import Dataset3d +from quantem.core.utils.compound_validators import ( + validate_pad_value, +) +from quantem.core.utils.imaging_utils import ( + cross_correlation_shift, +) +from quantem.imaging.drift.core.warping import DriftInterpolator +from quantem.imaging.drift_utils import ( + bilinear_kde_batch, + transform_coordinates_single_knot, +) + + +def preprocess( + self, + pad_fraction: float = 0.25, + pad_value: float | str | list[float] = "median", + kde_sigma: float = 0.5, + number_knots: int = 1, + show_merged: bool = False, + show_images: bool = False, + show_knots: bool = True, + **kwargs, +): + """Prepare images for drift correction by building the scanline model. + + Computes scan direction vectors, initializes Bezier knots that map + each scanline onto a padded canvas, and generates the initial warped + images. This must be called before any alignment step. + + Without preprocessing, there is no spatial model connecting the raw + images to the shared canvas - alignment methods would have no + coordinates to optimize. + + Parameters + ---------- + pad_fraction : float + Fraction of the image size to add as padding around the canvas. + Larger values give more room for drift but use more memory. + ``pad_fraction=0.25`` adds 25% on each side. + pad_value : float, str, or list[float] + Fill value for pixels outside the image footprint. Can be + ``'median'``, ``'mean'``, ``'min'``, ``'max'``, a quantile + (e.g. ``0.25``), or a per-image list of floats. + kde_sigma : float + Gaussian smoothing sigma (in pixels) applied after bilinear + scatter. Smooths the warped images to reduce scatter noise. + number_knots : int + Number of Bezier knots per scanline. Use ``1`` (recommended) + for linear drift correction. Higher values allow per-scanline + curvature but are slower and rarely needed. + show_merged : bool + Display the merged (averaged) warped images after preprocessing. + show_images : bool + Display each individual warped image after preprocessing. + show_knots : bool + Overlay knot positions on displayed images. + **kwargs + Additional keyword arguments passed to plotting functions. + + Returns + ------- + Self + For method chaining: ``drift.preprocess().align_affine()``. + + Examples + -------- + >>> drift = DriftCorrection.from_data( + ... images=[im0, im1], scan_direction_degrees=[0, 90]) + >>> drift.preprocess(pad_fraction=0.25, kde_sigma=0.5, number_knots=1) + """ + self.pad_fraction = float(pad_fraction) + self.pad_value = validate_pad_value(pad_value, self.images) + self.kde_sigma = float(kde_sigma) + self.number_knots = int(number_knots) + self.scan_direction = np.deg2rad(self.scan_direction_degrees) + self.scan_fast = np.stack( + [np.sin(-self.scan_direction), np.cos(-self.scan_direction)], axis=1) + self.scan_slow = np.stack( + [np.cos(-self.scan_direction), -np.sin(-self.scan_direction)], axis=1) + self.shape = ( + len(self.images), + int(np.round(self.images[0].shape[0] * (1 + self.pad_fraction) / 2) * 2), + int(np.round(self.images[1].shape[1] * (1 + self.pad_fraction) / 2) * 2), + ) + # Initialize knots - each image's scanlines mapped to the padded canvas + self.knots = [] + for img_idx in range(self.shape[0]): + shape = self.images[img_idx].shape + v_slow = np.linspace(-(shape[0] - 1) / 2, (shape[0] - 1) / 2, shape[0]) + u_fast = np.linspace(-(shape[1] - 1) / 2, (shape[1] - 1) / 2, self.number_knots) + row_knots = ((self.shape[1] - 1) / 2 + + u_fast[None, :] * self.scan_fast[img_idx, 0] + + v_slow[:, None] * self.scan_slow[img_idx, 0]) + col_knots = ((self.shape[2] - 1) / 2 + + u_fast[None, :] * self.scan_fast[img_idx, 1] + + v_slow[:, None] * self.scan_slow[img_idx, 1]) + self.knots.append(np.stack([row_knots, col_knots], axis=0)) + self.interpolator = [ + DriftInterpolator( + input_shape=self.images[i].shape, + output_shape=self.shape[1:], + scan_fast=self.scan_fast[i], + scan_slow=self.scan_slow[i], + pad_value=self.pad_value[i], + kde_sigma=self.kde_sigma, + ) + for i in range(self.shape[0]) + ] + # Cache source data on GPU and generate initial warped images + device = self._device + dtype = self._dtype + self.images_t = [ + torch.tensor(self.images[i].array, dtype=dtype, device=device) + for i in range(self.shape[0]) + ] + self.scan_fast_t = [ + torch.tensor(self.scan_fast[i], dtype=dtype, device=device) + for i in range(self.shape[0]) + ] + self.images_warped = Dataset3d.from_shape(self.shape) + self.weights_warped = Dataset3d.from_shape(self.shape) + canvas_shape = (self.shape[1], self.shape[2]) + warped_t = torch.zeros(self.shape[0], *canvas_shape, dtype=dtype, device=device) + for img_idx in range(self.shape[0]): + knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) + row_t, col_t = transform_coordinates_single_knot( + knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) + warped, weights = bilinear_kde_batch( + row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, + self.kde_sigma, self.pad_value[img_idx]) + warped_t[img_idx] = warped[0] + self.images_warped.array[img_idx] = warped[0].cpu().numpy() + self.weights_warped.array[img_idx] = weights[0].cpu().numpy() + self.calculate_error(0, _warped_t=warped_t) + kwargs.pop("title", None) + if show_merged: + self.plot_merged_images(show_knots=show_knots, title="Merged: initial", **kwargs) + if show_images: + self.plot_transformed_images( + show_knots=show_knots, + title=[f"Image {i}: initial" for i in range(self.shape[0])], + **kwargs, + ) + return self + + +def align_translation( + self, + upsample_factor: int = 8, + min_image_shift: float | None = None, + max_image_shift: float = 32, + show_merged: bool = True, + show_images: bool = False, + show_knots: bool = True, + **kwargs, +): + """ + Solve for the translation between all images in DriftCorrection.images_warped + """ + dxy = np.zeros((self.shape[0], 2)) + F_ref = np.fft.fft2(self.images_warped.array[0]) + for ind in range(1, self.shape[0]): + shifts, image_shift = cross_correlation_shift( + F_ref, + np.fft.fft2(self.images_warped.array[ind]), + upsample_factor=upsample_factor, + max_shift=max_image_shift, + fft_input=True, + fft_output=True, + return_shifted_image=True, + ) + dxy[ind, :] = shifts + F_ref = F_ref * ind / (ind + 1) + image_shift / (ind + 1) + dxy -= np.mean(dxy, axis=0) + if min_image_shift is not None: + if np.linalg.norm(dxy[ind]) < min_image_shift: + dxy[ind] = 0.0 + for ind in range(self.shape[0]): + self.knots[ind][0] += dxy[ind, 0] + self.knots[ind][1] += dxy[ind, 1] + for ind in range(self.shape[0]): + self.images_warped.array[ind], self.weights_warped.array[ind] = self.interpolator[ + ind + ].warp_image( + self.images[ind].array, + self.knots[ind], + ) + kwargs.pop("title", None) + if show_merged: + self.plot_merged_images(show_knots=show_knots, title="Merged: translation", **kwargs) + if show_images: + self.plot_transformed_images( + show_knots=show_knots, + title=[f"Image {i}: translation" for i in range(self.shape[0])], + **kwargs, + ) + return self From 27c6a80616b0696f42e6b1f08bf7b0f29bea8f7d Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Mon, 24 Aug 2026 21:51:20 -0700 Subject: [PATCH 03/11] refactor: relocate drift numerical helpers --- src/quantem/imaging/drift/__init__.py | 6 +- src/quantem/imaging/drift/apply.py | 87 ++- src/quantem/imaging/drift/core/affine.py | 6 +- src/quantem/imaging/drift/core/knots.py | 409 +++++++++++++ src/quantem/imaging/drift/core/nonrigid.py | 4 +- src/quantem/imaging/drift/core/strip.py | 1 + src/quantem/imaging/drift/core/warping.py | 511 ++++++++++------ src/quantem/imaging/drift/correction.py | 4 +- src/quantem/imaging/drift/preparation.py | 4 +- src/quantem/imaging/drift_utils.py | 664 --------------------- tests/imaging/test_drift_utils.py | 20 +- 11 files changed, 852 insertions(+), 864 deletions(-) create mode 100644 src/quantem/imaging/drift/core/knots.py create mode 100644 src/quantem/imaging/drift/core/strip.py delete mode 100644 src/quantem/imaging/drift_utils.py diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index a19bef65..e38d7311 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -1,10 +1,12 @@ from quantem.imaging.drift.correction import DriftCorrection as DriftCorrection -from quantem.imaging.drift.core.warping import ( - DriftInterpolator as DriftInterpolator, +from quantem.imaging.drift.apply import ( _bounded_sine_sigmoid_torch as _bounded_sine_sigmoid_torch, _fourier_crop_torch as _fourier_crop_torch, bounded_sine_sigmoid as bounded_sine_sigmoid, ) +from quantem.imaging.drift.core.knots import ( + DriftInterpolator as DriftInterpolator, +) DriftCorrection.__module__ = __name__ DriftInterpolator.__module__ = __name__ diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index 921223ee..1c3a0792 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -1,4 +1,6 @@ +import warnings + import matplotlib.pyplot as plt import numpy as np import torch @@ -9,12 +11,7 @@ fourier_cropping, ) from quantem.core.visualization import show_2d -from quantem.imaging.drift.core.warping import ( - _bounded_sine_sigmoid_torch, - _fourier_crop_torch, - bounded_sine_sigmoid, -) -from quantem.imaging.drift_utils import ( +from quantem.imaging.drift.core.knots import ( bilinear_kde_batch, transform_coordinates_single_knot, ) @@ -397,3 +394,81 @@ def calculate_error( self.error_track = error_current[None, :] # initialize with first row else: self.error_track = np.vstack((self.error_track, error_current)) + + +def bounded_sine_sigmoid(x, midpoint=0.5, width=1.0): + """ + Piecewise bounded sigmoid: zero, raised sine squared, one. + + Parameters + ---------- + x : array-like, shape (...,) + Input values in [0, 1]. + midpoint : float + Center of the sigmoid transition. + width : float + Width of the sigmoid (range over which it ramps from 0 to 1). + Returns + ------- + y : array-like + Output in [0, 1], same shape as x. + """ + x = np.asarray(x) + # Truncate width if midpoint too close to edge + left_max = midpoint - width / 2 + right_min = midpoint + width / 2 + if left_max < 0: + warnings.warn( + f"width={width} is too large for midpoint={midpoint}, " + f"clamping width to {2 * midpoint}.", + RuntimeWarning, + ) + width = 2 * midpoint + + if right_min > 1: + warnings.warn( + f"width={width} is too large for midpoint={midpoint}, " + f"clamping width to {2 * (1 - midpoint)}.", + RuntimeWarning, + ) + width = 2 * (1 - midpoint) + # Recalculate edges + left = midpoint - width / 2 + right = midpoint + width / 2 + + y = np.zeros_like(x, dtype=float) + in_band = (x >= left) & (x <= right) + # Map [left, right] to [0, pi/2] + t = (x[in_band] - left) / width # goes from 0 to 1 + y[in_band] = np.sin(t * np.pi / 2) ** 2 + y[x > right] = 1.0 + return y + + +def _bounded_sine_sigmoid_torch( + x: torch.Tensor, + midpoint: float = 0.5, + width: float = 1.0, +) -> torch.Tensor: + width = min(width, 2 * midpoint, 2 * (1 - midpoint)) + left = midpoint - width / 2 + right = midpoint + width / 2 + t = ((x - left) / width).clamp(0.0, 1.0) + return torch.where(x > right, torch.ones_like(x), torch.sin(t * (np.pi / 2)) ** 2) + + +def _fourier_crop_torch( + fft_array: torch.Tensor, + crop_shape: tuple[int, int], +) -> torch.Tensor: + crop_h, crop_w = crop_shape + h1 = crop_h // 2 + h2 = crop_h - h1 + w1 = crop_w // 2 + w2 = crop_w - w1 + result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) + result[:h1, :w1] = fft_array[:h1, :w1] + result[:h1, -w2:] = fft_array[:h1, -w2:] + result[-h2:, :w1] = fft_array[-h2:, :w1] + result[-h2:, -w2:] = fft_array[-h2:, -w2:] + return result diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index b226246f..933956bc 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -3,10 +3,12 @@ import torch from torch.fft import fftfreq -from quantem.imaging.drift_utils import ( +from quantem.imaging.drift.core.knots import ( bilinear_kde_batch, - cross_corr_batch, transform_coordinates_single_knot, +) +from quantem.imaging.drift.core.warping import ( + cross_corr_batch, translate_align, ) diff --git a/src/quantem/imaging/drift/core/knots.py b/src/quantem/imaging/drift/core/knots.py new file mode 100644 index 00000000..4e1ddb52 --- /dev/null +++ b/src/quantem/imaging/drift/core/knots.py @@ -0,0 +1,409 @@ +"""Knot interpolation and batched scan-coordinate projection.""" + +import numpy as np +import torch +from numpy.typing import NDArray +from scipy.interpolate import interp1d + +from quantem.core.utils.imaging_utils import bilinear_kde + + +class DriftInterpolator: + def __init__( + self, + input_shape, + output_shape, + scan_fast, + scan_slow, + pad_value, + kde_sigma, + ): + self.input_shape = input_shape + self.output_shape = output_shape + self.scan_fast = scan_fast + self.scan_slow = scan_slow + self.pad_value = pad_value + self.kde_sigma = kde_sigma + + self.rows_input = np.arange(input_shape[0]) + self.cols_input = np.arange(input_shape[1]) + self.u = np.linspace(0, 1, input_shape[1]) + + def transform_rows( + self, + knots_row: NDArray, + ): + num_knots = knots_row.shape[-1] + basis = np.linspace(0, 1, num_knots) + + if num_knots == 1: + xa = knots_row[0] + self.u[None, :] * self.scan_fast[0] * (self.input_shape[0] - 1) + ya = knots_row[1] + self.u[None, :] * self.scan_fast[1] * (self.input_shape[1] - 1) + elif num_knots == 2: + xa = interp1d(basis, knots_row[0], kind="linear", assume_sorted=True)(self.u) + ya = interp1d(basis, knots_row[1], kind="linear", assume_sorted=True)(self.u) + else: + kind = "quadratic" if num_knots == 3 else "cubic" + xa = interp1d( + basis, + knots_row[0], + kind=kind, + fill_value="extrapolate", + assume_sorted=True, + )(self.u) + ya = interp1d( + basis, + knots_row[1], + kind=kind, + fill_value="extrapolate", + assume_sorted=True, + )(self.u) + + return xa, ya + + def transform_coordinates( + self, + knots: NDArray, + ): + num_knots = knots.shape[-1] + + if num_knots == 1: + # vectorized version for speed + xa, ya = self.transform_rows(knots) + else: + xa = np.zeros(self.input_shape) + ya = np.zeros(self.input_shape) + for i in range(self.input_shape[0]): + xa[i], ya[i] = self.transform_rows(knots[:, i]) + + return xa, ya + + def warp_image( + self, + image: NDArray, + knots: NDArray, # shape: (2, rows, num_knots) + kde_sigma=None, + output_shape=None, + pad_value=None, + upsample_factor=None, + ) -> NDArray: + xa, ya = self.transform_coordinates( + knots, + ) + + if kde_sigma is None: + kde_sigma = self.kde_sigma + + if output_shape is None: + output_shape = self.output_shape + + if pad_value is None: + pad_value = self.pad_value + + if upsample_factor is None: + upsample_factor = 1.0 + + image_interp, weight_interp = bilinear_kde( + xa=xa * upsample_factor, # rows + ya=ya * upsample_factor, # cols + values=image, + output_shape=np.round(np.array(output_shape) * upsample_factor).astype("int"), + kde_sigma=kde_sigma * upsample_factor, + pad_value=pad_value, + return_pix_count=True, + ) + + return image_interp, weight_interp + + +def transform_coordinates_single_knot( + knots: torch.Tensor, + scan_fast: torch.Tensor, + input_shape: tuple[int, int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-knot fast path: map source pixels to canvas coordinates. + + **Single-knot only.** Each scanline has exactly one (row, col) anchor; + the fast-scan-direction position is filled in by linear interpolation + along the scanline. Multi-knot Bezier interpolation is intentionally + not supported here - that's the scipy backend's job. The pytorch path + optimizes for the common single-knot case (≥95% of real STEM workflows). + + Called by ``preprocess``, ``_affine_grid_search_batch``, and + ``_warp_and_translate_torch`` to map source image pixels onto the + padded output canvas. Without this, the warped images would have + no spatial mapping and the grid search couldn't score test drifts. + + Each input row maps to a line on the canvas: + ``row = knot_row + fraction * scan_fast[0] * (num_rows - 1)`` + ``col = knot_col + fraction * scan_fast[1] * (num_cols - 1)`` + where row and col dimensions scale independently for non-square images. + + Parameters + ---------- + knots : torch.Tensor + Knot positions, shape ``(2, num_rows, 1)``. First dim is (row, col). + The trailing 1 is the single-knot dimension; multi-knot inputs are + rejected by the caller before reaching this function. + scan_fast : torch.Tensor + Fast scan direction vector, shape ``(2,)``. + input_shape : tuple[int, int] + Original image shape ``(num_rows, num_cols)``. + + Returns + ------- + row_coords : torch.Tensor + Row coordinates on canvas, shape ``(num_rows, num_cols)``. + col_coords : torch.Tensor + Column coordinates on canvas, shape ``(num_rows, num_cols)``. + + Examples + -------- + >>> knots = torch.zeros(2, 64, 1) + >>> scan_fast = torch.tensor([0.0, 1.0]) + >>> r, c = transform_coordinates_single_knot(knots, scan_fast, (64, 64)) + >>> r.shape + torch.Size([64, 64]) + """ + num_rows, num_cols = input_shape + fast_fraction = torch.linspace(0, 1, num_cols, dtype=knots.dtype, device=knots.device) + row_coords = knots[0, :, 0:1] + fast_fraction[None, :] * scan_fast[0] * (num_rows - 1) + col_coords = knots[1, :, 0:1] + fast_fraction[None, :] * scan_fast[1] * (num_cols - 1) + return row_coords, col_coords + + +def bilinear_kde_batch( + row_coords: torch.Tensor, + col_coords: torch.Tensor, + source_image: torch.Tensor, + output_shape: tuple[int, int], + kde_sigma: float, + pad_value: float | torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Batched bilinear KDE: scatter N source images onto an output canvas. + + Each pixel scatters its value to its 4 nearest grid neighbors with + bilinear weights ``(1-dr)·(1-dc)``, ``dr·(1-dc)``, ``(1-dr)·dc``, + ``dr·dc`` where ``dr, dc`` are fractional row/col distances. + Accumulated counts and values are Gaussian-smoothed, then normalized: + ``output = pad_value·(1-coverage) + coverage·(values/counts)``. + + Used by both the affine grid search (N = candidate drift vectors, + single source image broadcast across drifts) and the nonrigid loop + (N = stacked source images, one per drift). + + Parameters + ---------- + row_coords : torch.Tensor + Row coordinates of input pixels, shape ``(N, rows, cols)``. + col_coords : torch.Tensor + Column coordinates of input pixels, shape ``(N, rows, cols)``. + source_image : torch.Tensor + Pixel values to scatter. Either ``(rows, cols)`` (same image used + for all N drifts - affine grid search) or ``(N, rows, cols)`` + (different image per drift - multi-image batched warping). + output_shape : tuple[int, int] + Canvas size ``(num_rows, num_cols)`` for the output images. + kde_sigma : float + Gaussian smoothing sigma in pixels. + pad_value : float or torch.Tensor + Fill value where pixel coverage is below threshold. If a tensor of + shape ``(N,)``, applies a different pad value per drift. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + ``(warped_images, sum_weights)`` - warped images and smoothed pixel coverage, + both shape ``(N, num_rows, num_cols)``. + """ + num_test_drifts = row_coords.shape[0] + num_rows, num_cols = output_shape + coverage_threshold = 1e-3 + # Flatten spatial dims - scatter_add_ works on 1D buffers + row_flat = row_coords.flatten(1) + col_flat = col_coords.flatten(1) + # Stay in float for fractional distance, convert to int only for scatter indices + row_floor = row_flat.floor() + col_floor = col_flat.floor() + frac_row = row_flat - row_floor + frac_col = col_flat - col_floor + row_floor = row_floor.int() + col_floor = col_floor.int() + if source_image.dim() == 3: + # Per-drift source images: each drift scatters its own pixel values + source_values_flat = source_image.flatten() + else: + source_values_flat = source_image.flatten().repeat(num_test_drifts) + # All N batch entries scatter into one flat buffer - offset separates them + batch_offsets = ( + torch.arange(num_test_drifts, device=row_coords.device, dtype=torch.int32) + * num_rows * num_cols + )[:, None] + # Float32 accumulators - scatter_add_ requires source dtype to match, + # so all input tensors must be float32 (raises on float64). + sum_weights = torch.zeros( + num_test_drifts * num_rows * num_cols, dtype=torch.float32, device=row_coords.device + ) + sum_values = torch.zeros_like(sum_weights) + # Periodic wrapping so pixels near edges scatter to the opposite side + row_wrapped = row_floor % num_rows + col_wrapped = col_floor % num_cols + row_next = (row_wrapped + 1) % num_rows + col_next = (col_wrapped + 1) % num_cols + # Each pixel distributes its value to the 4 nearest grid neighbors + # weighted by bilinear distance: (1-dr)(1-dc), dr(1-dc), (1-dr)dc, dr·dc + for corner_row, corner_col, corner_weight in [ + (row_wrapped, col_wrapped, ((1 - frac_row) * (1 - frac_col)).flatten()), + (row_next, col_wrapped, (frac_row * (1 - frac_col)).flatten()), + (row_wrapped, col_next, ((1 - frac_row) * frac_col).flatten()), + (row_next, col_next, (frac_row * frac_col).flatten()), + ]: + flat_indices = (corner_row * num_cols + corner_col + batch_offsets).flatten() + sum_weights.scatter_add_(0, flat_indices, corner_weight) + sum_values.scatter_add_(0, flat_indices, corner_weight * source_values_flat) + sum_weights = sum_weights.reshape(num_test_drifts, num_rows, num_cols) + sum_values = sum_values.reshape(num_test_drifts, num_rows, num_cols) + # Smooth the scattered counts and values to fill gaps between pixels + sum_weights = gaussian_smooth_batch(sum_weights, kde_sigma) + sum_values = gaussian_smooth_batch(sum_values, kde_sigma) + # Blend between pad_value (uncovered) and normalized values (covered), + # ramping linearly with coverage to avoid hard edges at the boundary + coverage_weight = torch.clamp(sum_weights / coverage_threshold, max=1.0) + if isinstance(pad_value, torch.Tensor) and pad_value.dim() == 1: + # Per-drift pad value: reshape (N,) → (N, 1, 1) for broadcasting + pad_value = pad_value[:, None, None] + warped_images = pad_value * (1 - coverage_weight) + coverage_weight * ( + sum_values / torch.clamp(sum_weights, min=1e-8) + ) + return warped_images, sum_weights + + +def gaussian_smooth_batch( + field_stack: torch.Tensor, + sigma: float, +) -> torch.Tensor: + """Batched 2D Gaussian smoothing matching ``scipy.ndimage.gaussian_filter``. + + Used by ``bilinear_kde_batch`` to smooth scattered counts and + values before normalization. Without smoothing, the warped images + have salt-and-pepper artifacts from the scatter step. + + Parameters + ---------- + field_stack : torch.Tensor + Input tensor of shape ``(N, num_rows, num_cols)``. + sigma : float + Standard deviation of the Gaussian kernel in pixels. + + Returns + ------- + torch.Tensor + Smoothed tensor of shape ``(N, num_rows, num_cols)``. + + """ + kernel, radius = _gaussian_kernel_1d(sigma, field_stack.dtype, field_stack.device) + # Separable kernel: column pass then row pass to halve FLOPs vs full 2D conv + kernel_col = kernel[None, None, None, :] + kernel_row = kernel[None, None, :, None] + field_stack = field_stack[:, None] + field_stack = torch.nn.functional.conv2d(_symmetric_pad(field_stack, 0, radius), kernel_col) + field_stack = torch.nn.functional.conv2d(_symmetric_pad(field_stack, radius, 0), kernel_row) + return field_stack[:, 0] + + +def gaussian_smooth_1d( + signal: torch.Tensor, + sigma: float, +) -> torch.Tensor: + """1D Gaussian smoothing matching ``scipy.ndimage.gaussian_filter``. + + Smooths each row of the input independently using a separable 1D kernel. + Used for regularizing knot displacement vectors in the nonrigid loop, + where the signal is 1D (one value per scan line). + + Parameters + ---------- + signal : torch.Tensor + Input tensor of shape ``(N, L)`` - N channels, L samples. + sigma : float + Standard deviation of the Gaussian kernel in pixels. + + Returns + ------- + torch.Tensor + Smoothed tensor of shape ``(N, L)``. + """ + kernel, radius = _gaussian_kernel_1d(sigma, signal.dtype, signal.device) + signal_padded = _symmetric_pad_1d(signal[:, None], radius) + return torch.nn.functional.conv1d(signal_padded, kernel[None, None, :])[:, 0] + + +def _symmetric_pad_1d(signal: torch.Tensor, pad: int) -> torch.Tensor: + """Symmetric 1D padding matching scipy's reflect mode. + + Same edge-repeat semantics as ``_symmetric_pad`` but for 1D signals. + Used by ``gaussian_smooth_1d`` for regularization of knot vectors. + """ + left = signal[:, :, :pad].flip(-1) + right = signal[:, :, -pad:].flip(-1) + return torch.cat([left, signal, right], dim=-1) + + +def _symmetric_pad( + field_stack: torch.Tensor, + pad_rows: int, + pad_cols: int, +) -> torch.Tensor: + """Symmetric padding matching scipy's reflect mode for parity. + + Without this, the torch and numpy Gaussian smoothing paths produce + different results near canvas edges, breaking numerical parity. + + Scipy's ``mode='reflect'`` repeats the edge pixel + (``[1,2,3]`` → ``[2,1,1,2,3,3,2]``), but PyTorch's + ``F.pad(mode='reflect')`` does not (``[1,2,3]`` → ``[3,2,1,2,3,2,1]``). + + Parameters + ---------- + field_stack : torch.Tensor + Input tensor of shape ``(N, C, num_rows, num_cols)``. + pad_rows : int + Number of rows to pad on top and bottom. + pad_cols : int + Number of columns to pad on left and right. + + Returns + ------- + torch.Tensor + Padded tensor. + + Examples + -------- + >>> t = torch.tensor([[[[1., 2., 3.]]]]) + >>> _symmetric_pad(t, 0, 2) + tensor([[[[2., 1., 1., 2., 3., 3., 2.]]]]) + """ + if pad_cols > 0: + left = field_stack[:, :, :, :pad_cols].flip(-1) + right = field_stack[:, :, :, -pad_cols:].flip(-1) + field_stack = torch.cat([left, field_stack, right], dim=-1) + if pad_rows > 0: + top = field_stack[:, :, :pad_rows, :].flip(-2) + bottom = field_stack[:, :, -pad_rows:, :].flip(-2) + field_stack = torch.cat([top, field_stack, bottom], dim=-2) + return field_stack + + +def _gaussian_kernel_1d(sigma, dtype, device, _cache={}): + """Normalized 1D Gaussian ``exp(-0.5*(x/sigma)^2)``, radius ``4*sigma``. + + Cached via mutable default arg - the grid search calls this ~800 times + with the same sigma, saving ~44ms of redundant kernel construction. + """ + key = (sigma, dtype, device) + if key not in _cache: + radius = int(4 * sigma + 0.5) + offsets = torch.arange(-radius, radius + 1, dtype=dtype, device=device) + kernel = torch.exp(-0.5 * (offsets / sigma) ** 2) + _cache[key] = (kernel / kernel.sum(), radius) + return _cache[key] diff --git a/src/quantem/imaging/drift/core/nonrigid.py b/src/quantem/imaging/drift/core/nonrigid.py index cd701fed..605cf4f4 100644 --- a/src/quantem/imaging/drift/core/nonrigid.py +++ b/src/quantem/imaging/drift/core/nonrigid.py @@ -5,9 +5,7 @@ from scipy.optimize import minimize from tqdm import tqdm -from quantem.imaging.drift_utils import ( - gaussian_smooth_1d, -) +from quantem.imaging.drift.core.knots import gaussian_smooth_1d def align_nonrigid( diff --git a/src/quantem/imaging/drift/core/strip.py b/src/quantem/imaging/drift/core/strip.py new file mode 100644 index 00000000..0c644ead --- /dev/null +++ b/src/quantem/imaging/drift/core/strip.py @@ -0,0 +1 @@ +"""Package boundary for the strip-correction stage added in a later PR.""" diff --git a/src/quantem/imaging/drift/core/warping.py b/src/quantem/imaging/drift/core/warping.py index 5d8017cf..64c2b7df 100644 --- a/src/quantem/imaging/drift/core/warping.py +++ b/src/quantem/imaging/drift/core/warping.py @@ -1,196 +1,357 @@ -import warnings +"""Cross-correlation and translation refinement for drift correction.""" + +import math -import numpy as np import torch -from numpy.typing import NDArray -from scipy.interpolate import interp1d - -from quantem.core.utils.imaging_utils import ( - bilinear_kde, -) - - -class DriftInterpolator: - def __init__( - self, - input_shape, - output_shape, - scan_fast, - scan_slow, - pad_value, - kde_sigma, - ): - self.input_shape = input_shape - self.output_shape = output_shape - self.scan_fast = scan_fast - self.scan_slow = scan_slow - self.pad_value = pad_value - self.kde_sigma = kde_sigma - - self.rows_input = np.arange(input_shape[0]) - self.cols_input = np.arange(input_shape[1]) - self.u = np.linspace(0, 1, input_shape[1]) - - def transform_rows( - self, - knots_row: NDArray, - ): - num_knots = knots_row.shape[-1] - basis = np.linspace(0, 1, num_knots) - - if num_knots == 1: - xa = knots_row[0] + self.u[None, :] * self.scan_fast[0] * (self.input_shape[0] - 1) - ya = knots_row[1] + self.u[None, :] * self.scan_fast[1] * (self.input_shape[1] - 1) - elif num_knots == 2: - xa = interp1d(basis, knots_row[0], kind="linear", assume_sorted=True)(self.u) - ya = interp1d(basis, knots_row[1], kind="linear", assume_sorted=True)(self.u) - else: - kind = "quadratic" if num_knots == 3 else "cubic" - xa = interp1d( - basis, - knots_row[0], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - ya = interp1d( - basis, - knots_row[1], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - - return xa, ya - - def transform_coordinates( - self, - knots: NDArray, - ): - num_knots = knots.shape[-1] - - if num_knots == 1: - # vectorized version for speed - xa, ya = self.transform_rows(knots) - else: - xa = np.zeros(self.input_shape) - ya = np.zeros(self.input_shape) - for i in range(self.input_shape[0]): - xa[i], ya[i] = self.transform_rows(knots[:, i]) - - return xa, ya - - def warp_image( - self, - image: NDArray, - knots: NDArray, # shape: (2, rows, num_knots) - kde_sigma=None, - output_shape=None, - pad_value=None, - upsample_factor=None, - ) -> NDArray: - xa, ya = self.transform_coordinates( - knots, - ) +from torch.fft import fft2, fftfreq, ifft2, ifftshift - if kde_sigma is None: - kde_sigma = self.kde_sigma - if output_shape is None: - output_shape = self.output_shape +def cross_corr_batch( + ref_images: torch.Tensor, + mov_images: torch.Tensor, + upsample_factor: int, + max_shift_mask: torch.Tensor | None = None, + freq_grids: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Score test drift vectors by cross-correlation alignment cost. - if pad_value is None: - pad_value = self.pad_value + Core cost function of the affine grid search. For each test drift, + measures how well the warped image pairs align after sub-pixel + translation correction. Without this, the grid search has no way + to rank test drifts - it is the signal that drives drift estimation. - if upsample_factor is None: - upsample_factor = 1.0 + Pipeline: FFT cross-correlation → parabolic coarse peak → + DFT upsample for sub-pixel refinement → Fourier-domain shift → + MAE between reference and aligned image. - image_interp, weight_interp = bilinear_kde( - xa=xa * upsample_factor, # rows - ya=ya * upsample_factor, # cols - values=image, - output_shape=np.round(np.array(output_shape) * upsample_factor).astype("int"), - kde_sigma=kde_sigma * upsample_factor, - pad_value=pad_value, - return_pix_count=True, - ) - - return image_interp, weight_interp + Parameters + ---------- + ref_images : torch.Tensor + Reference images, shape ``(N, num_rows, num_cols)``. + mov_images : torch.Tensor + Images to align, shape ``(N, num_rows, num_cols)``. + upsample_factor : int + Sub-pixel upsampling factor for DFT refinement. + max_shift_mask : torch.Tensor or None + Precomputed boolean mask, shape ``(num_rows, num_cols)``. + True where correlation peaks should be zeroed (beyond max shift). + freq_grids : tuple[torch.Tensor, torch.Tensor] or None, optional + Precomputed ``(freq_row, freq_col)`` from ``torch.fft.fftfreq``, + shapes ``(num_rows, 1)`` and ``(1, num_cols)``. Avoids + recomputing the same grids each call. Default is None. + Returns + ------- + torch.Tensor + MAE cost per pair, shape ``(N,)``. -def bounded_sine_sigmoid(x, midpoint=0.5, width=1.0): + Examples + -------- + >>> ref = torch.randn(5, 64, 64, dtype=torch.float64) + >>> mov = torch.randn(5, 64, 64, dtype=torch.float64) + >>> cost = cross_corr_batch(ref, mov, 8, 32.0) + >>> cost.shape + torch.Size([5]) """ - Piecewise bounded sigmoid: zero, raised sine squared, one. + num_test_drifts, num_rows, num_cols = ref_images.shape + dtype = ref_images.dtype + mov_fft = fft2(mov_images) + cross_corr_fft = fft2(ref_images) * mov_fft.conj() + cross_corr = ifft2(cross_corr_fft).real + # Reject correlation peaks beyond max shift to avoid locking onto + # periodic lattice repeats or noise peaks far from the true shift + if max_shift_mask is not None: + cross_corr.masked_fill_(max_shift_mask[None], 0.0) + # Find best-matching shift: integer peak → parabola (~0.1 px) → DFT (~0.01 px) + peak_flat_idx = cross_corr.flatten(1).argmax(dim=1) + peak_row = peak_flat_idx // num_cols + peak_col = peak_flat_idx % num_cols + batch_idx = torch.arange(num_test_drifts, device=ref_images.device) + refined_row, refined_col = _parabolic_peak_2d( + cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx + ) + image_shifts = _dft_refine_shifts( + cross_corr_fft, refined_row, refined_col, upsample_factor + ) + # Wrap from [0, N) to [-N/2, N/2) so shifts represent actual displacement + image_shifts[:, 0] = ((image_shifts[:, 0] + num_rows / 2) % num_rows) - num_rows / 2 + image_shifts[:, 1] = ((image_shifts[:, 1] + num_cols / 2) % num_cols) - num_cols / 2 + if freq_grids is not None: + freq_row, freq_col = freq_grids + else: + freq_row = fftfreq(num_rows, device=ref_images.device, dtype=dtype)[:, None] + freq_col = fftfreq(num_cols, device=ref_images.device, dtype=dtype)[None, :] + phase = -2j * math.pi * ( + freq_row[None] * image_shifts[:, 0, None, None] + + freq_col[None] * image_shifts[:, 1, None, None] + ) + aligned_images = ifft2(mov_fft * torch.exp(phase)).real + return torch.mean(torch.abs(ref_images - aligned_images), dim=(1, 2)) + + +def translate_align( + warped_images: torch.Tensor, + upsample_factor: int, + max_image_shift: float | None, +) -> torch.Tensor: + """Pairwise translation alignment of warped images via cross-correlation. + + Called by ``_warp_and_translate_torch`` after each affine warp to + remove residual translational misalignment between the image pair. + Without this step, the merged image would be blurred by the remaining + translation offset even after the affine drift is corrected. + + Starting from image 0 as reference, sequentially aligns each image + using FFT cross-correlation with parabolic + DFT sub-pixel refinement. + The reference is updated as a running Fourier-domain average. Parameters ---------- - x : array-like, shape (...,) - Input values in [0, 1]. - midpoint : float - Center of the sigmoid transition. - width : float - Width of the sigmoid (range over which it ramps from 0 to 1). + warped_images : torch.Tensor + Warped images, shape ``(num_images, num_rows, num_cols)``. + upsample_factor : int + Sub-pixel precision (1/N pixel) for DFT refinement. + max_image_shift : float or None + Maximum allowed shift in pixels. Peaks beyond this radius are masked. + Returns ------- - y : array-like - Output in [0, 1], same shape as x. + torch.Tensor + Zero-mean shifts, shape ``(num_images, 2)`` in (row, col) order. """ - x = np.asarray(x) - # Truncate width if midpoint too close to edge - left_max = midpoint - width / 2 - right_min = midpoint + width / 2 - if left_max < 0: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * midpoint}.", - RuntimeWarning, + num_images, num_rows, num_cols = warped_images.shape + dtype = warped_images.dtype + device = warped_images.device + image_shifts = torch.zeros(num_images, 2, dtype=dtype, device=device) + ref_fft = fft2(warped_images[0]) + # Reject bad correlation peaks from noise or periodicity + # by zeroing everything beyond max_image_shift pixels from origin + shift_mask = None + if max_image_shift is not None: + dist_row = fftfreq(num_rows, 1.0 / num_rows, device=device, dtype=dtype) + dist_col = fftfreq(num_cols, 1.0 / num_cols, device=device, dtype=dtype) + shift_mask = dist_row[:, None] ** 2 + dist_col[None, :] ** 2 >= max_image_shift ** 2 + freq_row = fftfreq(num_rows, device=device, dtype=dtype)[:, None] + freq_col = fftfreq(num_cols, device=device, dtype=dtype)[None, :] + for img_idx in range(1, num_images): + mov_fft = fft2(warped_images[img_idx]) + cross_corr_fft = ref_fft * mov_fft.conj() + cross_corr = ifft2(cross_corr_fft).real + if shift_mask is not None: + cross_corr.masked_fill_(shift_mask, 0.0) + # Find the integer peak, refine with parabola to ~0.1 px + peak_flat_idx = cross_corr.flatten().argmax() + peak_row = peak_flat_idx[None] // num_cols + peak_col = peak_flat_idx[None] % num_cols + batch_idx = torch.zeros(1, dtype=torch.long, device=device) + refined_row, refined_col = _parabolic_peak_2d( + cross_corr[None], peak_row, peak_col, num_rows, num_cols, batch_idx ) - width = 2 * midpoint - - if right_min > 1: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * (1 - midpoint)}.", - RuntimeWarning, + # Zoom into peak neighborhood with DFT to get ~0.01 px precision, + # then wrap from [0, N) to centered [-N/2, N/2) convention + refined_shift = _dft_refine_shifts( + cross_corr_fft[None], refined_row, refined_col, upsample_factor ) - width = 2 * (1 - midpoint) - # Recalculate edges - left = midpoint - width / 2 - right = midpoint + width / 2 - - y = np.zeros_like(x, dtype=float) - in_band = (x >= left) & (x <= right) - # Map [left, right] to [0, pi/2] - t = (x[in_band] - left) / width # goes from 0 to 1 - y[in_band] = np.sin(t * np.pi / 2) ** 2 - y[x > right] = 1.0 - return y - - -def _bounded_sine_sigmoid_torch( - x: torch.Tensor, - midpoint: float = 0.5, - width: float = 1.0, -) -> torch.Tensor: - width = min(width, 2 * midpoint, 2 * (1 - midpoint)) - left = midpoint - width / 2 - right = midpoint + width / 2 - t = ((x - left) / width).clamp(0.0, 1.0) - return torch.where(x > right, torch.ones_like(x), torch.sin(t * (np.pi / 2)) ** 2) + image_shifts[img_idx, 0] = ((refined_shift[0, 0] + num_rows / 2) % num_rows) - num_rows / 2 + image_shifts[img_idx, 1] = ((refined_shift[0, 1] + num_cols / 2) % num_cols) - num_cols / 2 + # Apply the recovered shift to current image via Fourier shift theorem, + # then blend into running average so later images align to the cumulative mean + phase = torch.exp( + -2j * math.pi * ( + freq_row * image_shifts[img_idx, 0] + freq_col * image_shifts[img_idx, 1] + ) + ) + ref_fft = ref_fft * img_idx / (img_idx + 1) + mov_fft * phase / (img_idx + 1) + # Remove mean so shifts are relative (no absolute reference frame) + image_shifts -= image_shifts.mean(dim=0) + return image_shifts + + +def _dft_refine_shifts( + cross_corr_fft, + peak_row, + peak_col, + upsample_factor, +): + """Refine coarse sub-pixel shifts using DFT upsampling + parabolic fit. + + After ``_parabolic_peak_2d`` gives a coarse sub-pixel position, this + function zooms into a small neighborhood via the matrix-multiply DFT + and applies a second parabolic refinement on the upsampled patch. + The result is sub-pixel shifts with ``1 / upsample_factor`` precision. + + Without this step, shifts would have only ~0.1 px precision from + parabolic fitting alone. With ``upsample_factor=8``, precision + improves to ~0.01 px. + + Parameters + ---------- + cross_corr_fft : torch.Tensor + Complex cross-correlation in Fourier domain, ``(N, num_rows, num_cols)``. + peak_row, peak_col : torch.Tensor + Coarse sub-pixel peak positions in [0, N) from ``_parabolic_peak_2d``. + upsample_factor : int + Sub-pixel precision factor. + + Returns + ------- + image_shifts : torch.Tensor + Sub-pixel shifts in [0, N) coordinates, shape ``(N, 2)``. + """ + num_test_drifts = cross_corr_fft.shape[0] + dtype = peak_row.dtype + batch_idx = torch.arange(num_test_drifts, device=cross_corr_fft.device) + # Evaluate the correlation surface at 1/upsample_factor pixel spacing + # in a small window around each coarse peak - gives actual values, + # not the parabolic approximation from step 1 + upsampled_corr = _dft_upsample_batch( + cross_corr_fft, upsample_factor, torch.stack([peak_row, peak_col], dim=1) + ) + upsample_size = upsampled_corr.shape[1] + peak_flat_idx = upsampled_corr.flatten(1).argmax(dim=1) + local_row = peak_flat_idx // upsample_size + local_col = peak_flat_idx % upsample_size + # Final parabolic fit on the dense grid for last fraction of precision. + # Peaks at the edge of the upsampled window can't use the 3-point stencil + # (no neighbor on one side), so those are masked and kept at integer position + can_refine = ( + (local_row >= 1) + & (local_row < upsample_size - 1) + & (local_col >= 1) + & (local_col < upsample_size - 1) + ) + peak_val = upsampled_corr[batch_idx, local_row, local_col] + d_row_fine = _parabolic_sub_pixel( + upsampled_corr[batch_idx, (local_row - 1).clamp(min=0), local_col], + peak_val, + upsampled_corr[batch_idx, (local_row + 1).clamp(max=upsample_size - 1), local_col], + mask=can_refine, + ) + d_col_fine = _parabolic_sub_pixel( + upsampled_corr[batch_idx, local_row, (local_col - 1).clamp(min=0)], + peak_val, + upsampled_corr[batch_idx, local_row, (local_col + 1).clamp(max=upsample_size - 1)], + mask=can_refine, + ) + # Convert upsampled-grid position back to image-pixel coordinates: + # patch center is at index patch_radius in the upsampled grid, + # so (local_row - patch_radius) / upsample_factor = offset from coarse peak + patch_radius = math.ceil(1.5 * upsample_factor) + image_shifts = torch.zeros(num_test_drifts, 2, dtype=dtype, device=cross_corr_fft.device) + # local_row/col are int from argmax - cast to float for sub-pixel arithmetic + image_shifts[:, 0] = peak_row + (local_row.to(dtype) - patch_radius + d_row_fine) / upsample_factor + image_shifts[:, 1] = peak_col + (local_col.to(dtype) - patch_radius + d_col_fine) / upsample_factor + return image_shifts -def _fourier_crop_torch( - fft_array: torch.Tensor, - crop_shape: tuple[int, int], +def _dft_upsample_batch( + cross_corr_fft: torch.Tensor, + upsample_factor: int, + peak_positions: torch.Tensor, ) -> torch.Tensor: - crop_h, crop_w = crop_shape - h1 = crop_h // 2 - h2 = crop_h - h1 - w1 = crop_w // 2 - w2 = crop_w - w1 - result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) - result[:h1, :w1] = fft_array[:h1, :w1] - result[:h1, -w2:] = fft_array[:h1, -w2:] - result[-h2:, :w1] = fft_array[-h2:, :w1] - result[-h2:, -w2:] = fft_array[-h2:, -w2:] - return result + """Sub-pixel peak refinement for all test drifts in one pass. + + After the coarse FFT cross-correlation finds integer-pixel peaks, + zooms into a small neighborhood using the Guizar-Sicairos + matrix-multiply DFT. Without DFT upsampling, shift precision is + limited to ~0.1 px from parabolic fitting alone. With + ``upsample_factor=8``, precision improves to ~0.01 px. + + Parameters + ---------- + cross_corr_fft : torch.Tensor + Complex 2D cross-correlation in Fourier domain, shape ``(N, num_rows, num_cols)``. + upsample_factor : int + Upsampling factor (typically 8). + peak_positions : torch.Tensor + Coarse peak locations ``(row, col)`` per test drift, shape ``(N, 2)``. + + Returns + ------- + torch.Tensor + Real-valued upsampled correlation neighborhoods, shape ``(N, P, P)`` + where ``P = 2 * ceil(1.5 * upsample_factor) + 1``. + + """ + num_test_drifts, num_rows, num_cols = cross_corr_fft.shape + real_dtype = torch.float32 if cross_corr_fft.dtype == torch.complex64 else torch.float64 + # 1.5x radius ensures the patch captures the true peak after parabolic shift + patch_radius = math.ceil(1.5 * upsample_factor) + # Upsampled grid positions centered at zero: [-radius, ..., 0, ..., +radius] + upsample_grid = torch.arange(-patch_radius, patch_radius + 1, dtype=real_dtype, device=cross_corr_fft.device) + # ifftshift reorders [0,1,...,N-1] to match FFT output ordering, + # then subtract N//2 to center at zero + freq_row_base = ifftshift( + torch.arange(num_rows, dtype=real_dtype, device=cross_corr_fft.device) + ) - num_rows // 2 + freq_col_base = ifftshift( + torch.arange(num_cols, dtype=real_dtype, device=cross_corr_fft.device) + ) - num_cols // 2 + freq_row = freq_row_base[None, :] + (peak_positions[:, 0] - num_rows // 2)[:, None] + freq_col = freq_col_base[None, :] + (peak_positions[:, 1] - num_cols // 2)[:, None] + # Guizar-Sicairos matrix-multiply DFT: K_row @ CC @ K_col + kern_row = torch.exp( + -2j * math.pi / (num_rows * upsample_factor) + * upsample_grid[None, :, None] * freq_row[:, None, :] + ).to(cross_corr_fft.dtype) # real → complex for matrix multiply + kern_col = torch.exp( + -2j * math.pi / (num_cols * upsample_factor) + * freq_col[:, :, None] * upsample_grid[None, None, :] + ).to(cross_corr_fft.dtype) # real → complex for matrix multiply + # (N,P,M) @ (N,M,K) @ (N,K,P) -> (N,P,P) + return (kern_row @ cross_corr_fft @ kern_col).real + + +def _parabolic_peak_2d(cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx): + """Refine an integer cross-correlation peak to sub-pixel precision. + + Extracts the 3-point stencil along each axis and fits a parabola. + Without this, the DFT upsample window would be centered on the + integer peak which may be up to 0.5 px away from the true peak, + causing the upsampled patch to miss the true maximum. + + Parameters + ---------- + cross_corr : torch.Tensor + Batched correlation map, shape ``(N, num_rows, num_cols)``. + peak_row, peak_col : torch.Tensor + Integer peak positions, shape ``(N,)``. + num_rows, num_cols : int + Dimensions for periodic wrapping. + batch_idx : torch.Tensor + Batch indices, ``torch.arange(N)``. + + Returns + ------- + refined_row, refined_col : torch.Tensor + Sub-pixel peak positions in [0, N) coordinates. + """ + dtype = cross_corr.dtype + val_center = cross_corr[batch_idx, peak_row, peak_col] + val_row_m1 = cross_corr[batch_idx, (peak_row - 1) % num_rows, peak_col] + val_row_p1 = cross_corr[batch_idx, (peak_row + 1) % num_rows, peak_col] + val_col_m1 = cross_corr[batch_idx, peak_row, (peak_col - 1) % num_cols] + val_col_p1 = cross_corr[batch_idx, peak_row, (peak_col + 1) % num_cols] + # peak_row/col are int from argmax - cast to float for sub-pixel addition. + # Double modulo handles tiny negative offsets from float32 rounding + # that would otherwise wrap to N instead of 0 (e.g. -4e-8 % 64 = 64.0) + refined_row = ((peak_row.to(dtype) + _parabolic_sub_pixel(val_row_m1, val_center, val_row_p1)) % num_rows) % num_rows + refined_col = ((peak_col.to(dtype) + _parabolic_sub_pixel(val_col_m1, val_center, val_col_p1)) % num_cols) % num_cols + return refined_row, refined_col + + +def _parabolic_sub_pixel(val_m1, val_0, val_p1, mask=None): + """Sub-pixel offset from a 3-point stencil via parabolic interpolation. + + Cross-correlation peaks fall on integer pixel positions, but the true + shift is usually between pixels. Fitting a parabola through the peak + and its two neighbors gives ~0.1 px precision cheaply: + ``offset = (val_p1 - val_m1) / (4·val_0 - 2·val_p1 - 2·val_m1)``. + Without this, the DFT upsample window may be centered on the wrong + pixel and miss the true peak. + """ + denom = 4 * val_0 - 2 * val_p1 - 2 * val_m1 + valid = denom != 0 + if mask is not None: + valid = valid & mask + return torch.where(valid, (val_p1 - val_m1) / denom, torch.zeros_like(denom)) diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index c75f1cda..eefcda3a 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -95,9 +95,9 @@ class DriftCorrection(AutoSerialize): - **Batched grid search**: all ~97 candidate drift vectors are warped and scored in parallel, instead of one-at-a-time in a Python loop. - - **Batched bilinear KDE** (``drift_utils.bilinear_kde_batch``): + - **Batched bilinear KDE** (``core.knots.bilinear_kde_batch``): scatter-based image warping via ``scatter_add_`` with int32 indices. - - **Batched FFT cross-correlation** (``drift_utils.cross_corr_batch``): + - **Batched FFT cross-correlation** (``core.warping.cross_corr_batch``): sub-pixel alignment using DFT upsampling across all candidates at once. - **Zero CPU round-trips**: coordinate transforms, Gaussian smoothing, translation alignment, and error computation all stay on GPU until diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py index 445b6654..ad0a76e4 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preparation.py @@ -9,8 +9,8 @@ from quantem.core.utils.imaging_utils import ( cross_correlation_shift, ) -from quantem.imaging.drift.core.warping import DriftInterpolator -from quantem.imaging.drift_utils import ( +from quantem.imaging.drift.core.knots import ( + DriftInterpolator, bilinear_kde_batch, transform_coordinates_single_knot, ) diff --git a/src/quantem/imaging/drift_utils.py b/src/quantem/imaging/drift_utils.py deleted file mode 100644 index bd6280ae..00000000 --- a/src/quantem/imaging/drift_utils.py +++ /dev/null @@ -1,664 +0,0 @@ -"""Torch helper functions for drift correction.""" - -import math - -import torch -from torch.fft import fft2, fftfreq, ifft2, ifftshift - - -# --------------------------------------------------------------------------- -# Public API - called by DriftCorrection in drift.py -# --------------------------------------------------------------------------- - - -def bilinear_kde_batch( - row_coords: torch.Tensor, - col_coords: torch.Tensor, - source_image: torch.Tensor, - output_shape: tuple[int, int], - kde_sigma: float, - pad_value: float | torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Batched bilinear KDE: scatter N source images onto an output canvas. - - Each pixel scatters its value to its 4 nearest grid neighbors with - bilinear weights ``(1-dr)·(1-dc)``, ``dr·(1-dc)``, ``(1-dr)·dc``, - ``dr·dc`` where ``dr, dc`` are fractional row/col distances. - Accumulated counts and values are Gaussian-smoothed, then normalized: - ``output = pad_value·(1-coverage) + coverage·(values/counts)``. - - Used by both the affine grid search (N = candidate drift vectors, - single source image broadcast across drifts) and the nonrigid loop - (N = stacked source images, one per drift). - - Parameters - ---------- - row_coords : torch.Tensor - Row coordinates of input pixels, shape ``(N, rows, cols)``. - col_coords : torch.Tensor - Column coordinates of input pixels, shape ``(N, rows, cols)``. - source_image : torch.Tensor - Pixel values to scatter. Either ``(rows, cols)`` (same image used - for all N drifts - affine grid search) or ``(N, rows, cols)`` - (different image per drift - multi-image batched warping). - output_shape : tuple[int, int] - Canvas size ``(num_rows, num_cols)`` for the output images. - kde_sigma : float - Gaussian smoothing sigma in pixels. - pad_value : float or torch.Tensor - Fill value where pixel coverage is below threshold. If a tensor of - shape ``(N,)``, applies a different pad value per drift. - - Returns - ------- - tuple[torch.Tensor, torch.Tensor] - ``(warped_images, sum_weights)`` - warped images and smoothed pixel coverage, - both shape ``(N, num_rows, num_cols)``. - """ - num_test_drifts = row_coords.shape[0] - num_rows, num_cols = output_shape - coverage_threshold = 1e-3 - # Flatten spatial dims - scatter_add_ works on 1D buffers - row_flat = row_coords.flatten(1) - col_flat = col_coords.flatten(1) - # Stay in float for fractional distance, convert to int only for scatter indices - row_floor = row_flat.floor() - col_floor = col_flat.floor() - frac_row = row_flat - row_floor - frac_col = col_flat - col_floor - row_floor = row_floor.int() - col_floor = col_floor.int() - if source_image.dim() == 3: - # Per-drift source images: each drift scatters its own pixel values - source_values_flat = source_image.flatten() - else: - source_values_flat = source_image.flatten().repeat(num_test_drifts) - # All N batch entries scatter into one flat buffer - offset separates them - batch_offsets = ( - torch.arange(num_test_drifts, device=row_coords.device, dtype=torch.int32) - * num_rows * num_cols - )[:, None] - # Float32 accumulators - scatter_add_ requires source dtype to match, - # so all input tensors must be float32 (raises on float64). - sum_weights = torch.zeros( - num_test_drifts * num_rows * num_cols, dtype=torch.float32, device=row_coords.device - ) - sum_values = torch.zeros_like(sum_weights) - # Periodic wrapping so pixels near edges scatter to the opposite side - row_wrapped = row_floor % num_rows - col_wrapped = col_floor % num_cols - row_next = (row_wrapped + 1) % num_rows - col_next = (col_wrapped + 1) % num_cols - # Each pixel distributes its value to the 4 nearest grid neighbors - # weighted by bilinear distance: (1-dr)(1-dc), dr(1-dc), (1-dr)dc, dr·dc - for corner_row, corner_col, corner_weight in [ - (row_wrapped, col_wrapped, ((1 - frac_row) * (1 - frac_col)).flatten()), - (row_next, col_wrapped, (frac_row * (1 - frac_col)).flatten()), - (row_wrapped, col_next, ((1 - frac_row) * frac_col).flatten()), - (row_next, col_next, (frac_row * frac_col).flatten()), - ]: - flat_indices = (corner_row * num_cols + corner_col + batch_offsets).flatten() - sum_weights.scatter_add_(0, flat_indices, corner_weight) - sum_values.scatter_add_(0, flat_indices, corner_weight * source_values_flat) - sum_weights = sum_weights.reshape(num_test_drifts, num_rows, num_cols) - sum_values = sum_values.reshape(num_test_drifts, num_rows, num_cols) - # Smooth the scattered counts and values to fill gaps between pixels - sum_weights = gaussian_smooth_batch(sum_weights, kde_sigma) - sum_values = gaussian_smooth_batch(sum_values, kde_sigma) - # Blend between pad_value (uncovered) and normalized values (covered), - # ramping linearly with coverage to avoid hard edges at the boundary - coverage_weight = torch.clamp(sum_weights / coverage_threshold, max=1.0) - if isinstance(pad_value, torch.Tensor) and pad_value.dim() == 1: - # Per-drift pad value: reshape (N,) → (N, 1, 1) for broadcasting - pad_value = pad_value[:, None, None] - warped_images = pad_value * (1 - coverage_weight) + coverage_weight * ( - sum_values / torch.clamp(sum_weights, min=1e-8) - ) - return warped_images, sum_weights - - -def cross_corr_batch( - ref_images: torch.Tensor, - mov_images: torch.Tensor, - upsample_factor: int, - max_shift_mask: torch.Tensor | None = None, - freq_grids: tuple[torch.Tensor, torch.Tensor] | None = None, -) -> torch.Tensor: - """Score test drift vectors by cross-correlation alignment cost. - - Core cost function of the affine grid search. For each test drift, - measures how well the warped image pairs align after sub-pixel - translation correction. Without this, the grid search has no way - to rank test drifts - it is the signal that drives drift estimation. - - Pipeline: FFT cross-correlation → parabolic coarse peak → - DFT upsample for sub-pixel refinement → Fourier-domain shift → - MAE between reference and aligned image. - - Parameters - ---------- - ref_images : torch.Tensor - Reference images, shape ``(N, num_rows, num_cols)``. - mov_images : torch.Tensor - Images to align, shape ``(N, num_rows, num_cols)``. - upsample_factor : int - Sub-pixel upsampling factor for DFT refinement. - max_shift_mask : torch.Tensor or None - Precomputed boolean mask, shape ``(num_rows, num_cols)``. - True where correlation peaks should be zeroed (beyond max shift). - freq_grids : tuple[torch.Tensor, torch.Tensor] or None, optional - Precomputed ``(freq_row, freq_col)`` from ``torch.fft.fftfreq``, - shapes ``(num_rows, 1)`` and ``(1, num_cols)``. Avoids - recomputing the same grids each call. Default is None. - - Returns - ------- - torch.Tensor - MAE cost per pair, shape ``(N,)``. - - Examples - -------- - >>> ref = torch.randn(5, 64, 64, dtype=torch.float64) - >>> mov = torch.randn(5, 64, 64, dtype=torch.float64) - >>> cost = cross_corr_batch(ref, mov, 8, 32.0) - >>> cost.shape - torch.Size([5]) - """ - num_test_drifts, num_rows, num_cols = ref_images.shape - dtype = ref_images.dtype - mov_fft = fft2(mov_images) - cross_corr_fft = fft2(ref_images) * mov_fft.conj() - cross_corr = ifft2(cross_corr_fft).real - # Reject correlation peaks beyond max shift to avoid locking onto - # periodic lattice repeats or noise peaks far from the true shift - if max_shift_mask is not None: - cross_corr.masked_fill_(max_shift_mask[None], 0.0) - # Find best-matching shift: integer peak → parabola (~0.1 px) → DFT (~0.01 px) - peak_flat_idx = cross_corr.flatten(1).argmax(dim=1) - peak_row = peak_flat_idx // num_cols - peak_col = peak_flat_idx % num_cols - batch_idx = torch.arange(num_test_drifts, device=ref_images.device) - refined_row, refined_col = _parabolic_peak_2d( - cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx - ) - image_shifts = _dft_refine_shifts( - cross_corr_fft, refined_row, refined_col, upsample_factor - ) - # Wrap from [0, N) to [-N/2, N/2) so shifts represent actual displacement - image_shifts[:, 0] = ((image_shifts[:, 0] + num_rows / 2) % num_rows) - num_rows / 2 - image_shifts[:, 1] = ((image_shifts[:, 1] + num_cols / 2) % num_cols) - num_cols / 2 - if freq_grids is not None: - freq_row, freq_col = freq_grids - else: - freq_row = fftfreq(num_rows, device=ref_images.device, dtype=dtype)[:, None] - freq_col = fftfreq(num_cols, device=ref_images.device, dtype=dtype)[None, :] - phase = -2j * math.pi * ( - freq_row[None] * image_shifts[:, 0, None, None] - + freq_col[None] * image_shifts[:, 1, None, None] - ) - aligned_images = ifft2(mov_fft * torch.exp(phase)).real - return torch.mean(torch.abs(ref_images - aligned_images), dim=(1, 2)) - - -def translate_align( - warped_images: torch.Tensor, - upsample_factor: int, - max_image_shift: float | None, -) -> torch.Tensor: - """Pairwise translation alignment of warped images via cross-correlation. - - Called by ``_warp_and_translate_torch`` after each affine warp to - remove residual translational misalignment between the image pair. - Without this step, the merged image would be blurred by the remaining - translation offset even after the affine drift is corrected. - - Starting from image 0 as reference, sequentially aligns each image - using FFT cross-correlation with parabolic + DFT sub-pixel refinement. - The reference is updated as a running Fourier-domain average. - - Parameters - ---------- - warped_images : torch.Tensor - Warped images, shape ``(num_images, num_rows, num_cols)``. - upsample_factor : int - Sub-pixel precision (1/N pixel) for DFT refinement. - max_image_shift : float or None - Maximum allowed shift in pixels. Peaks beyond this radius are masked. - - Returns - ------- - torch.Tensor - Zero-mean shifts, shape ``(num_images, 2)`` in (row, col) order. - """ - num_images, num_rows, num_cols = warped_images.shape - dtype = warped_images.dtype - device = warped_images.device - image_shifts = torch.zeros(num_images, 2, dtype=dtype, device=device) - ref_fft = fft2(warped_images[0]) - # Reject bad correlation peaks from noise or periodicity - # by zeroing everything beyond max_image_shift pixels from origin - shift_mask = None - if max_image_shift is not None: - dist_row = fftfreq(num_rows, 1.0 / num_rows, device=device, dtype=dtype) - dist_col = fftfreq(num_cols, 1.0 / num_cols, device=device, dtype=dtype) - shift_mask = dist_row[:, None] ** 2 + dist_col[None, :] ** 2 >= max_image_shift ** 2 - freq_row = fftfreq(num_rows, device=device, dtype=dtype)[:, None] - freq_col = fftfreq(num_cols, device=device, dtype=dtype)[None, :] - for img_idx in range(1, num_images): - mov_fft = fft2(warped_images[img_idx]) - cross_corr_fft = ref_fft * mov_fft.conj() - cross_corr = ifft2(cross_corr_fft).real - if shift_mask is not None: - cross_corr.masked_fill_(shift_mask, 0.0) - # Find the integer peak, refine with parabola to ~0.1 px - peak_flat_idx = cross_corr.flatten().argmax() - peak_row = peak_flat_idx[None] // num_cols - peak_col = peak_flat_idx[None] % num_cols - batch_idx = torch.zeros(1, dtype=torch.long, device=device) - refined_row, refined_col = _parabolic_peak_2d( - cross_corr[None], peak_row, peak_col, num_rows, num_cols, batch_idx - ) - # Zoom into peak neighborhood with DFT to get ~0.01 px precision, - # then wrap from [0, N) to centered [-N/2, N/2) convention - refined_shift = _dft_refine_shifts( - cross_corr_fft[None], refined_row, refined_col, upsample_factor - ) - image_shifts[img_idx, 0] = ((refined_shift[0, 0] + num_rows / 2) % num_rows) - num_rows / 2 - image_shifts[img_idx, 1] = ((refined_shift[0, 1] + num_cols / 2) % num_cols) - num_cols / 2 - # Apply the recovered shift to current image via Fourier shift theorem, - # then blend into running average so later images align to the cumulative mean - phase = torch.exp( - -2j * math.pi * ( - freq_row * image_shifts[img_idx, 0] + freq_col * image_shifts[img_idx, 1] - ) - ) - ref_fft = ref_fft * img_idx / (img_idx + 1) + mov_fft * phase / (img_idx + 1) - # Remove mean so shifts are relative (no absolute reference frame) - image_shifts -= image_shifts.mean(dim=0) - return image_shifts - - -def transform_coordinates_single_knot( - knots: torch.Tensor, - scan_fast: torch.Tensor, - input_shape: tuple[int, int], -) -> tuple[torch.Tensor, torch.Tensor]: - """Single-knot fast path: map source pixels to canvas coordinates. - - **Single-knot only.** Each scanline has exactly one (row, col) anchor; - the fast-scan-direction position is filled in by linear interpolation - along the scanline. Multi-knot Bezier interpolation is intentionally - not supported here - that's the scipy backend's job. The pytorch path - optimizes for the common single-knot case (≥95% of real STEM workflows). - - Called by ``preprocess``, ``_affine_grid_search_batch``, and - ``_warp_and_translate_torch`` to map source image pixels onto the - padded output canvas. Without this, the warped images would have - no spatial mapping and the grid search couldn't score test drifts. - - Each input row maps to a line on the canvas: - ``row = knot_row + fraction * scan_fast[0] * (num_rows - 1)`` - ``col = knot_col + fraction * scan_fast[1] * (num_cols - 1)`` - where row and col dimensions scale independently for non-square images. - - Parameters - ---------- - knots : torch.Tensor - Knot positions, shape ``(2, num_rows, 1)``. First dim is (row, col). - The trailing 1 is the single-knot dimension; multi-knot inputs are - rejected by the caller before reaching this function. - scan_fast : torch.Tensor - Fast scan direction vector, shape ``(2,)``. - input_shape : tuple[int, int] - Original image shape ``(num_rows, num_cols)``. - - Returns - ------- - row_coords : torch.Tensor - Row coordinates on canvas, shape ``(num_rows, num_cols)``. - col_coords : torch.Tensor - Column coordinates on canvas, shape ``(num_rows, num_cols)``. - - Examples - -------- - >>> knots = torch.zeros(2, 64, 1) - >>> scan_fast = torch.tensor([0.0, 1.0]) - >>> r, c = transform_coordinates_single_knot(knots, scan_fast, (64, 64)) - >>> r.shape - torch.Size([64, 64]) - """ - num_rows, num_cols = input_shape - fast_fraction = torch.linspace(0, 1, num_cols, dtype=knots.dtype, device=knots.device) - row_coords = knots[0, :, 0:1] + fast_fraction[None, :] * scan_fast[0] * (num_rows - 1) - col_coords = knots[1, :, 0:1] + fast_fraction[None, :] * scan_fast[1] * (num_cols - 1) - return row_coords, col_coords - - -def gaussian_smooth_batch( - field_stack: torch.Tensor, - sigma: float, -) -> torch.Tensor: - """Batched 2D Gaussian smoothing matching ``scipy.ndimage.gaussian_filter``. - - Used by ``bilinear_kde_batch`` to smooth scattered counts and - values before normalization. Without smoothing, the warped images - have salt-and-pepper artifacts from the scatter step. - - Parameters - ---------- - field_stack : torch.Tensor - Input tensor of shape ``(N, num_rows, num_cols)``. - sigma : float - Standard deviation of the Gaussian kernel in pixels. - - Returns - ------- - torch.Tensor - Smoothed tensor of shape ``(N, num_rows, num_cols)``. - - """ - kernel, radius = _gaussian_kernel_1d(sigma, field_stack.dtype, field_stack.device) - # Separable kernel: column pass then row pass to halve FLOPs vs full 2D conv - kernel_col = kernel[None, None, None, :] - kernel_row = kernel[None, None, :, None] - field_stack = field_stack[:, None] - field_stack = torch.nn.functional.conv2d(_symmetric_pad(field_stack, 0, radius), kernel_col) - field_stack = torch.nn.functional.conv2d(_symmetric_pad(field_stack, radius, 0), kernel_row) - return field_stack[:, 0] - - -def gaussian_smooth_1d( - signal: torch.Tensor, - sigma: float, -) -> torch.Tensor: - """1D Gaussian smoothing matching ``scipy.ndimage.gaussian_filter``. - - Smooths each row of the input independently using a separable 1D kernel. - Used for regularizing knot displacement vectors in the nonrigid loop, - where the signal is 1D (one value per scan line). - - Parameters - ---------- - signal : torch.Tensor - Input tensor of shape ``(N, L)`` - N channels, L samples. - sigma : float - Standard deviation of the Gaussian kernel in pixels. - - Returns - ------- - torch.Tensor - Smoothed tensor of shape ``(N, L)``. - """ - kernel, radius = _gaussian_kernel_1d(sigma, signal.dtype, signal.device) - signal_padded = _symmetric_pad_1d(signal[:, None], radius) - return torch.nn.functional.conv1d(signal_padded, kernel[None, None, :])[:, 0] - - -# --------------------------------------------------------------------------- -# Building blocks - used internally by the public API functions above -# --------------------------------------------------------------------------- - - -def _dft_refine_shifts( - cross_corr_fft, - peak_row, - peak_col, - upsample_factor, -): - """Refine coarse sub-pixel shifts using DFT upsampling + parabolic fit. - - After ``_parabolic_peak_2d`` gives a coarse sub-pixel position, this - function zooms into a small neighborhood via the matrix-multiply DFT - and applies a second parabolic refinement on the upsampled patch. - The result is sub-pixel shifts with ``1 / upsample_factor`` precision. - - Without this step, shifts would have only ~0.1 px precision from - parabolic fitting alone. With ``upsample_factor=8``, precision - improves to ~0.01 px. - - Parameters - ---------- - cross_corr_fft : torch.Tensor - Complex cross-correlation in Fourier domain, ``(N, num_rows, num_cols)``. - peak_row, peak_col : torch.Tensor - Coarse sub-pixel peak positions in [0, N) from ``_parabolic_peak_2d``. - upsample_factor : int - Sub-pixel precision factor. - - Returns - ------- - image_shifts : torch.Tensor - Sub-pixel shifts in [0, N) coordinates, shape ``(N, 2)``. - """ - num_test_drifts = cross_corr_fft.shape[0] - dtype = peak_row.dtype - batch_idx = torch.arange(num_test_drifts, device=cross_corr_fft.device) - # Evaluate the correlation surface at 1/upsample_factor pixel spacing - # in a small window around each coarse peak - gives actual values, - # not the parabolic approximation from step 1 - upsampled_corr = _dft_upsample_batch( - cross_corr_fft, upsample_factor, torch.stack([peak_row, peak_col], dim=1) - ) - upsample_size = upsampled_corr.shape[1] - peak_flat_idx = upsampled_corr.flatten(1).argmax(dim=1) - local_row = peak_flat_idx // upsample_size - local_col = peak_flat_idx % upsample_size - # Final parabolic fit on the dense grid for last fraction of precision. - # Peaks at the edge of the upsampled window can't use the 3-point stencil - # (no neighbor on one side), so those are masked and kept at integer position - can_refine = ( - (local_row >= 1) - & (local_row < upsample_size - 1) - & (local_col >= 1) - & (local_col < upsample_size - 1) - ) - peak_val = upsampled_corr[batch_idx, local_row, local_col] - d_row_fine = _parabolic_sub_pixel( - upsampled_corr[batch_idx, (local_row - 1).clamp(min=0), local_col], - peak_val, - upsampled_corr[batch_idx, (local_row + 1).clamp(max=upsample_size - 1), local_col], - mask=can_refine, - ) - d_col_fine = _parabolic_sub_pixel( - upsampled_corr[batch_idx, local_row, (local_col - 1).clamp(min=0)], - peak_val, - upsampled_corr[batch_idx, local_row, (local_col + 1).clamp(max=upsample_size - 1)], - mask=can_refine, - ) - # Convert upsampled-grid position back to image-pixel coordinates: - # patch center is at index patch_radius in the upsampled grid, - # so (local_row - patch_radius) / upsample_factor = offset from coarse peak - patch_radius = math.ceil(1.5 * upsample_factor) - image_shifts = torch.zeros(num_test_drifts, 2, dtype=dtype, device=cross_corr_fft.device) - # local_row/col are int from argmax - cast to float for sub-pixel arithmetic - image_shifts[:, 0] = peak_row + (local_row.to(dtype) - patch_radius + d_row_fine) / upsample_factor - image_shifts[:, 1] = peak_col + (local_col.to(dtype) - patch_radius + d_col_fine) / upsample_factor - return image_shifts - - -def _dft_upsample_batch( - cross_corr_fft: torch.Tensor, - upsample_factor: int, - peak_positions: torch.Tensor, -) -> torch.Tensor: - """Sub-pixel peak refinement for all test drifts in one pass. - - After the coarse FFT cross-correlation finds integer-pixel peaks, - zooms into a small neighborhood using the Guizar-Sicairos - matrix-multiply DFT. Without DFT upsampling, shift precision is - limited to ~0.1 px from parabolic fitting alone. With - ``upsample_factor=8``, precision improves to ~0.01 px. - - Parameters - ---------- - cross_corr_fft : torch.Tensor - Complex 2D cross-correlation in Fourier domain, shape ``(N, num_rows, num_cols)``. - upsample_factor : int - Upsampling factor (typically 8). - peak_positions : torch.Tensor - Coarse peak locations ``(row, col)`` per test drift, shape ``(N, 2)``. - - Returns - ------- - torch.Tensor - Real-valued upsampled correlation neighborhoods, shape ``(N, P, P)`` - where ``P = 2 * ceil(1.5 * upsample_factor) + 1``. - - """ - num_test_drifts, num_rows, num_cols = cross_corr_fft.shape - real_dtype = torch.float32 if cross_corr_fft.dtype == torch.complex64 else torch.float64 - # 1.5x radius ensures the patch captures the true peak after parabolic shift - patch_radius = math.ceil(1.5 * upsample_factor) - # Upsampled grid positions centered at zero: [-radius, ..., 0, ..., +radius] - upsample_grid = torch.arange(-patch_radius, patch_radius + 1, dtype=real_dtype, device=cross_corr_fft.device) - # ifftshift reorders [0,1,...,N-1] to match FFT output ordering, - # then subtract N//2 to center at zero - freq_row_base = ifftshift( - torch.arange(num_rows, dtype=real_dtype, device=cross_corr_fft.device) - ) - num_rows // 2 - freq_col_base = ifftshift( - torch.arange(num_cols, dtype=real_dtype, device=cross_corr_fft.device) - ) - num_cols // 2 - freq_row = freq_row_base[None, :] + (peak_positions[:, 0] - num_rows // 2)[:, None] - freq_col = freq_col_base[None, :] + (peak_positions[:, 1] - num_cols // 2)[:, None] - # Guizar-Sicairos matrix-multiply DFT: K_row @ CC @ K_col - kern_row = torch.exp( - -2j * math.pi / (num_rows * upsample_factor) - * upsample_grid[None, :, None] * freq_row[:, None, :] - ).to(cross_corr_fft.dtype) # real → complex for matrix multiply - kern_col = torch.exp( - -2j * math.pi / (num_cols * upsample_factor) - * freq_col[:, :, None] * upsample_grid[None, None, :] - ).to(cross_corr_fft.dtype) # real → complex for matrix multiply - # (N,P,M) @ (N,M,K) @ (N,K,P) -> (N,P,P) - return (kern_row @ cross_corr_fft @ kern_col).real - -# --------------------------------------------------------------------------- -# Primitives - lowest-level operations -# --------------------------------------------------------------------------- - - -def _parabolic_peak_2d(cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx): - """Refine an integer cross-correlation peak to sub-pixel precision. - - Extracts the 3-point stencil along each axis and fits a parabola. - Without this, the DFT upsample window would be centered on the - integer peak which may be up to 0.5 px away from the true peak, - causing the upsampled patch to miss the true maximum. - - Parameters - ---------- - cross_corr : torch.Tensor - Batched correlation map, shape ``(N, num_rows, num_cols)``. - peak_row, peak_col : torch.Tensor - Integer peak positions, shape ``(N,)``. - num_rows, num_cols : int - Dimensions for periodic wrapping. - batch_idx : torch.Tensor - Batch indices, ``torch.arange(N)``. - - Returns - ------- - refined_row, refined_col : torch.Tensor - Sub-pixel peak positions in [0, N) coordinates. - """ - dtype = cross_corr.dtype - val_center = cross_corr[batch_idx, peak_row, peak_col] - val_row_m1 = cross_corr[batch_idx, (peak_row - 1) % num_rows, peak_col] - val_row_p1 = cross_corr[batch_idx, (peak_row + 1) % num_rows, peak_col] - val_col_m1 = cross_corr[batch_idx, peak_row, (peak_col - 1) % num_cols] - val_col_p1 = cross_corr[batch_idx, peak_row, (peak_col + 1) % num_cols] - # peak_row/col are int from argmax - cast to float for sub-pixel addition. - # Double modulo handles tiny negative offsets from float32 rounding - # that would otherwise wrap to N instead of 0 (e.g. -4e-8 % 64 = 64.0) - refined_row = ((peak_row.to(dtype) + _parabolic_sub_pixel(val_row_m1, val_center, val_row_p1)) % num_rows) % num_rows - refined_col = ((peak_col.to(dtype) + _parabolic_sub_pixel(val_col_m1, val_center, val_col_p1)) % num_cols) % num_cols - return refined_row, refined_col - - -def _parabolic_sub_pixel(val_m1, val_0, val_p1, mask=None): - """Sub-pixel offset from a 3-point stencil via parabolic interpolation. - - Cross-correlation peaks fall on integer pixel positions, but the true - shift is usually between pixels. Fitting a parabola through the peak - and its two neighbors gives ~0.1 px precision cheaply: - ``offset = (val_p1 - val_m1) / (4·val_0 - 2·val_p1 - 2·val_m1)``. - Without this, the DFT upsample window may be centered on the wrong - pixel and miss the true peak. - """ - denom = 4 * val_0 - 2 * val_p1 - 2 * val_m1 - valid = denom != 0 - if mask is not None: - valid = valid & mask - return torch.where(valid, (val_p1 - val_m1) / denom, torch.zeros_like(denom)) - - -def _symmetric_pad_1d(signal: torch.Tensor, pad: int) -> torch.Tensor: - """Symmetric 1D padding matching scipy's reflect mode. - - Same edge-repeat semantics as ``_symmetric_pad`` but for 1D signals. - Used by ``gaussian_smooth_1d`` for regularization of knot vectors. - """ - left = signal[:, :, :pad].flip(-1) - right = signal[:, :, -pad:].flip(-1) - return torch.cat([left, signal, right], dim=-1) - - -def _symmetric_pad( - field_stack: torch.Tensor, - pad_rows: int, - pad_cols: int, -) -> torch.Tensor: - """Symmetric padding matching scipy's reflect mode for parity. - - Without this, the torch and numpy Gaussian smoothing paths produce - different results near canvas edges, breaking numerical parity. - - Scipy's ``mode='reflect'`` repeats the edge pixel - (``[1,2,3]`` → ``[2,1,1,2,3,3,2]``), but PyTorch's - ``F.pad(mode='reflect')`` does not (``[1,2,3]`` → ``[3,2,1,2,3,2,1]``). - - Parameters - ---------- - field_stack : torch.Tensor - Input tensor of shape ``(N, C, num_rows, num_cols)``. - pad_rows : int - Number of rows to pad on top and bottom. - pad_cols : int - Number of columns to pad on left and right. - - Returns - ------- - torch.Tensor - Padded tensor. - - Examples - -------- - >>> t = torch.tensor([[[[1., 2., 3.]]]]) - >>> _symmetric_pad(t, 0, 2) - tensor([[[[2., 1., 1., 2., 3., 3., 2.]]]]) - """ - if pad_cols > 0: - left = field_stack[:, :, :, :pad_cols].flip(-1) - right = field_stack[:, :, :, -pad_cols:].flip(-1) - field_stack = torch.cat([left, field_stack, right], dim=-1) - if pad_rows > 0: - top = field_stack[:, :, :pad_rows, :].flip(-2) - bottom = field_stack[:, :, -pad_rows:, :].flip(-2) - field_stack = torch.cat([top, field_stack, bottom], dim=-2) - return field_stack - - -def _gaussian_kernel_1d(sigma, dtype, device, _cache={}): - """Normalized 1D Gaussian ``exp(-0.5*(x/sigma)^2)``, radius ``4*sigma``. - - Cached via mutable default arg - the grid search calls this ~800 times - with the same sigma, saving ~44ms of redundant kernel construction. - """ - key = (sigma, dtype, device) - if key not in _cache: - radius = int(4 * sigma + 0.5) - offsets = torch.arange(-radius, radius + 1, dtype=dtype, device=device) - kernel = torch.exp(-0.5 * (offsets / sigma) ** 2) - _cache[key] = (kernel / kernel.sum(), radius) - return _cache[key] diff --git a/tests/imaging/test_drift_utils.py b/tests/imaging/test_drift_utils.py index a917ef17..bf540b4e 100644 --- a/tests/imaging/test_drift_utils.py +++ b/tests/imaging/test_drift_utils.py @@ -1,4 +1,4 @@ -"""Parity tests for drift_utils.py torch functions. +"""Parity tests for Drift Torch numerical helpers. Tests the core building blocks against numpy/scipy equivalents. The full pipeline is covered by frozen baselines in test_drift.py. @@ -10,18 +10,22 @@ from scipy.ndimage import gaussian_filter from quantem.core.utils.imaging_utils import bilinear_kde, fourier_cropping -from quantem.imaging.drift import _bounded_sine_sigmoid_torch, _fourier_crop_torch -from quantem.imaging.drift import bounded_sine_sigmoid -from quantem.imaging.drift_utils import ( - _parabolic_peak_2d, - _parabolic_sub_pixel, +from quantem.imaging.drift import ( + _bounded_sine_sigmoid_torch, + _fourier_crop_torch, + bounded_sine_sigmoid, +) +from quantem.imaging.drift.core.knots import ( _symmetric_pad, bilinear_kde_batch, - cross_corr_batch, gaussian_smooth_1d, gaussian_smooth_batch, ) - +from quantem.imaging.drift.core.warping import ( + _parabolic_peak_2d, + _parabolic_sub_pixel, + cross_corr_batch, +) # --------------------------------------------------------------------------- # High-level: cross-correlation and warping From ef0cd61e226eac884bb0a45eac4fb9ca352b6359 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Mon, 24 Aug 2026 23:27:50 -0700 Subject: [PATCH 04/11] feat: add regional diagnostics for affine drift correction (#15) * feat: add drift registration diagnostics * feat: add structured drift reports * feat: add static drift diagnostic plots * test: validate drift diagnostics on frozen results * fix: keep drift diagnostics consistent with fitted state --- src/quantem/imaging/drift/core/affine.py | 3 + src/quantem/imaging/drift/core/nonrigid.py | 3 + src/quantem/imaging/drift/diagnostics.py | 276 +++++++++++++++++++++ src/quantem/imaging/drift/plot.py | 160 ++++++++++++ src/quantem/imaging/drift/preparation.py | 3 + src/quantem/imaging/drift/report.py | 69 ++++++ tests/imaging/test_drift.py | 168 +++++++++++++ 7 files changed, 682 insertions(+) create mode 100644 src/quantem/imaging/drift/diagnostics.py create mode 100644 src/quantem/imaging/drift/report.py diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index 933956bc..82af9334 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -3,6 +3,7 @@ import torch from torch.fft import fftfreq +from quantem.imaging.drift import diagnostics from quantem.imaging.drift.core.knots import ( bilinear_kde_batch, transform_coordinates_single_knot, @@ -156,6 +157,8 @@ def _search_and_apply(candidates, label): print(f"Error: {err[0, 1]:.2f} -> {err[-1, 1]:.2f} " f"({(err[0, 1] - err[-1, 1]) / err[0, 1] * 100:+.1f}%)") + diagnostics._record_stage(self, "affine") + # Plots kwargs.pop("title", None) if show_merged: diff --git a/src/quantem/imaging/drift/core/nonrigid.py b/src/quantem/imaging/drift/core/nonrigid.py index 605cf4f4..7620af48 100644 --- a/src/quantem/imaging/drift/core/nonrigid.py +++ b/src/quantem/imaging/drift/core/nonrigid.py @@ -5,6 +5,7 @@ from scipy.optimize import minimize from tqdm import tqdm +from quantem.imaging.drift import diagnostics from quantem.imaging.drift.core.knots import gaussian_smooth_1d @@ -301,6 +302,8 @@ def align_nonrigid( warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor=8) self.calculate_error(2, _warped_t=warped_t) + diagnostics._record_stage(self, "nonrigid") + if show_merged: self.plot_merged_images( show_knots=show_knots, diff --git a/src/quantem/imaging/drift/diagnostics.py b/src/quantem/imaging/drift/diagnostics.py new file mode 100644 index 00000000..b4d52343 --- /dev/null +++ b/src/quantem/imaging/drift/diagnostics.py @@ -0,0 +1,276 @@ +"""Pure numerical diagnostics for drift-correction results. + +This module measures existing correction checkpoints. It does not run an +optimizer, update knots, or replace the current warped-image cache. Public API +decisions are intentionally deferred to the final API integration change. +""" + +from collections.abc import Sequence + +import numpy as np +import torch +from numpy.typing import NDArray + +_COVERAGE_THRESHOLD = 1e-3 +_MIN_CORRELATION_PIXELS = 64 + + +def _copy_knots(knots: Sequence[NDArray | torch.Tensor]) -> tuple[NDArray, ...]: + """Copy knot arrays to CPU without retaining mutable solver storage.""" + copied = [] + for knot in knots: + if isinstance(knot, torch.Tensor): + knot = knot.detach().cpu().numpy() + copied.append(np.asarray(knot).copy()) + return tuple(copied) + + +def _record_stage(correction, stage: str) -> None: + """Record one lightweight knot checkpoint after a completed stage.""" + if not hasattr(correction, "knots"): + raise RuntimeError( + "A diagnostic checkpoint requires prepared knots. Run preprocess() first." + ) + if not hasattr(correction, "_diagnostic_knots"): + correction._diagnostic_knots = {} + stage = str(stage) + recorded = list(correction._diagnostic_knots) + if stage in recorded: + for stale_stage in recorded[recorded.index(stage) :]: + del correction._diagnostic_knots[stale_stage] + correction._diagnostic_knots[stage] = _copy_knots(correction.knots) + + +def _available_stages(correction) -> tuple[str, ...]: + """Return recorded correction stages in completion order.""" + snapshots = getattr(correction, "_diagnostic_knots", {}) + return tuple(snapshots) + + +def _select_stages( + correction, + stages: Sequence[str] | None, +) -> tuple[str, ...]: + """Validate a requested stage sequence against recorded checkpoints.""" + available = _available_stages(correction) + if not available: + raise RuntimeError( + "Drift diagnostics require a prepared correction. Run preprocess() first." + ) + if stages is None: + return available + requested = tuple(dict.fromkeys(str(stage) for stage in stages)) + missing = [stage for stage in requested if stage not in available] + if missing: + raise ValueError( + f"Diagnostic stages are not available: {missing}. Available stages: {list(available)}." + ) + return requested + + +def _stage_knots(correction, stage: str) -> tuple[NDArray, ...]: + """Return isolated knot arrays for one recorded correction stage.""" + selected = _select_stages(correction, (stage,)) + return _copy_knots(correction._diagnostic_knots[selected[0]]) + + +def _warp_stage( + correction, + stage: str, +) -> tuple[NDArray[np.float32], NDArray[np.float32]]: + """Render a checkpoint without changing the live correction object.""" + knots = _stage_knots(correction, stage) + images = [] + weights = [] + for image_index, stage_knots in enumerate(knots): + warped, coverage = correction.interpolator[image_index].warp_image( + np.asarray(correction.images[image_index].array), + stage_knots, + ) + images.append(np.asarray(warped, dtype=np.float32)) + weights.append(np.asarray(coverage, dtype=np.float32)) + return np.stack(images), np.stack(weights) + + +def _masked_ncc( + reference: NDArray, + moving: NDArray, + mask: NDArray[np.bool_], +) -> float: + """Calculate normalized cross-correlation on explicitly measured pixels.""" + if int(mask.sum()) < _MIN_CORRELATION_PIXELS: + return float("nan") + first = np.asarray(reference, dtype=np.float64)[mask] + second = np.asarray(moving, dtype=np.float64)[mask] + first -= first.mean() + second -= second.mean() + denominator = np.linalg.norm(first) * np.linalg.norm(second) + if denominator <= np.finfo(np.float64).eps: + return float("nan") + return float(first @ second / denominator) + + +def _masked_errors( + reference: NDArray, + moving: NDArray, + mask: NDArray[np.bool_], +) -> tuple[float, float]: + """Return mean absolute and root-mean-square disagreement.""" + if not np.any(mask): + return float("nan"), float("nan") + difference = ( + np.asarray(reference, dtype=np.float64)[mask] - np.asarray(moving, dtype=np.float64)[mask] + ) + return float(np.mean(np.abs(difference))), float(np.sqrt(np.mean(difference**2))) + + +def _scan_comparisons(stack: NDArray) -> list[tuple[NDArray, NDArray]]: + """Build pair or leave-one-out comparisons for a registered scan stack.""" + if stack.shape[0] < 2: + raise ValueError( + f"Registration diagnostics require at least two scans, got {stack.shape[0]}." + ) + if stack.shape[0] == 2: + return [(stack[0], stack[1])] + return [ + ( + np.mean(np.delete(stack, image_index, axis=0), axis=0), + stack[image_index], + ) + for image_index in range(stack.shape[0]) + ] + + +def _pair_metrics( + reference: NDArray, + moving: NDArray, + mask: NDArray[np.bool_], +) -> dict[str, float]: + """Measure one registered pair over common and regional coverage.""" + num_rows = mask.shape[0] + metrics = {"common_ncc": _masked_ncc(reference, moving, mask)} + for name, (row_start, row_stop) in zip( + ("top_ncc", "middle_ncc", "bottom_ncc"), + ( + (0, num_rows // 3), + (num_rows // 3, 2 * num_rows // 3), + (2 * num_rows // 3, num_rows), + ), + strict=True, + ): + regional_mask = np.zeros_like(mask) + regional_mask[row_start:row_stop] = mask[row_start:row_stop] + metrics[name] = _masked_ncc(reference, moving, regional_mask) + mean_absolute, root_mean_square = _masked_errors(reference, moving, mask) + metrics["mean_absolute_difference"] = mean_absolute + metrics["root_mean_square_difference"] = root_mean_square + return metrics + + +def _finite_mean(values: Sequence[float]) -> float: + """Average finite measurements without emitting all-NaN warnings.""" + finite = np.asarray(values, dtype=np.float64) + finite = finite[np.isfinite(finite)] + return float(finite.mean()) if finite.size else float("nan") + + +def _registration_data( + correction, + stages: Sequence[str] | None = None, +) -> tuple[ + tuple[str, ...], + dict[str, NDArray[np.float32]], + NDArray[np.bool_], + list[dict[str, float | str]], +]: + """Build common-mask registration evidence for completed stages.""" + selected = _select_stages(correction, stages) + stacks = {} + coverage_by_stage = {} + for stage in selected: + stack, weights = _warp_stage(correction, stage) + stacks[stage] = stack + coverage_by_stage[stage] = np.all(weights >= _COVERAGE_THRESHOLD, axis=0) + common_mask = np.logical_and.reduce([coverage_by_stage[stage] for stage in selected]) + + rows: list[dict[str, float | str]] = [] + for stage in selected: + pair_rows = [ + _pair_metrics(reference, moving, common_mask) + for reference, moving in _scan_comparisons(stacks[stage]) + ] + row: dict[str, float | str] = {"stage": stage} + for key in pair_rows[0]: + row[key] = _finite_mean([pair[key] for pair in pair_rows]) + row["coverage"] = float(common_mask.mean()) + rows.append(row) + return selected, stacks, common_mask, rows + + +def _displacement_fields( + correction, + stage: str, +) -> NDArray[np.float64]: + """Return row/column knot displacement from the initial checkpoint.""" + initial = _stage_knots(correction, "initial") + selected = _stage_knots(correction, stage) + if any( + current.shape != baseline.shape + for current, baseline in zip(selected, initial, strict=True) + ): + raise ValueError("Displacement diagnostics require matching knot topology across stages.") + return np.stack( + [ + np.asarray(current, dtype=np.float64) - np.asarray(baseline, dtype=np.float64) + for current, baseline in zip(selected, initial, strict=True) + ] + ) + + +def _displacement_rows( + correction, + stages: Sequence[str] | None = None, +) -> list[dict[str, float | int | str]]: + """Summarize scan-line-origin displacement for each image and stage. + + Adjacent-change values are component-wise RMS measurements. Fast-knot + differences are not normalized by knot separation, so they are comparable + only when the fast-direction knot topology and spacing are unchanged. + """ + selected = _select_stages(correction, stages) + rows: list[dict[str, float | int | str]] = [] + for stage in selected: + fields = _displacement_fields(correction, stage) + for image_index, field in enumerate(fields): + magnitude = np.linalg.norm(field, axis=0) + endpoint_vector = np.mean(field[:, -1] - field[:, 0], axis=1) + adjacent_line_change = np.diff(field, axis=1) + adjacent_fast_knot_change = np.diff(field, axis=2) + fast_knot_second_difference = np.diff(field, n=2, axis=2) + rows.append( + { + "stage": stage, + "image": image_index, + "endpoint_row_displacement_px": float(endpoint_vector[0]), + "endpoint_column_displacement_px": float(endpoint_vector[1]), + "endpoint_displacement_px": float(np.linalg.norm(endpoint_vector)), + "rms_displacement_px": float(np.sqrt(np.mean(magnitude**2))), + "max_displacement_px": float(np.max(magnitude)), + "component_rms_adjacent_line_change_px": ( + float(np.sqrt(np.mean(adjacent_line_change**2))) + if adjacent_line_change.size + else 0.0 + ), + "component_rms_adjacent_fast_knot_change_px": ( + float(np.sqrt(np.mean(adjacent_fast_knot_change**2))) + if adjacent_fast_knot_change.size + else 0.0 + ), + "component_rms_fast_knot_second_difference_px": ( + float(np.sqrt(np.mean(fast_knot_second_difference**2))) + if fast_knot_second_difference.size + else 0.0 + ), + } + ) + return rows diff --git a/src/quantem/imaging/drift/plot.py b/src/quantem/imaging/drift/plot.py index ee587567..2ac83bde 100644 --- a/src/quantem/imaging/drift/plot.py +++ b/src/quantem/imaging/drift/plot.py @@ -3,6 +3,7 @@ import numpy as np from quantem.core.visualization import show_2d +from quantem.imaging.drift import diagnostics def plot_transformed_images(self, show_knots: bool = True, **kwargs): @@ -109,3 +110,162 @@ def plot_merged_images(self, show_knots: bool = True, **kwargs): y, x, ) + + +def _comparison_pair(stack: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Compare scan 0 with scan 1, or with the mean of all remaining scans.""" + if stack.shape[0] == 2: + return stack[0], stack[1] + return np.mean(stack[1:], axis=0), stack[0] + + +def _registration_overlay( + reference: np.ndarray, + moving: np.ndarray, + mask: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Build percentile-normalized RGB agreement and difference images.""" + values = np.concatenate((reference[mask], moving[mask])) + if values.size: + low, high = np.percentile(values, (1.0, 99.0)) + else: + low, high = 0.0, 1.0 + scale = max(float(high - low), np.finfo(np.float32).eps) + first = np.clip((reference - low) / scale, 0.0, 1.0) + second = np.clip((moving - low) / scale, 0.0, 1.0) + overlay = np.stack((first, second, np.zeros_like(first)), axis=-1) + difference = np.abs(first - second) + overlay[~mask] = 0.0 + difference[~mask] = 0.0 + return overlay, difference + + +def _plot_registration_diagnostics( + correction, + *, + stages: tuple[str, ...] | list[str] | None = None, + figsize: tuple[float, float] | None = None, +): + """Plot stage-by-stage registration on one fixed measured footprint.""" + selected, stacks, common_mask, _ = diagnostics._registration_data( + correction, + stages, + ) + if figsize is None: + figsize = (11.0, 3.4 * len(selected)) + figure, axes = plt.subplots( + len(selected), + 3, + figsize=figsize, + squeeze=False, + ) + for row_index, stage in enumerate(selected): + reference, moving = _comparison_pair(stacks[stage]) + metrics = diagnostics._pair_metrics(reference, moving, common_mask) + overlay, difference = _registration_overlay(reference, moving, common_mask) + axes[row_index, 0].imshow(overlay) + axes[row_index, 0].set_title( + f"{stage}: RGB registration\ncommon NCC {float(metrics['common_ncc']):.4f}" + ) + axes[row_index, 1].imshow(difference, cmap="magma", vmin=0.0, vmax=1.0) + axes[row_index, 1].set_title( + "Percentile-normalized |difference|\n" + f"native MAD {float(metrics['mean_absolute_difference']):.4g}; " + f"RMS {float(metrics['root_mean_square_difference']):.4g}" + ) + axes[row_index, 2].imshow(common_mask, cmap="gray", vmin=0, vmax=1) + axes[row_index, 2].set_title( + f"Fixed common coverage\n{float(common_mask.mean()):.1%} of canvas" + ) + for axis in axes[row_index]: + axis.set_xticks([]) + axis.set_yticks([]) + figure.suptitle( + "Drift registration diagnostics", + fontsize=14, + fontweight="semibold", + ) + figure.tight_layout() + return figure, axes + + +def _plot_displacement_diagnostics( + correction, + *, + stage: str | None = None, + figsize: tuple[float, float] | None = None, +): + """Plot scan-line-origin displacement without changing fitted knots.""" + if stage is None: + stage = diagnostics._select_stages(correction, None)[-1] + else: + stage = diagnostics._select_stages(correction, (stage,))[0] + fields = diagnostics._displacement_fields(correction, stage) + rows = diagnostics._displacement_rows(correction, (stage,)) + if figsize is None: + figsize = (11.0, 3.6 * len(fields)) + figure, axes = plt.subplots( + len(fields), + 2, + figsize=figsize, + squeeze=False, + ) + for image_index, (field, metrics) in enumerate(zip(fields, rows, strict=True)): + scan_lines = np.arange(field.shape[1]) + row_displacement = np.mean(field[0], axis=1) + column_displacement = np.mean(field[1], axis=1) + axes[image_index, 0].plot( + scan_lines, + row_displacement, + label="row displacement", + ) + axes[image_index, 0].plot( + scan_lines, + column_displacement, + label="column displacement", + ) + axes[image_index, 0].axhline(0.0, color="0.3", linewidth=0.7) + axes[image_index, 0].set_xlabel("Scan line") + axes[image_index, 0].set_ylabel("Displacement (px)") + axes[image_index, 0].set_title( + f"Image {image_index}: scan-line origins\n" + f"endpoint displacement {float(metrics['endpoint_displacement_px']):.3g} px" + ) + axes[image_index, 0].legend(frameon=False) + axes[image_index, 0].grid(alpha=0.25) + + magnitude = np.linalg.norm(field, axis=0) + if magnitude.shape[1] == 1: + axes[image_index, 1].plot( + scan_lines, + magnitude[:, 0], + color="#0072B2", + ) + axes[image_index, 1].set_xlabel("Scan line") + axes[image_index, 1].set_ylabel("Displacement magnitude (px)") + axes[image_index, 1].grid(alpha=0.25) + else: + image = axes[image_index, 1].imshow( + magnitude, + cmap="viridis", + aspect="auto", + origin="upper", + ) + axes[image_index, 1].set_xlabel("Knot along fast scan") + axes[image_index, 1].set_ylabel("Scan line") + figure.colorbar( + image, + ax=axes[image_index, 1], + label="Displacement (px)", + ) + axes[image_index, 1].set_title( + f"Image {image_index}: displacement magnitude\n" + f"maximum {float(metrics['max_displacement_px']):.3g} px" + ) + figure.suptitle( + f"Drift displacement diagnostics: {stage}", + fontsize=14, + fontweight="semibold", + ) + figure.tight_layout() + return figure, axes diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py index ad0a76e4..e679b340 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preparation.py @@ -9,6 +9,7 @@ from quantem.core.utils.imaging_utils import ( cross_correlation_shift, ) +from quantem.imaging.drift import diagnostics from quantem.imaging.drift.core.knots import ( DriftInterpolator, bilinear_kde_batch, @@ -137,6 +138,7 @@ def preprocess( warped_t[img_idx] = warped[0] self.images_warped.array[img_idx] = warped[0].cpu().numpy() self.weights_warped.array[img_idx] = weights[0].cpu().numpy() + diagnostics._record_stage(self, "initial") self.calculate_error(0, _warped_t=warped_t) kwargs.pop("title", None) if show_merged: @@ -191,6 +193,7 @@ def align_translation( self.images[ind].array, self.knots[ind], ) + diagnostics._record_stage(self, "translation") kwargs.pop("title", None) if show_merged: self.plot_merged_images(show_knots=show_knots, title="Merged: translation", **kwargs) diff --git a/src/quantem/imaging/drift/report.py b/src/quantem/imaging/drift/report.py new file mode 100644 index 00000000..812fd544 --- /dev/null +++ b/src/quantem/imaging/drift/report.py @@ -0,0 +1,69 @@ +"""Structured evidence assembled from pure drift diagnostics. + +These helpers intentionally remain internal during staged integration. The +final API change will decide which table becomes ``DriftCorrection.report``. +""" + +from collections.abc import Mapping, Sequence + +from . import diagnostics + + +def _rename_rows( + rows: Sequence[Mapping[str, float | int | str]], + columns: Mapping[str, str], +) -> list[dict[str, float | int | str]]: + """Return independent records with stable, reader-facing labels.""" + return [{columns.get(name, name): value for name, value in row.items()} for row in rows] + + +def _registration_report( + correction, + *, + stages: Sequence[str] | None = None, +) -> list[dict[str, float | int | str]]: + """Return common-coverage registration records by correction stage.""" + _, _, _, rows = diagnostics._registration_data(correction, stages) + return _rename_rows( + rows, + { + "stage": "Correction stage", + "common_ncc": "Common NCC", + "top_ncc": "Top third NCC", + "middle_ncc": "Middle third NCC", + "bottom_ncc": "Bottom third NCC", + "mean_absolute_difference": "Mean absolute difference (native intensity units)", + "root_mean_square_difference": "Root mean square difference (native intensity units)", + "coverage": "Coverage", + }, + ) + + +def _displacement_report( + correction, + *, + stages: Sequence[str] | None = None, +) -> list[dict[str, float | int | str]]: + """Return explicit scan-line-origin displacement records by image.""" + rows = diagnostics._displacement_rows(correction, stages) + return _rename_rows( + rows, + { + "stage": "Correction stage", + "image": "Image", + "endpoint_row_displacement_px": "Endpoint row displacement (px)", + "endpoint_column_displacement_px": "Endpoint column displacement (px)", + "endpoint_displacement_px": "Endpoint displacement (px)", + "rms_displacement_px": "RMS displacement (px)", + "max_displacement_px": "Maximum displacement (px)", + "component_rms_adjacent_line_change_px": ( + "Component RMS adjacent-line displacement change (px)" + ), + "component_rms_adjacent_fast_knot_change_px": ( + "Component RMS adjacent-fast-knot displacement change (px; knot-spacing dependent)" + ), + "component_rms_fast_knot_second_difference_px": ( + "Component RMS fast-knot second displacement difference (px; knot-spacing dependent)" + ), + }, + ) diff --git a/tests/imaging/test_drift.py b/tests/imaging/test_drift.py index 3ac72a06..a9bd5b5c 100644 --- a/tests/imaging/test_drift.py +++ b/tests/imaging/test_drift.py @@ -7,9 +7,14 @@ import numpy as np import pytest +from matplotlib import pyplot as plt from scipy.ndimage import gaussian_filter + from quantem.core.datastructures.dataset2d import Dataset2d from quantem.imaging.drift import DriftCorrection +from quantem.imaging.drift import diagnostics as drift_diagnostics +from quantem.imaging.drift import plot as drift_plot +from quantem.imaging.drift import report as drift_report def make_synthetic_drift_data(scale=1, seed=42): @@ -241,3 +246,166 @@ def test_align_nonrigid_lbfgs_matches_frozen_baseline(scale, expected_error, exp drift.knots[0].sum(), expected_k0, decimal=6) np.testing.assert_almost_equal( drift.knots[1].sum(), expected_k1, decimal=6) + + +def _make_affine_diagnostic_correction(): + """Build the standard synthetic affine result used by diagnostics tests.""" + image_0, image_1, _ = make_synthetic_drift_data(scale=1, seed=42) + correction = DriftCorrection.from_data( + images=[image_0, image_1], + scan_direction_degrees=[0.0, 90.0], + ).preprocess(show_merged=False, show_images=False) + correction.align_affine( + step=0.02, + num_tests=5, + refine=False, + show_merged=False, + show_images=False, + ) + return correction + + +def test_registration_report_matches_frozen_common_coverage_metrics(): + """Registration evidence uses one measured footprint across stages.""" + correction = _make_affine_diagnostic_correction() + registration = drift_report._registration_report(correction) + + assert [row["Correction stage"] for row in registration] == ["initial", "affine"] + np.testing.assert_allclose( + [row["Coverage"] for row in registration], + [0.6423046875, 0.6423046875], + rtol=0, + atol=1e-12, + ) + np.testing.assert_allclose( + [ + [row["Common NCC"], row["Mean absolute difference (native intensity units)"]] + for row in registration + ], + [ + [0.48240861115393513, 0.456299068925209], + [0.7480021414803405, 0.2941893335769055], + ], + rtol=0, + # The affine solver's sub-pixel translation varies slightly across + # CPU thread counts; the report itself is deterministic for fixed knots. + atol=1e-4, + ) + + +def test_displacement_report_names_endpoint_and_adjacent_change_metrics(): + """Displacement evidence distinguishes endpoint motion from adjacent changes.""" + correction = _make_affine_diagnostic_correction() + displacement = drift_report._displacement_report(correction) + + initial = [row for row in displacement if row["Correction stage"] == "initial"] + affine = [row for row in displacement if row["Correction stage"] == "affine"] + assert all( + value == 0 + for row in initial + for name, value in row.items() + if name not in {"Correction stage", "Image"} + ) + np.testing.assert_allclose( + [row["Endpoint displacement (px)"] for row in affine], + [5.679612662972729, 5.679612662972729], + rtol=0, + atol=1e-7, + ) + np.testing.assert_allclose( + [row["Component RMS adjacent-line displacement change (px)"] for row in affine], + [0.0316227766016838, 0.0316227766016838], + rtol=0, + atol=1e-8, + ) + np.testing.assert_array_equal( + [ + row[ + "Component RMS adjacent-fast-knot displacement change (px; knot-spacing dependent)" + ] + for row in affine + ], + 0.0, + ) + + +def test_diagnostic_reports_and_plots_do_not_mutate_correction(): + """Reports and static figures leave fitted fields and caches untouched.""" + correction = _make_affine_diagnostic_correction() + knots_before = [np.asarray(knot).copy() for knot in correction.knots] + images_before = correction.images_warped.array.copy() + weights_before = correction.weights_warped.array.copy() + errors_before = correction.error_track.copy() + + drift_report._registration_report(correction) + drift_report._displacement_report(correction) + registration_figure, registration_axes = drift_plot._plot_registration_diagnostics(correction) + displacement_figure, displacement_axes = drift_plot._plot_displacement_diagnostics(correction) + + assert registration_axes.shape == (2, 3) + assert displacement_axes.shape == (2, 2) + assert registration_figure.number in plt.get_fignums() + assert displacement_figure.number in plt.get_fignums() + for before, after in zip(knots_before, correction.knots, strict=True): + np.testing.assert_array_equal(after, before) + np.testing.assert_array_equal(correction.images_warped.array, images_before) + np.testing.assert_array_equal(correction.weights_warped.array, weights_before) + np.testing.assert_array_equal(correction.error_track, errors_before) + assert drift_diagnostics._available_stages(correction) == ("initial", "affine") + plt.close(registration_figure) + plt.close(displacement_figure) + + +def test_translation_records_current_knots_and_affine_rerun_discards_it(): + """Translation evidence follows the fitted state through an affine rerun.""" + correction = _make_affine_diagnostic_correction() + correction.align_translation(show_merged=False, show_images=False) + + assert drift_diagnostics._available_stages(correction) == ( + "initial", + "affine", + "translation", + ) + for recorded, current in zip( + drift_diagnostics._stage_knots(correction, "translation"), + correction.knots, + strict=True, + ): + np.testing.assert_array_equal(recorded, current) + + correction.align_affine( + step=0.02, + num_tests=5, + refine=False, + show_merged=False, + show_images=False, + ) + + assert drift_diagnostics._available_stages(correction) == ("initial", "affine") + + +def test_multiscan_plot_metric_describes_the_rendered_comparison(): + """A multi-scan plot labels the exact scan-0-versus-rest image it shows.""" + image_0, image_1, _ = make_synthetic_drift_data(scale=1, seed=42) + image_2 = np.roll(image_0, shift=7, axis=0) + correction = DriftCorrection.from_data( + images=[image_0, image_1, image_2], + scan_direction_degrees=[0.0, 90.0, 45.0], + ).preprocess(show_merged=False, show_images=False) + + figure, axes = drift_plot._plot_registration_diagnostics( + correction, + stages=("initial",), + ) + _, stacks, common_mask, aggregate_rows = drift_diagnostics._registration_data( + correction, + ("initial",), + ) + reference, moving = drift_plot._comparison_pair(stacks["initial"]) + expected_overlay, _ = drift_plot._registration_overlay(reference, moving, common_mask) + displayed_metrics = drift_diagnostics._pair_metrics(reference, moving, common_mask) + + np.testing.assert_allclose(axes[0, 0].images[0].get_array(), expected_overlay) + assert f"{displayed_metrics['common_ncc']:.4f}" in axes[0, 0].get_title() + assert not np.isclose(displayed_metrics["common_ncc"], aggregate_rows[0]["common_ncc"]) + plt.close(figure) From 3bd3aac2ada254c5b30015f8af983c9d0226b82a Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 10:40:36 -0700 Subject: [PATCH 05/11] Correct spectrum-image drift without mixing spectral channels (#17) * feat: add reference-based spectrum-image correction * fix: reject ambiguous spectrum-image reference matches --- src/quantem/imaging/__init__.py | 6 + src/quantem/imaging/drift/__init__.py | 6 + src/quantem/imaging/drift/apply.py | 117 ++++++ src/quantem/imaging/drift/core/affine.py | 46 ++- src/quantem/imaging/drift/core/warping.py | 9 +- src/quantem/imaging/drift/correction.py | 80 ++++ src/quantem/imaging/drift/io.py | 458 ++++++++++++++++++++++ src/quantem/imaging/drift/preparation.py | 1 + tests/imaging/test_drift3d.py | 281 +++++++++++++ 9 files changed, 996 insertions(+), 8 deletions(-) create mode 100644 src/quantem/imaging/drift/io.py create mode 100644 tests/imaging/test_drift3d.py diff --git a/src/quantem/imaging/__init__.py b/src/quantem/imaging/__init__.py index 3637acdc..88a7aa59 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,3 +1,9 @@ from quantem.imaging.drift import DriftCorrection as DriftCorrection +from quantem.imaging.drift import ( + pair_spectrum_image_references as pair_spectrum_image_references, + read_emd as read_emd, + read_emd_eds as read_emd_eds, + read_emd_metadata as read_emd_metadata, +) from quantem.imaging.lattice import Lattice as Lattice from quantem.imaging.lattice_visualization import PLOT_REGISTRY as PLOT_REGISTRY diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index e38d7311..10eba339 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -7,6 +7,12 @@ from quantem.imaging.drift.core.knots import ( DriftInterpolator as DriftInterpolator, ) +from quantem.imaging.drift.io import ( + pair_spectrum_image_references as pair_spectrum_image_references, + read_emd as read_emd, + read_emd_eds as read_emd_eds, + read_emd_metadata as read_emd_metadata, +) DriftCorrection.__module__ = __name__ DriftInterpolator.__module__ = __name__ diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index 1c3a0792..a1745880 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -1,4 +1,5 @@ +import copy import warnings import matplotlib.pyplot as plt @@ -7,6 +8,7 @@ from scipy.ndimage import distance_transform_edt from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset3d import Dataset3d from quantem.core.utils.imaging_utils import ( fourier_cropping, ) @@ -17,6 +19,121 @@ ) +@torch.inference_mode() +def apply_correction(self, spectrum_image: Dataset3d) -> Dataset3d: + """Apply one learned spatial field to every spectrum-image channel. + + Parameters + ---------- + spectrum_image : Dataset3d + Spectrum image with axes ``(scan_row, scan_col, energy)``. + + Returns + ------- + Dataset3d + Corrected spectrum image with unchanged calibration, units, metadata, + and spectral-axis ordering. + + Examples + -------- + >>> corrected = drift.apply_correction(spectrum_image) + >>> corrected.shape == spectrum_image.shape + True + """ + if not getattr(self, "_reference_mode", False): + raise RuntimeError( + "apply_correction requires DriftCorrection.from_reference(...)." + ) + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "No drift field found. Call preprocess() and align_affine() first." + ) + if "affine" not in getattr(self, "_diagnostic_knots", {}): + raise RuntimeError( + "No fitted reference field found. Call align_affine() before " + "applying the correction." + ) + if not isinstance(spectrum_image, Dataset3d): + raise TypeError("spectrum_image must be a Dataset3d.") + dataset = spectrum_image + scan_rows, scan_cols, num_channels = dataset.shape + if (scan_rows, scan_cols) != tuple(self.images[1].shape): + raise ValueError( + "The spectrum-image scan axes must match the fitted alignment image; " + f"got {(scan_rows, scan_cols)} and {self.images[1].shape}." + ) + if self.knots[1].shape[-1] != 1: + raise ValueError( + "Spectrum-image correction currently requires number_knots=1." + ) + + device = self._device + delta_canvas_t = torch.as_tensor( + self.knots[1] - self._initial_knots[1], + dtype=torch.float32, + device=device, + )[:, :, 0] + scan_fast_t = torch.as_tensor( + self.scan_fast[1], dtype=torch.float32, device=device + ) + scan_slow_t = torch.as_tensor( + self.scan_slow[1], dtype=torch.float32, device=device + ) + aspect = float(scan_rows - 1) / float(scan_cols - 1) if scan_cols > 1 else 1.0 + determinant = ( + scan_slow_t[0] * scan_fast_t[1] + - scan_fast_t[0] * aspect * scan_slow_t[1] + ) + drift_row_t = ( + scan_fast_t[1] * delta_canvas_t[0] + - scan_fast_t[0] * aspect * delta_canvas_t[1] + ) / determinant + drift_col_t = ( + -scan_slow_t[1] * delta_canvas_t[0] + + scan_slow_t[0] * delta_canvas_t[1] + ) / determinant + row_t = torch.arange(scan_rows, dtype=torch.float32, device=device) + col_t = torch.arange(scan_cols, dtype=torch.float32, device=device) + sample_row_t = row_t[:, None] - drift_row_t[:, None] + sample_col_t = col_t[None, :] - drift_col_t[:, None] + grid_t = torch.stack( + ( + 2.0 * sample_col_t.expand(scan_rows, scan_cols) / (scan_cols - 1) - 1.0, + 2.0 * sample_row_t.expand(scan_rows, scan_cols) / (scan_rows - 1) - 1.0, + ), + dim=-1, + )[None] + + corrected_array = np.empty(dataset.shape, dtype=np.float32) + chunk_size = min(num_channels, 64) + for start in range(0, num_channels, chunk_size): + stop = min(start + chunk_size, num_channels) + channels_t = torch.as_tensor( + np.ascontiguousarray(dataset.array[:, :, start:stop]), + dtype=torch.float32, + device=device, + ).permute(2, 0, 1)[None] + corrected_t = torch.nn.functional.grid_sample( + channels_t, + grid_t, + mode="bilinear", + padding_mode="border", + align_corners=True, + )[0].permute(1, 2, 0) + corrected_array[:, :, start:stop] = corrected_t.cpu().numpy() + + corrected = Dataset3d.from_array( + corrected_array, + name=f"drift-corrected {dataset.name}", + origin=dataset.origin.copy(), + sampling=dataset.sampling.copy(), + units=list(dataset.units), + signal_units=dataset.signal_units, + ) + corrected.metadata.update(copy.deepcopy(dataset.metadata)) + return corrected + + @torch.inference_mode() def generate_corrected( self, diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index 82af9334..5724f1db 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -98,6 +98,8 @@ def align_affine( raise ValueError( f"num_tests must be odd (got {num_tests}). Try {num_tests + 1}." ) + fixed_reference = getattr(self, "_reference_mode", False) + # Build candidate grid with circular mask (~21% fewer than square) grid_axis = np.arange(-(num_tests - 1) / 2, (num_tests + 1) / 2) row_grid, col_grid = np.meshgrid(grid_axis, grid_axis, indexing="ij") @@ -120,16 +122,28 @@ def _print_top_candidates(label, candidates, costs_tensor): def _apply_drift(drift_vec): for img_idx in range(self.shape[0]): + if fixed_reference and img_idx == 0: + continue scanline_offset = np.arange(self.knots[img_idx].shape[1]) - (self.knots[img_idx].shape[1] - 1) / 2 self.knots[img_idx][0] += drift_vec[0] * scanline_offset[:, None] self.knots[img_idx][1] += drift_vec[1] * scanline_offset[:, None] def _search_and_apply(candidates, label): - best_idx, costs = self._affine_grid_search_batch(candidates, upsample_factor, max_image_shift, chunk_size) + best_idx, costs = self._affine_grid_search_batch( + candidates, + upsample_factor, + max_image_shift, + chunk_size, + fixed_reference=fixed_reference, + ) _apply_drift(candidates[best_idx]) if verbose: _print_top_candidates(label, candidates, costs) - warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor) + warped_t = self._warp_and_translate_torch( + max_image_shift, + upsample_factor, + fixed_reference=fixed_reference, + ) self.calculate_error(1, _warped_t=warped_t) return candidates[best_idx] @@ -178,7 +192,14 @@ def _search_and_apply(candidates, label): @torch.inference_mode() -def _affine_grid_search_batch(self, drift_vectors, upsample_factor, max_image_shift, chunk_size=None): +def _affine_grid_search_batch( + self, + drift_vectors, + upsample_factor, + max_image_shift, + chunk_size=None, + fixed_reference=False, +): """Evaluate all candidate drift vectors in parallel. Warps both images for each candidate using ``bilinear_kde_batch`` @@ -245,8 +266,13 @@ def _affine_grid_search_batch(self, drift_vectors, upsample_factor, max_image_sh warped_pair = [] for img_idx in range(2): image_t, row_base, col_base, scanline_offset = base_data[img_idx] - row_candidates = row_base[None] + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] - col_candidates = col_base[None] + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] + if fixed_reference and img_idx == 0: + batch_size = drift_chunk.shape[0] + row_candidates = row_base[None].expand(batch_size, -1, -1) + col_candidates = col_base[None].expand(batch_size, -1, -1) + else: + row_candidates = row_base[None] + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] + col_candidates = col_base[None] + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] warped, _ = bilinear_kde_batch( row_candidates, col_candidates, image_t, canvas_shape, self.kde_sigma, @@ -307,6 +333,7 @@ def _warp_and_translate_torch( upsample_factor: int = 8, knots_batch: torch.Tensor | None = None, solve_translation: bool = True, + fixed_reference: bool = False, ) -> torch.Tensor: """Regenerate warped images and solve translation on GPU. @@ -331,6 +358,8 @@ def _warp_and_translate_torch( solve_translation : bool If False, skip translation alignment (Phase 2+3). Only warp once using current knots and sync to CPU. + fixed_reference : bool, default False + Keep image 0 fixed while translating later images. Returns ------- @@ -366,7 +395,12 @@ def _warp_all(warped_t, weights_t): self.weights_warped.array[:] = weights_t.cpu().numpy() return warped_t # Solve translation shifts and apply to knots - shifts_t = translate_align(warped_t, upsample_factor, max_image_shift) + shifts_t = translate_align( + warped_t, + upsample_factor, + max_image_shift, + anchor_first=fixed_reference, + ) if knots_batch is not None: knots_batch[:, 0] += shifts_t[:, 0:1] knots_batch[:, 1] += shifts_t[:, 1:2] diff --git a/src/quantem/imaging/drift/core/warping.py b/src/quantem/imaging/drift/core/warping.py index 64c2b7df..8b741044 100644 --- a/src/quantem/imaging/drift/core/warping.py +++ b/src/quantem/imaging/drift/core/warping.py @@ -93,6 +93,7 @@ def translate_align( warped_images: torch.Tensor, upsample_factor: int, max_image_shift: float | None, + anchor_first: bool = False, ) -> torch.Tensor: """Pairwise translation alignment of warped images via cross-correlation. @@ -113,6 +114,9 @@ def translate_align( Sub-pixel precision (1/N pixel) for DFT refinement. max_image_shift : float or None Maximum allowed shift in pixels. Peaks beyond this radius are masked. + anchor_first : bool, default False + Keep image 0 fixed and apply the full relative translation to later + images. Used when a reference image defines the output coordinates. Returns ------- @@ -162,8 +166,9 @@ def translate_align( ) ) ref_fft = ref_fft * img_idx / (img_idx + 1) + mov_fft * phase / (img_idx + 1) - # Remove mean so shifts are relative (no absolute reference frame) - image_shifts -= image_shifts.mean(dim=0) + if not anchor_first: + # Mutual correction has no absolute reference frame. + image_shifts -= image_shifts.mean(dim=0) return image_shifts diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index eefcda3a..63497337 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -133,6 +133,7 @@ def __init__( self.images = images self.scan_direction_degrees = ensure_valid_array(scan_direction_degrees, ndim=1) + self._reference_mode = False device, _ = validate_device(None) self._device = device @@ -165,6 +166,84 @@ def from_data( _token=cls._token, ) + @classmethod + def from_reference( + cls, + reference_image: Dataset2d | NDArray | Self, + alignment_image: Dataset2d | NDArray, + *, + scan_direction_degrees: float | None = None, + ) -> Self: + """Fit a fixed-reference drift field from two-dimensional images. + + The two-dimensional ``reference_image`` defines the output frame. + ``alignment_image`` is the HAADF image acquired with a spectrum image. + Pass the spectrum image to :meth:`apply_correction` after fitting. + + Parameters + ---------- + reference_image : Dataset2d, numpy.ndarray, or DriftCorrection + Fixed two-dimensional reference image. A solved correction is + converted to its corrected native-size image automatically. + alignment_image : Dataset2d or numpy.ndarray + HAADF image acquired with the spectrum image. + scan_direction_degrees : float or None, default None + Recorded scan rotation. Omit when ``alignment_image`` carries + ``metadata["scan_rotation_deg"]``. + + Returns + ------- + DriftCorrection + Correction ready for ``preprocess`` and ``align_affine``. + + Examples + -------- + >>> drift = DriftCorrection.from_reference( + ... reference, + ... haadf, + ... scan_direction_degrees=0.0, + ... ) + >>> drift.preprocess(show_merged=False, show_images=False) + >>> drift.align_affine(show_merged=False, show_images=False) + >>> corrected = drift.apply_correction(spectrum_image) + """ + if isinstance(reference_image, cls): + reference_image = reference_image.generate_corrected( + upsample_factor=1, + strip_padding=True, + mask_output=False, + show_merged=False, + ) + reference_array = ( + reference_image.array + if isinstance(reference_image, Dataset2d) + else ensure_valid_array(reference_image, ndim=2) + ) + alignment_array = ( + alignment_image.array + if isinstance(alignment_image, Dataset2d) + else ensure_valid_array(alignment_image, ndim=2) + ) + if reference_array.shape != alignment_array.shape: + raise ValueError( + "reference_image and alignment_image must have the same " + f"shape; got {reference_array.shape} and {alignment_array.shape}." + ) + if scan_direction_degrees is None: + metadata = getattr(alignment_image, "metadata", {}) + scan_direction_degrees = metadata.get("scan_rotation_deg") + if scan_direction_degrees is None: + raise TypeError( + "scan_direction_degrees is required when alignment_image " + "does not carry scan_rotation_deg metadata." + ) + result = cls.from_data( + images=[reference_image, alignment_image], + scan_direction_degrees=[scan_direction_degrees, scan_direction_degrees], + ) + result._reference_mode = True + return result + preprocess = preparation.preprocess align_translation = preparation.align_translation align_affine = affine.align_affine @@ -179,6 +258,7 @@ def from_data( _optimize_knots_scipy = nonrigid._optimize_knots_scipy generate_corrected = drift_apply.generate_corrected generate_corrected_image = drift_apply.generate_corrected_image + apply_correction = drift_apply.apply_correction calculate_error = drift_apply.calculate_error plot_transformed_images = drift_plot.plot_transformed_images plot_convergence = drift_plot.plot_convergence diff --git a/src/quantem/imaging/drift/io.py b/src/quantem/imaging/drift/io.py new file mode 100644 index 00000000..a7f19e49 --- /dev/null +++ b/src/quantem/imaging/drift/io.py @@ -0,0 +1,458 @@ +"""Velox EMD input and metadata pairing for drift correction.""" + +import json +import math +from pathlib import Path + +import h5py +import numpy as np + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset3d import Dataset3d + + +def _deinterleave_velox_metadata(raw: bytes) -> bytes: + """Undo repeated-byte encodings found in some Velox metadata records.""" + for stride in range(2, 17): + candidate = raw[::stride] + if candidate.startswith(b"{"): + return candidate + return raw + + +def _read_velox_metadata_record(file_path: str | Path) -> dict: + """Decode the first image metadata record without loading image pixels.""" + with h5py.File(file_path, "r") as handle: + image_group = handle.get("Data/Image") + if image_group is None: + return {} + for image_name in image_group: + metadata = image_group[image_name].get("Metadata") + if metadata is None: + continue + raw = bytes(metadata[:]).replace(b"\x00", b"") + for candidate in (raw, _deinterleave_velox_metadata(raw)): + try: + return json.loads(candidate.decode("utf-8", errors="ignore")) + except json.JSONDecodeError: + continue + return {} + + +def read_emd_metadata(file_path: str | Path) -> dict[str, object]: + """Read normalized Velox acquisition metadata without loading pixels. + + Parameters + ---------- + file_path : str or Path + Velox EMD file. + + Returns + ------- + dict[str, object] + Scan rotation, stage position, calibration, field of view, acquisition + time, acquisition context, and the decoded original metadata. + + Examples + -------- + >>> metadata = read_emd_metadata("scan.emd") + >>> metadata["scan_rotation_deg"] + 0.0 + """ + file_path = Path(file_path) + metadata = _read_velox_metadata_record(file_path) + with h5py.File(file_path, "r") as handle: + acquisition_context = ( + "spectrum_image" + if handle.get("Data/SpectrumImage") is not None + or handle.get("Data/SpectrumStream") is not None + else "image" + ) + rotation = metadata.get("Scan", {}).get("ScanRotation") + magnification = metadata.get("Optics", {}).get("NominalMagnification") + position = metadata.get("Stage", {}).get("Position", {}) or {} + stage_position = ( + (float(position["x"]), float(position["y"])) + if "x" in position and "y" in position + else None + ) + pixel_size_m = metadata.get("BinaryResult", {}).get("PixelSize", {}).get("width") + scan_size = metadata.get("Scan", {}).get("ScanSize", {}) or {} + scan_columns = scan_size.get("width") + scan_rows = scan_size.get("height", scan_columns) + scan_shape = ( + (int(scan_rows), int(scan_columns)) + if scan_rows and scan_columns + else None + ) + timestamp = metadata.get("Acquisition", {}).get( + "AcquisitionStartDatetime", {} + ).get("DateTime") + return { + "path": str(file_path), + "scan_rotation_deg": None if rotation is None else math.degrees(float(rotation)), + "magnification": None if magnification is None else float(magnification), + "stage_position_m": stage_position, + "pixel_size_nm": None if pixel_size_m is None else float(pixel_size_m) * 1e9, + "scan_shape": scan_shape, + "field_of_view_m": ( + float(pixel_size_m) * float(scan_columns) + if pixel_size_m and scan_columns + else None + ), + "acquisition_timestamp": ( + int(timestamp) if timestamp and str(timestamp).isdigit() else None + ), + "acquisition_context": acquisition_context, + "original_metadata": metadata, + } + + +def _axis_calibration(axes: list[dict], ndim: int) -> tuple[list[float], list[float], list[str]]: + """Return origin, sampling, and units in array-axis order.""" + indexed = sorted( + ( + (int(axis.get("index_in_array", fallback)), axis) + for fallback, axis in enumerate(axes[:ndim]) + ), + key=lambda item: item[0], + ) + ordered = [axis for _, axis in indexed] + origin = [float(axis.get("offset", 0.0)) for axis in ordered] + sampling = [float(axis.get("scale", 1.0)) for axis in ordered] + units = [(axis.get("units") or "pixels") for axis in ordered] + return origin, sampling, units + + +def _image_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset2d: + """Build a calibrated image dataset from one RosettaSciIO stream.""" + origin, sampling, units = _axis_calibration(stream.get("axes", []), 2) + title = stream.get("metadata", {}).get("General", {}).get("title") + image = Dataset2d.from_array( + np.ascontiguousarray(stream["data"], dtype=np.float32), + name=title or Path(path).stem, + origin=origin, + sampling=sampling, + units=units, + ) + image.file_path = str(path) + image.metadata.update(metadata) + return image + + +def read_emd(path: str | Path) -> Dataset2d: + """Read the calibrated HAADF image and recorded scan geometry. + + Parameters + ---------- + path : str or Path + Velox EMD acquisition. + + Returns + ------- + Dataset2d + HAADF image carrying normalized acquisition metadata. + + Examples + -------- + >>> image = read_emd("scan.emd") + >>> image.metadata["scan_rotation_deg"] + 0.0 + """ + from rsciio.emd import file_reader + + streams = file_reader(str(path), select_type="images") + stream = next( + ( + item + for item in streams + if item["data"].ndim == 2 + and item.get("metadata", {}).get("General", {}).get("title") == "HAADF" + ), + next((item for item in streams if item["data"].ndim == 2), None), + ) + if stream is None: + raise ValueError(f"{path} contains no two-dimensional image stream.") + return _image_dataset(stream, path, read_emd_metadata(path)) + + +def _spectrum_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset3d: + """Build a scan-row, scan-column, energy Dataset3d.""" + array = np.asarray(stream["data"]) + axes = stream.get("axes", []) + indexed_axes = [ + axis + for _, axis in sorted( + ( + (int(axis.get("index_in_array", fallback)), axis) + for fallback, axis in enumerate(axes[:3]) + ), + key=lambda item: item[0], + ) + ] + energy_axis = next( + ( + index + for index, axis in enumerate(indexed_axes) + if (axis.get("units") or "").lower() in {"ev", "kev"} + or "energy" in (axis.get("name") or "").lower() + ), + 2, + ) + if energy_axis != 2: + array = np.moveaxis(array, energy_axis, 2) + axis = indexed_axes.pop(energy_axis) + indexed_axes.append(axis) + origin = [float(axis.get("offset", 0.0)) for axis in indexed_axes] + sampling = [float(axis.get("scale", 1.0)) for axis in indexed_axes] + units = [(axis.get("units") or "pixels") for axis in indexed_axes] + title = stream.get("metadata", {}).get("General", {}).get("title") + spectrum = Dataset3d.from_array( + np.ascontiguousarray(array, dtype=np.float32), + name=title or Path(path).stem, + origin=origin, + sampling=sampling, + units=units, + signal_units="counts", + ) + spectrum.file_path = str(path) + spectrum.metadata.update(metadata) + return spectrum + + +def read_emd_eds( + path: str | Path, + *, + load_spectrum: bool = False, + verbose: bool = True, +) -> dict[str, object]: + """Read one Velox EDS or EELS spectrum-image acquisition. + + The full spectrum is opt-in because it can be large. When loaded, it is a + calibrated :class:`Dataset3d` with ``(scan_row, scan_col, energy)`` axes. + Energy calibration and units are retained exactly as recorded. + + Parameters + ---------- + path : str or Path + Velox spectrum-image EMD file. + load_spectrum : bool, default False + Load the full spectrum in addition to the HAADF and element maps. + verbose : bool, default True + Print a compact inventory. + + Returns + ------- + dict[str, object] + HAADF image, optional spectrum, element maps, and acquisition metadata. + + Examples + -------- + >>> acquisition = read_emd_eds("spectrum_image.emd", load_spectrum=True) + >>> acquisition["spectrum"].units[-1] + 'keV' + """ + from rsciio.emd import file_reader + + metadata = read_emd_metadata(path) + streams = list( + file_reader( + str(path), + select_type=None if load_spectrum else "images", + ) + ) + spectrum_stream = next( + (stream for stream in streams if np.asarray(stream["data"]).ndim == 3), + None, + ) + if load_spectrum and spectrum_stream is None: + raise ValueError(f"{path} contains no three-dimensional spectrum.") + haadf_stream = next( + ( + stream + for stream in streams + if np.asarray(stream["data"]).ndim == 2 + and stream.get("metadata", {}).get("General", {}).get("title") == "HAADF" + ), + None, + ) + element_maps = { + stream.get("metadata", {}).get("General", {}).get("title", "map"): + np.ascontiguousarray(stream["data"], dtype=np.float32) + for stream in streams + if np.asarray(stream["data"]).ndim == 2 and stream is not haadf_stream + } + haadf = ( + None + if haadf_stream is None + else _image_dataset(haadf_stream, path, metadata) + ) + spectrum = ( + None + if spectrum_stream is None + else _spectrum_dataset(spectrum_stream, path, metadata) + ) + if haadf is None and not element_maps: + raise ValueError(f"{path} contains no HAADF image or element maps.") + if verbose: + print( + f"HAADF {None if haadf is None else haadf.shape} " + f"elements {list(element_maps)} " + f"spectrum {None if spectrum is None else spectrum.shape} " + f"scan {metadata['scan_rotation_deg']} deg" + ) + return { + "haadf": haadf, + "spectrum": spectrum, + "element_maps": element_maps, + "metadata": metadata, + "path": str(path), + } + + +def _relative_difference(first: float | None, second: float | None) -> float: + """Return relative difference, treating missing metadata as compatible.""" + if first is None or second is None: + return 0.0 + scale = max(abs(float(first)), abs(float(second)), 1e-30) + return abs(float(first) - float(second)) / scale + + +def _same_specimen_area(first: dict, second: dict) -> bool: + """Match scan grid, calibration, field of view, and stage position.""" + if ( + first["scan_shape"] is not None + and second["scan_shape"] is not None + and tuple(first["scan_shape"]) != tuple(second["scan_shape"]) + ): + return False + if _relative_difference(first["pixel_size_nm"], second["pixel_size_nm"]) > 0.02: + return False + if _relative_difference(first["field_of_view_m"], second["field_of_view_m"]) > 0.02: + return False + if _relative_difference(first["magnification"], second["magnification"]) > 0.02: + return False + first_stage = first["stage_position_m"] + second_stage = second["stage_position_m"] + if first_stage is None or second_stage is None: + return False + distance = float(np.linalg.norm(np.subtract(first_stage, second_stage))) + fields = [ + float(value) + for value in (first["field_of_view_m"], second["field_of_view_m"]) + if value is not None and float(value) > 0 + ] + tolerance = max(0.25 * min(fields), 10e-9) if fields else 10e-9 + return distance <= tolerance + + +def pair_spectrum_image_references( + folder: str | Path, + *, + rotation_tolerance_degrees: float = 5.0, +) -> list[dict[str, object]]: + """Match spectrum images to orthogonal HAADF references from metadata. + + Filenames and acquisition order are never used. Each spectrum image must + have one compatible near-zero reference and one compatible near-90-degree + reference from the same specimen area, scan grid, and calibration. The + assignment must also be mutual: a reference pair cannot silently match + multiple spectrum images. Ambiguous and incomplete matches are returned + with a reason rather than guessed. + + Parameters + ---------- + folder : str or Path + Folder containing Velox EMD acquisitions. + rotation_tolerance_degrees : float, default 5.0 + Allowed deviation from 0 and 90 degrees. + + Returns + ------- + list[dict[str, object]] + One record per spectrum image with paths, status, and reason. + + Examples + -------- + >>> matches = pair_spectrum_image_references("~/data/session") + >>> ready = [match for match in matches if match["status"] == "ready"] + """ + folder = Path(folder).expanduser() + records = [ + read_emd_metadata(path) + for path in sorted(folder.glob("*.emd")) + if not path.name.startswith("._") + ] + references = [ + record for record in records if record["acquisition_context"] == "image" + ] + spectrum_images = [ + record + for record in records + if record["acquisition_context"] == "spectrum_image" + ] + tolerance = float(rotation_tolerance_degrees) + matches = [] + for spectrum_image in spectrum_images: + candidates = [ + reference + for reference in references + if _same_specimen_area(spectrum_image, reference) + ] + zero = [ + item + for item in candidates + if item["scan_rotation_deg"] is not None + and abs(float(item["scan_rotation_deg"])) <= tolerance + ] + orthogonal = [ + item + for item in candidates + if item["scan_rotation_deg"] is not None + and abs(abs(float(item["scan_rotation_deg"])) - 90.0) <= tolerance + ] + if len(zero) == 1 and len(orthogonal) == 1: + status = "ready" + reason = "" + elif len(zero) > 1 or len(orthogonal) > 1: + status = "ambiguous" + reason = ( + f"Found {len(zero)} near-zero and {len(orthogonal)} " + "near-90-degree compatible references." + ) + else: + status = "unpaired" + reason = ( + f"Found {len(zero)} near-zero and {len(orthogonal)} " + "near-90-degree compatible references." + ) + matches.append( + { + "spectrum_image": Path(spectrum_image["path"]), + "reference_zero": Path(zero[0]["path"]) if len(zero) == 1 else None, + "reference_orthogonal": ( + Path(orthogonal[0]["path"]) if len(orthogonal) == 1 else None + ), + "status": status, + "reason": reason, + } + ) + pair_users: dict[tuple[Path, Path], list[int]] = {} + for index, match in enumerate(matches): + if match["status"] != "ready": + continue + pair = (match["reference_zero"], match["reference_orthogonal"]) + pair_users.setdefault(pair, []).append(index) + for indices in pair_users.values(): + if len(indices) == 1: + continue + reason = ( + "Ambiguous: the same compatible reference pair matches " + f"{len(indices)} spectrum images." + ) + for index in indices: + matches[index]["status"] = "ambiguous" + matches[index]["reason"] = reason + matches[index]["reference_zero"] = None + matches[index]["reference_orthogonal"] = None + return matches diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py index e679b340..cbb66594 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preparation.py @@ -102,6 +102,7 @@ def preprocess( + u_fast[None, :] * self.scan_fast[img_idx, 1] + v_slow[:, None] * self.scan_slow[img_idx, 1]) self.knots.append(np.stack([row_knots, col_knots], axis=0)) + self._initial_knots = [knots.copy() for knots in self.knots] self.interpolator = [ DriftInterpolator( input_shape=self.images[i].shape, diff --git a/tests/imaging/test_drift3d.py b/tests/imaging/test_drift3d.py new file mode 100644 index 00000000..3035492c --- /dev/null +++ b/tests/imaging/test_drift3d.py @@ -0,0 +1,281 @@ +"""Reference-based spectrum-image drift correction workflows.""" + +import json +import math + +import h5py +import numpy as np +from scipy.ndimage import gaussian_filter + +from quantem.core.datastructures.dataset3d import Dataset3d +from quantem.imaging.drift import ( + DriftCorrection, + pair_spectrum_image_references, + read_emd_eds, +) + + +def _column_drift(array: np.ndarray, rate: float) -> np.ndarray: + """Apply a known scan-row-dependent column displacement.""" + scan_rows, scan_cols = array.shape[:2] + columns = np.arange(scan_cols) + drifted = np.empty_like(array) + for scan_row in range(scan_rows): + sample = np.clip(columns + rate * scan_row, 0, scan_cols - 1.001) + lower = np.floor(sample).astype(int) + fraction = sample - lower + if array.ndim == 3: + fraction = fraction[:, None] + drifted[scan_row] = ( + array[scan_row, lower] * (1.0 - fraction) + + array[scan_row, lower + 1] * fraction + ) + return drifted + + +def test_reference_correction_preserves_spectra_and_calibration(): + """One HAADF-derived field corrects every spectrum channel identically.""" + scan_size = 80 + row, column = np.mgrid[:scan_size, :scan_size] + rng = np.random.default_rng(4) + reference = gaussian_filter( + rng.normal(size=(scan_size, scan_size)).astype(np.float32), + 1.2, + ) + for center_row, center_column in ((18, 20), (29, 61), (63, 24), (58, 65)): + reference += 4.0 * np.exp( + -( + (row - center_row) ** 2 + + (column - center_column) ** 2 + ) + / (2.0 * 4.0**2) + ) + clean_spectrum = np.stack( + ( + reference, + 2.0 * reference + 3.0, + np.zeros_like(reference), + np.full_like(reference, 11.0), + ), + axis=-1, + ).astype(np.float32) + drift_rate = 0.16 + alignment_image = _column_drift(reference, drift_rate).astype(np.float32) + spectrum_image = Dataset3d.from_array( + _column_drift(clean_spectrum, drift_rate).astype(np.float32), + name="SrTiO3 spectrum image", + origin=[1.2, 2.4, 0.35], + sampling=[0.08, 0.08, 0.01], + units=["nm", "nm", "keV"], + signal_units="counts", + ) + spectrum_image.metadata.update( + {"scan_rotation_deg": 0.0, "detector": "Super-X"} + ) + + drift = DriftCorrection.from_reference( + reference, + alignment_image, + scan_direction_degrees=0.0, + ).preprocess(show_merged=False, show_images=False) + initial_reference_knots = drift.knots[0].copy() + drift.align_affine( + step=0.04, + num_tests=11, + refine=True, + max_image_shift=16, + show_merged=False, + show_images=False, + ) + corrected = drift.apply_correction(spectrum_image) + + interior = np.s_[12:-12, 12:-12] + raw_ncc = np.corrcoef( + clean_spectrum[interior].ravel(), + spectrum_image.array[interior].ravel(), + )[0, 1] + corrected_ncc = np.corrcoef( + clean_spectrum[interior].ravel(), + corrected.array[interior].ravel(), + )[0, 1] + assert corrected_ncc > raw_ncc + 0.02 + assert corrected_ncc > 0.99 + np.testing.assert_allclose(drift.knots[0], initial_reference_knots, atol=0.0) + np.testing.assert_allclose( + corrected.array[..., 1], + 2.0 * corrected.array[..., 0] + 3.0, + atol=2e-6, + ) + np.testing.assert_array_equal(corrected.array[..., 2], 0.0) + np.testing.assert_allclose(corrected.array[..., 3], 11.0, atol=2e-6) + np.testing.assert_array_equal(corrected.origin, spectrum_image.origin) + np.testing.assert_array_equal(corrected.sampling, spectrum_image.sampling) + assert corrected.units == spectrum_image.units + assert corrected.signal_units == spectrum_image.signal_units + assert corrected.metadata == spectrum_image.metadata + + +def _write_metadata_emd( + path, + *, + rotation_degrees: float, + stage_position: tuple[float, float], + spectrum_image: bool, + timestamp: int, + scan_shape: tuple[int, int] = (256, 256), + pixel_size_m: float = 6.74e-9 / 256, +): + """Write the minimal Velox metadata used by the pairing workflow.""" + metadata = { + "Scan": { + "ScanRotation": math.radians(rotation_degrees), + "ScanSize": {"width": scan_shape[1], "height": scan_shape[0]}, + }, + "Optics": {"NominalMagnification": 15_000_000}, + "Stage": { + "Position": {"x": stage_position[0], "y": stage_position[1]} + }, + "BinaryResult": {"PixelSize": {"width": pixel_size_m}}, + "Acquisition": { + "AcquisitionStartDatetime": {"DateTime": str(timestamp)} + }, + } + encoded = np.frombuffer(json.dumps(metadata).encode(), dtype=np.uint8) + with h5py.File(path, "w") as handle: + image = handle.create_group("Data/Image/0") + image.create_dataset("Metadata", data=encoded) + if spectrum_image: + handle.create_group("Data/SpectrumImage") + + +def test_spectrum_image_pairing_uses_metadata_not_names_or_order(tmp_path): + """A spectrum image is paired by rotation and stage metadata.""" + stage = (1.2e-6, -3.4e-6) + _write_metadata_emd( + tmp_path / "zzz_last_name.emd", + rotation_degrees=0.0, + stage_position=stage, + spectrum_image=False, + timestamp=30, + ) + _write_metadata_emd( + tmp_path / "aaa_first_name.emd", + rotation_degrees=-90.0, + stage_position=stage, + spectrum_image=False, + timestamp=10, + ) + _write_metadata_emd( + tmp_path / "middle_name.emd", + rotation_degrees=0.0, + stage_position=stage, + spectrum_image=True, + timestamp=20, + ) + _write_metadata_emd( + tmp_path / "nearby_but_wrong_area.emd", + rotation_degrees=90.0, + stage_position=(stage[0] + 100e-9, stage[1]), + spectrum_image=False, + timestamp=40, + ) + _write_metadata_emd( + tmp_path / "same_area_wrong_grid.emd", + rotation_degrees=90.0, + stage_position=stage, + spectrum_image=False, + timestamp=50, + scan_shape=(128, 128), + pixel_size_m=6.74e-9 / 128, + ) + + match = pair_spectrum_image_references(tmp_path)[0] + + assert match["status"] == "ready" + assert match["spectrum_image"].name == "middle_name.emd" + assert match["reference_zero"].name == "zzz_last_name.emd" + assert match["reference_orthogonal"].name == "aaa_first_name.emd" + + +def test_spectrum_image_pairing_rejects_shared_reference_assignment(tmp_path): + """One reference pair cannot be silently reused for multiple acquisitions.""" + stage = (1.2e-6, -3.4e-6) + _write_metadata_emd( + tmp_path / "reference_zero.emd", + rotation_degrees=0.0, + stage_position=stage, + spectrum_image=False, + timestamp=10, + ) + _write_metadata_emd( + tmp_path / "reference_orthogonal.emd", + rotation_degrees=90.0, + stage_position=stage, + spectrum_image=False, + timestamp=20, + ) + for index in range(2): + _write_metadata_emd( + tmp_path / f"spectrum_{index}.emd", + rotation_degrees=0.0, + stage_position=stage, + spectrum_image=True, + timestamp=30 + index, + ) + + matches = pair_spectrum_image_references(tmp_path) + + assert len(matches) == 2 + assert {match["status"] for match in matches} == {"ambiguous"} + assert all("matches 2 spectrum images" in match["reason"] for match in matches) + assert all(match["reference_zero"] is None for match in matches) + assert all(match["reference_orthogonal"] is None for match in matches) + + +def test_read_emd_eds_preserves_native_energy_axis(tmp_path, monkeypatch): + """The EMD loader returns scan axes first and native energy calibration.""" + path = tmp_path / "spectrum.emd" + _write_metadata_emd( + path, + rotation_degrees=0.0, + stage_position=(0.0, 0.0), + spectrum_image=True, + timestamp=1, + ) + native = np.arange(4 * 5 * 6, dtype=np.uint32).reshape(4, 5, 6) + streams = [ + { + "data": np.ones((5, 6), dtype=np.float32), + "metadata": {"General": {"title": "HAADF"}}, + "axes": [ + {"index_in_array": 0, "scale": 0.2, "offset": 1.0, "units": "nm"}, + {"index_in_array": 1, "scale": 0.2, "offset": 2.0, "units": "nm"}, + ], + }, + { + "data": native, + "metadata": {"General": {"title": "EDS"}}, + "axes": [ + { + "index_in_array": 0, + "name": "Energy", + "scale": 0.01, + "offset": 0.35, + "units": "keV", + }, + {"index_in_array": 1, "scale": 0.2, "offset": 1.0, "units": "nm"}, + {"index_in_array": 2, "scale": 0.2, "offset": 2.0, "units": "nm"}, + ], + }, + ] + monkeypatch.setattr("rsciio.emd.file_reader", lambda *args, **kwargs: streams) + + acquisition = read_emd_eds(path, load_spectrum=True, verbose=False) + spectrum = acquisition["spectrum"] + + assert isinstance(spectrum, Dataset3d) + assert spectrum.shape == (5, 6, 4) + np.testing.assert_array_equal(spectrum.array, np.moveaxis(native, 0, 2)) + np.testing.assert_array_equal(spectrum.origin, [1.0, 2.0, 0.35]) + np.testing.assert_array_equal(spectrum.sampling, [0.2, 0.2, 0.01]) + assert spectrum.units == ["nm", "nm", "keV"] From 0f16b6a71441ba9936d66216ce000f37b409ef92 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 10:55:20 -0700 Subject: [PATCH 06/11] Correct 4D-STEM scan axes without altering detector pixels (#18) * feat: add reference-based spectrum-image correction * fix: reject ambiguous spectrum-image reference matches * feat: propagate drift across 4D-STEM scan axes * test: protect 4D-STEM detector coordinates * test: compare full detector signature field * test: verify 4D-STEM coverage and output residency * fix: stream native 4D-STEM scan correction * test: compare native detector frames across backends * test: set float32 backend parity tolerance --- src/quantem/imaging/__init__.py | 5 +- src/quantem/imaging/drift/__init__.py | 3 + src/quantem/imaging/drift/apply.py | 251 ++++++-- src/quantem/imaging/drift/correction.py | 117 +++- src/quantem/imaging/drift/fourdstem.py | 782 ++++++++++++++++++++++++ tests/imaging/test_drift4d.py | 425 +++++++++++++ 6 files changed, 1517 insertions(+), 66 deletions(-) create mode 100644 src/quantem/imaging/drift/fourdstem.py create mode 100644 tests/imaging/test_drift4d.py diff --git a/src/quantem/imaging/__init__.py b/src/quantem/imaging/__init__.py index 88a7aa59..929f549b 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,4 +1,7 @@ -from quantem.imaging.drift import DriftCorrection as DriftCorrection +from quantem.imaging.drift import ( + CorrectionResult as CorrectionResult, + DriftCorrection as DriftCorrection, +) from quantem.imaging.drift import ( pair_spectrum_image_references as pair_spectrum_image_references, read_emd as read_emd, diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index 10eba339..5e8f720d 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -7,6 +7,9 @@ from quantem.imaging.drift.core.knots import ( DriftInterpolator as DriftInterpolator, ) +from quantem.imaging.drift.fourdstem import ( + CorrectionResult as CorrectionResult, +) from quantem.imaging.drift.io import ( pair_spectrum_image_references as pair_spectrum_image_references, read_emd as read_emd, diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index a1745880..3f93f428 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -19,65 +19,24 @@ ) -@torch.inference_mode() -def apply_correction(self, spectrum_image: Dataset3d) -> Dataset3d: - """Apply one learned spatial field to every spectrum-image channel. - - Parameters - ---------- - spectrum_image : Dataset3d - Spectrum image with axes ``(scan_row, scan_col, energy)``. - - Returns - ------- - Dataset3d - Corrected spectrum image with unchanged calibration, units, metadata, - and spectral-axis ordering. - - Examples - -------- - >>> corrected = drift.apply_correction(spectrum_image) - >>> corrected.shape == spectrum_image.shape - True - """ - if not getattr(self, "_reference_mode", False): - raise RuntimeError( - "apply_correction requires DriftCorrection.from_reference(...)." - ) - if not hasattr(self, "_initial_knots"): - raise RuntimeError( - "No drift field found. Call preprocess() and align_affine() first." - ) - if "affine" not in getattr(self, "_diagnostic_knots", {}): - raise RuntimeError( - "No fitted reference field found. Call align_affine() before " - "applying the correction." - ) - if not isinstance(spectrum_image, Dataset3d): - raise TypeError("spectrum_image must be a Dataset3d.") - dataset = spectrum_image - scan_rows, scan_cols, num_channels = dataset.shape - if (scan_rows, scan_cols) != tuple(self.images[1].shape): - raise ValueError( - "The spectrum-image scan axes must match the fitted alignment image; " - f"got {(scan_rows, scan_cols)} and {self.images[1].shape}." - ) - if self.knots[1].shape[-1] != 1: +def _raw_frame_displacement(self, image_index: int) -> torch.Tensor: + """Return a one-knot trajectory displacement in raw scan coordinates.""" + index = image_index % len(self.images) + if self.knots[index].shape[-1] != 1: raise ValueError( - "Spectrum-image correction currently requires number_knots=1." + "Multidimensional correction currently requires number_knots=1." ) - - device = self._device + scan_rows, scan_cols = self.images[index].shape delta_canvas_t = torch.as_tensor( - self.knots[1] - self._initial_knots[1], + self.knots[index] - self._initial_knots[index], dtype=torch.float32, - device=device, + device=self._device, )[:, :, 0] scan_fast_t = torch.as_tensor( - self.scan_fast[1], dtype=torch.float32, device=device + self.scan_fast[index], dtype=torch.float32, device=self._device ) scan_slow_t = torch.as_tensor( - self.scan_slow[1], dtype=torch.float32, device=device + self.scan_slow[index], dtype=torch.float32, device=self._device ) aspect = float(scan_rows - 1) / float(scan_cols - 1) if scan_cols > 1 else 1.0 determinant = ( @@ -92,11 +51,19 @@ def apply_correction(self, spectrum_image: Dataset3d) -> Dataset3d: -scan_slow_t[1] * delta_canvas_t[0] + scan_slow_t[0] * delta_canvas_t[1] ) / determinant - row_t = torch.arange(scan_rows, dtype=torch.float32, device=device) - col_t = torch.arange(scan_cols, dtype=torch.float32, device=device) + return torch.stack((drift_row_t, drift_col_t)) + + +def _scan_warp_grid(self, image_index: int) -> torch.Tensor: + """Return one raw-frame sampling grid for a solved scan trajectory.""" + index = image_index % len(self.images) + scan_rows, scan_cols = self.images[index].shape + drift_row_t, drift_col_t = _raw_frame_displacement(self, index) + row_t = torch.arange(scan_rows, dtype=torch.float32, device=self._device) + col_t = torch.arange(scan_cols, dtype=torch.float32, device=self._device) sample_row_t = row_t[:, None] - drift_row_t[:, None] sample_col_t = col_t[None, :] - drift_col_t[:, None] - grid_t = torch.stack( + return torch.stack( ( 2.0 * sample_col_t.expand(scan_rows, scan_cols) / (scan_cols - 1) - 1.0, 2.0 * sample_row_t.expand(scan_rows, scan_cols) / (scan_rows - 1) - 1.0, @@ -104,26 +71,184 @@ def apply_correction(self, spectrum_image: Dataset3d) -> Dataset3d: dim=-1, )[None] - corrected_array = np.empty(dataset.shape, dtype=np.float32) - chunk_size = min(num_channels, 64) + +@torch.inference_mode() +def _apply_scan_field( + self, + data: np.ndarray | torch.Tensor, + *, + image_index: int, + chunk_size: int | None, + output_dtype: np.dtype | torch.dtype | str | None, + mode: str = "bilinear", + output: np.ndarray | None = None, +) -> np.ndarray | torch.Tensor: + """Warp only the two leading scan axes of an array.""" + index = image_index % len(self.images) + scan_rows, scan_cols = self.images[index].shape + if data.ndim < 2 or tuple(data.shape[:2]) != (scan_rows, scan_cols): + raise ValueError( + "The leading scan axes must match the fitted alignment image; " + f"got {tuple(data.shape[:2])} and {(scan_rows, scan_cols)}." + ) + trailing_shape = tuple(data.shape[2:]) + num_channels = int(np.prod(trailing_shape, dtype=np.int64)) if trailing_shape else 1 + source_is_torch = isinstance(data, torch.Tensor) + source_dtype = data.dtype + if output is not None: + if not isinstance(output, np.ndarray) or output.shape != tuple(data.shape): + raise ValueError("output must be a NumPy array matching the input shape.") + output_flat = output.reshape(scan_rows, scan_cols, num_channels) + else: + output_flat = None + flat = ( + data.reshape(scan_rows, scan_cols, num_channels) + if source_is_torch + else torch.from_numpy( + np.ascontiguousarray(data).reshape(scan_rows, scan_cols, num_channels) + ) + ) + grid_t = _scan_warp_grid(self, index) + chunk_size = min(num_channels, 64 if chunk_size is None else int(chunk_size)) + if chunk_size < 1: + raise ValueError(f"chunk_size must be positive, got {chunk_size}.") + corrected_t = None + if output_flat is None: + corrected_t = torch.empty( + (scan_rows, scan_cols, num_channels), + dtype=torch.float32, + device=self._device, + ) for start in range(0, num_channels, chunk_size): stop = min(start + chunk_size, num_channels) - channels_t = torch.as_tensor( - np.ascontiguousarray(dataset.array[:, :, start:stop]), + channels_t = flat[:, :, start:stop].to( + device=self._device, dtype=torch.float32, - device=device, ).permute(2, 0, 1)[None] - corrected_t = torch.nn.functional.grid_sample( + corrected_block_t = torch.nn.functional.grid_sample( channels_t, grid_t, - mode="bilinear", + mode=mode, padding_mode="border", align_corners=True, )[0].permute(1, 2, 0) - corrected_array[:, :, start:stop] = corrected_t.cpu().numpy() + if output_flat is None: + corrected_t[:, :, start:stop] = corrected_block_t + else: + output_dtype_np = output_flat.dtype + if np.issubdtype(output_dtype_np, np.integer): + limits = np.iinfo(output_dtype_np) + corrected_block_t = corrected_block_t.round().clamp( + limits.min, + limits.max, + ) + output_flat[:, :, start:stop] = corrected_block_t.cpu().numpy().astype( + output_dtype_np, + copy=False, + ) + if output_flat is not None: + return output + assert corrected_t is not None + corrected_t = corrected_t.reshape((scan_rows, scan_cols, *trailing_shape)) + preserve_dtype = output_dtype == "same" + requested_dtype = source_dtype if preserve_dtype else output_dtype + if requested_dtype is not None: + if source_is_torch or isinstance(requested_dtype, torch.dtype): + torch_dtype = source_dtype if preserve_dtype else requested_dtype + if not isinstance(torch_dtype, torch.dtype): + torch_dtype = torch.from_numpy( + np.empty((), dtype=np.dtype(torch_dtype)) + ).dtype + if not torch_dtype.is_floating_point: + limits = torch.iinfo(torch_dtype) + corrected_t = corrected_t.round().clamp(limits.min, limits.max) + corrected_t = corrected_t.to(torch_dtype) + else: + numpy_dtype = np.dtype(source_dtype if preserve_dtype else requested_dtype) + if np.issubdtype(numpy_dtype, np.integer): + limits = np.iinfo(numpy_dtype) + corrected_t = corrected_t.round().clamp(limits.min, limits.max) + return corrected_t.cpu().numpy().astype(numpy_dtype, copy=False) + if source_is_torch: + return corrected_t.to(data.device) + return corrected_t.cpu().numpy() + +@torch.inference_mode() +def apply_correction( + self, + data: Dataset3d | np.ndarray | torch.Tensor | None = None, + image_index: int = -1, + *, + mode: str = "bilinear", + chunk_size: int | None = None, + output_dtype: np.dtype | torch.dtype | str | None = None, + output_device: str | torch.device | None = None, + output: np.ndarray | None = None, + verbose: bool = True, +): + """Apply one learned spatial field without mixing trailing data axes. + + Parameters + ---------- + data : Dataset3d, numpy.ndarray, torch.Tensor, or None + Data with leading ``(scan_row, scan_col)`` axes. ``None`` uses a raw + dataset retained by :meth:`DriftCorrection.from_4dstem`. + image_index : int, default -1 + Fitted scan trajectory to apply. + chunk_size : int or None, default None + Number of spectral or detector channels corrected per GPU batch. + output_dtype : numpy dtype, torch dtype, "same", or None + Output dtype. The default returns floating-point interpolation. + output_device : str, torch.device, or None + Device for a tensor result. NumPy inputs return NumPy by default. + output : numpy.ndarray or None + Optional preallocated NumPy destination. + verbose : bool, default True + Reserved for progress reporting in larger workflows. + + Returns + ------- + Dataset3d, numpy.ndarray, or torch.Tensor + Corrected data with unchanged shape and trailing-axis order. + + Examples + -------- + >>> corrected = drift.apply_correction(spectrum_image, image_index=1) + """ + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "No drift field found. Call preprocess() and align_affine() first." + ) + if "affine" not in getattr(self, "_diagnostic_knots", {}): + raise RuntimeError( + "No fitted reference field found. Call align_affine() before " + "applying the correction." + ) + if mode not in {"bilinear", "bicubic"}: + raise ValueError(f"mode must be 'bilinear' or 'bicubic', got {mode!r}.") + if data is None: + datasets = getattr(self, "_datasets", None) + if datasets is None: + raise ValueError("No data supplied and no 4D-STEM dataset is retained.") + data = datasets[image_index % len(datasets)] + dataset = data if isinstance(data, Dataset3d) else None + array = dataset.array if dataset is not None else data + corrected_array = _apply_scan_field( + self, + array, + image_index=image_index, + chunk_size=chunk_size, + output_dtype=output_dtype, + mode=mode, + output=output, + ) + if output is None and output_device is not None: + corrected_array = torch.as_tensor(corrected_array).to(output_device) + if dataset is None: + return corrected_array corrected = Dataset3d.from_array( - corrected_array, + np.asarray(corrected_array), name=f"drift-corrected {dataset.name}", origin=dataset.origin.copy(), sampling=dataset.sampling.copy(), diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index 63497337..c0beb7e1 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -1,6 +1,7 @@ from collections.abc import Sequence from typing import Self +import numpy as np import torch from numpy.typing import NDArray @@ -14,8 +15,8 @@ from quantem.core.utils.validators import ensure_valid_array from . import apply as drift_apply +from . import fourdstem, preparation from . import plot as drift_plot -from . import preparation from .core import affine, nonrigid @@ -124,6 +125,7 @@ def __init__( self, images: list[Dataset2d], scan_direction_degrees: NDArray, + device: str | None = None, _token: object | None = None, ): if _token is not self._token: @@ -135,7 +137,7 @@ def __init__( self.scan_direction_degrees = ensure_valid_array(scan_direction_degrees, ndim=1) self._reference_mode = False - device, _ = validate_device(None) + device, _ = validate_device(device) self._device = device self._dtype = torch.float32 @@ -157,12 +159,14 @@ def from_data( cls, images: list[Dataset2d] | list[NDArray] | Dataset3d | NDArray, scan_direction_degrees: list[float] | NDArray, + device: str | None = None, ) -> Self: validated_images = validate_list_of_dataset2d(images) return cls( images=validated_images, scan_direction_degrees=scan_direction_degrees, + device=device, _token=cls._token, ) @@ -173,6 +177,7 @@ def from_reference( alignment_image: Dataset2d | NDArray, *, scan_direction_degrees: float | None = None, + device: str | None = None, ) -> Self: """Fit a fixed-reference drift field from two-dimensional images. @@ -240,10 +245,100 @@ def from_reference( result = cls.from_data( images=[reference_image, alignment_image], scan_direction_degrees=[scan_direction_degrees, scan_direction_degrees], + device=device, ) result._reference_mode = True return result + @classmethod + def from_4dstem( + cls, + *datasets, + scan_direction_degrees: Sequence[float] | NDArray | None = None, + scan_sampling: float | tuple[float, float] | None = None, + scan_units: str | tuple[str, str] = "pixels", + device: str | None = None, + ) -> Self: + """Fit one correction from virtual images of two 4D-STEM scans. + + Only scan axes are transformed when the fitted fields are later + applied. Detector coordinates and native detector sampling are never + rotated, resampled, or binned by this constructor. + """ + if len(datasets) != 2: + raise ValueError( + f"from_4dstem requires exactly two datasets, got {len(datasets)}." + ) + raw_datasets = [fourdstem.data_array(dataset) for dataset in datasets] + if any(data.ndim != 4 for data in raw_datasets): + raise ValueError( + "from_4dstem expects scan-axis-leading 4-D inputs; got " + f"{[tuple(data.shape) for data in raw_datasets]}." + ) + detector_shapes = [tuple(data.shape[2:]) for data in raw_datasets] + if detector_shapes[0] != detector_shapes[1]: + raise ValueError( + "4D-STEM inputs must share one detector shape; got " + f"{detector_shapes}." + ) + if scan_direction_degrees is None: + angles = [ + getattr(dataset, "metadata", {}).get("scan_rotation_deg") + for dataset in datasets + ] + if any(angle is None for angle in angles): + raise TypeError( + "scan_direction_degrees is required when 4D-STEM inputs " + "do not carry scan_rotation_deg metadata." + ) + scan_direction_degrees = angles + virtual_images = [] + for index, (dataset, raw_dataset) in enumerate( + zip(datasets, raw_datasets, strict=True) + ): + image = Dataset2d.from_array( + fourdstem.integrate_virtual_detector(raw_dataset), + name=f"4D-STEM virtual image {index}", + ) + if scan_sampling is None and hasattr(dataset, "sampling"): + image.sampling = np.asarray(dataset.sampling[:2], dtype=float) + if isinstance(scan_units, str) and hasattr(dataset, "units"): + image.units = list(dataset.units[:2]) + image.metadata.update(dict(getattr(dataset, "metadata", {}))) + virtual_images.append(image) + if scan_sampling is not None: + sampling = ( + (float(scan_sampling), float(scan_sampling)) + if np.isscalar(scan_sampling) + else tuple(float(value) for value in scan_sampling) + ) + units = ( + (scan_units, scan_units) + if isinstance(scan_units, str) + else tuple(scan_units) + ) + for image in virtual_images: + image.sampling = sampling + image.units = units + result = cls.from_data( + virtual_images, + scan_direction_degrees=scan_direction_degrees, + device=device, + ) + result._datasets = raw_datasets + result._datasets_consumed = False + return result + + @property + def imgs(self): + """Compatibility alias for the scan images used by the solver.""" + return self.images + + @property + def device(self) -> str: + """Compute device selected for drift fitting and data propagation.""" + return str(self._device) + preprocess = preparation.preprocess align_translation = preparation.align_translation align_affine = affine.align_affine @@ -259,6 +354,24 @@ def from_reference( generate_corrected = drift_apply.generate_corrected generate_corrected_image = drift_apply.generate_corrected_image apply_correction = drift_apply.apply_correction + drift_field = fourdstem.drift_field + probe_positions = fourdstem.probe_positions + corrected_virtual_images = fourdstem.corrected_virtual_images + regional_diffraction_patterns = fourdstem.regional_diffraction_patterns + corrected_4dstem = fourdstem.corrected_4dstem + integrate_virtual_detector = staticmethod(fourdstem.integrate_virtual_detector) + + def save(self, path, mode="w", store="auto", skip=(), compression_level=4): + """Serialize fitted state without embedding raw multidimensional data.""" + if isinstance(skip, (str, type)): + skip = [skip] + super().save( + path, + mode=mode, + store=store, + skip=[*skip, "_datasets"], + compression_level=compression_level, + ) calculate_error = drift_apply.calculate_error plot_transformed_images = drift_plot.plot_transformed_images plot_convergence = drift_plot.plot_convergence diff --git a/src/quantem/imaging/drift/fourdstem.py b/src/quantem/imaging/drift/fourdstem.py new file mode 100644 index 00000000..7dcfe2bb --- /dev/null +++ b/src/quantem/imaging/drift/fourdstem.py @@ -0,0 +1,782 @@ +"""Virtual-detector and paired-dataset products for 4D-STEM correction.""" +from collections.abc import Mapping +from dataclasses import dataclass + +import numpy as np +import torch +from tqdm.auto import tqdm + +from quantem.imaging.drift.apply import ( + _apply_scan_field, + _raw_frame_displacement, +) + + +def data_array(dataset): + """Return the resident array or tensor from a QuantEM dataset wrapper.""" + if isinstance(dataset, (np.ndarray, torch.Tensor)): + return dataset + array = getattr(dataset, "array", None) + if array is not None: + return array + tensor = getattr(dataset, "_tensor", None) + if tensor is not None: + return tensor + raise TypeError( + "4D-STEM input must be an array, tensor, or QuantEM dataset with " + "resident .array/.tensor data." + ) + + +def padding_offset( + canvas_shape: tuple[int, int], + scan_shape: tuple[int, int], +) -> tuple[float, float]: + """Return the centered scan offset in the padded solver canvas.""" + return ( + (canvas_shape[0] - scan_shape[0]) / 2.0, + (canvas_shape[1] - scan_shape[1]) / 2.0, + ) + + +@dataclass +class CorrectionResult: + """Container returned by 0/90 4D-STEM collection correction. + + This result represents the scan-derived corrected coordinate system: + both input 4D-STEM datasets are treated as drifted scans and corrected toward a + shared consensus frame before optional diffraction-pattern-level merge. + + Attributes + ---------- + corrected_4dstem_0, corrected_4dstem_1 : np.ndarray | torch.Tensor + Per-side drift-corrected 4D-STEM datasets, scan-axis-leading layout. + Dataset 1 has already been oriented into dataset 0's display frame. + corrected_4dstem : np.ndarray | torch.Tensor | None + Diffraction-pattern-level average of ``corrected_4dstem_0`` and the + oriented ``corrected_4dstem_1``. ``None`` when ``merge=False``. + scalar_corrected_vdf : np.ndarray | None + Scan-derived corrected VDF computed by correcting the raw alignment + VDFs with the same operator used for 4D-STEM channels. This is not an + external ground truth; it is the scalar virtual-image result implied + by the learned scan drift fields. + """ + corrected_4dstem_0: np.ndarray | torch.Tensor + corrected_4dstem_1: np.ndarray | torch.Tensor + corrected_4dstem: np.ndarray | torch.Tensor | None = None + scalar_corrected_vdf: np.ndarray | None = None + + +def _rot90_to_image0_frame( + dc, image_index: int = 1, reference_index: int = 0 +) -> int: + """Return the scan-axis rot90 count to show ``image_index`` like the reference. + + Two 0/90 scans differ by a 90-degree scan-axis rotation, so displaying one in + the other's frame is a ``rot90`` whose count is the signed angle difference in + quarter turns. ``reference_index`` defaults to image 0 (the consensus frame); + a caller comparing against a different reference passes its index. + """ + delta = float( + dc.scan_direction_degrees[image_index] - dc.scan_direction_degrees[reference_index] + ) + return (-int(round(delta / 90.0))) % 4 + + +def integrate_virtual_detector( + dataset, + detector_mask: np.ndarray | torch.Tensor | None = None, + *, + reduce: str = "mean", +) -> np.ndarray: + """Integrate a virtual image from a scan-axis-leading dataset. + + Sums (or averages) the selected trailing detector/channel pixels for every + scan position, producing a 2-D scan image. With ``detector_mask=None`` it + integrates the *whole* detector (a full/total virtual image, bright-field + dominated for thin samples); pass a disk mask for bright field or an annulus + for dark field. This is the single canonical virtual-image reduction for the + drift package; every other virtual-image helper delegates here. + + QuantEM's detector backend selects the resident NumPy, Torch, CuPy, CUDA, + or MPS reduction path. Three-dimensional spectrum images are treated as a + one-row detector so the same backend also integrates their channel axis. + + Parameters + ---------- + dataset : ndarray or torch.Tensor, shape ``(H, W, ...channels)`` + 3-D/4-D dataset with scan axes first. numpy may be a ``np.memmap``. + detector_mask : ndarray or torch.Tensor, optional + Boolean mask over the trailing detector / channel axes. ``None`` + selects every channel. + reduce : {"mean", "sum"} + Average or sum the selected detector pixels. + Returns + ------- + numpy.ndarray, shape ``(H, W)``, dtype float32 + + Examples + -------- + >>> adf = integrate_virtual_detector(data, detector_mask=annulus) + """ + dataset = data_array(dataset) + if reduce not in {"mean", "sum"}: + raise ValueError(f"reduce must be 'mean' or 'sum', got {reduce!r}") + scan_shape = tuple(int(value) for value in dataset.shape[:2]) + detector_shape = tuple(int(value) for value in dataset.shape[2:]) + if len(detector_shape) == 1: + detector_shape = (1, detector_shape[0]) + dataset = dataset.reshape(-1, *detector_shape) + elif len(detector_shape) != 2: + raise ValueError( + "integrate_virtual_detector expects (row, col, channel) or " + f"(row, col, detector_row, detector_col), got {tuple(dataset.shape)}" + ) + mask = ( + np.ones(detector_shape, dtype=bool) + if detector_mask is None + else to_numpy(detector_mask, dtype=bool).reshape(detector_shape) + ) + num_selected = int(mask.sum()) + if num_selected == 0: + raise ValueError("detector_mask selects zero detector pixels") + if isinstance(dataset, torch.Tensor): + mask_t = torch.as_tensor(mask, device=dataset.device) + flattened = dataset.reshape(*scan_shape, -1).to(torch.float32) + image = flattened[..., mask_t.reshape(-1)].sum(-1) + image = image.detach().cpu().numpy() + else: + array = dataset.get() if hasattr(dataset, "get") else np.asarray(dataset) + image = array.reshape(*scan_shape, -1)[..., mask.reshape(-1)].sum( + -1, + dtype=np.float64, + ).astype(np.float32) + return image / num_selected if reduce == "mean" else image + + +def drift_field(self, idx: int) -> torch.Tensor: + """Return fitted raw-frame drift for one scan in ``(row, col)`` order. + + A one-knot model returns ``(2, scan_rows)`` because its displacement is + constant along each scanline. Multi-knot models return + ``(2, scan_rows, scan_columns)``. + + Parameters + ---------- + idx : int + Scan index in the correction pair. + + Returns + ------- + torch.Tensor + Row and column displacement in raw scan coordinates. + + Examples + -------- + >>> field_90 = drift.drift_field(1) + """ + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "drift_field() requires preprocess() and correct_affine() first. " + "Run dc.preprocess().correct_affine() (and optionally " + ".correct_nonrigid()) before drift_field()." + ) + return _raw_frame_displacement(self, idx) + + +def probe_positions( + self, + image_index: int = 0, + *, + corrected: bool = True, + strip_padding: bool = True, + plot: bool = True, + stride: int = 16, +) -> np.ndarray: + """Return nominal or drift-updated probe positions for one scan image. + + Positions retain the raw diffraction-pattern indexing and use ``(row, + col)`` coordinates in the shared corrected frame. This lets iterative + ptychography consume corrected positions without interpolating detector + data. + + Parameters + ---------- + image_index : int, default 0 + Scan image or 4D-STEM acquisition to describe. + corrected : bool, default True + Return fitted positions instead of the nominal scan grid. + strip_padding : bool, default True + Express positions in the original image-0 frame instead of the padded + solver canvas. + plot : bool, default True + Draw the nominal and corrected positions for inspection. + stride : int, default 16 + Subsampling used only by the plot. + + Returns + ------- + numpy.ndarray + Position array with shape ``(scan_rows, scan_columns, 2)``. + + Examples + -------- + >>> positions = drift.probe_positions(image_index=0, plot=False) + """ + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "probe_positions() requires preprocess() first. Run " + "dc.preprocess() before exporting nominal or corrected " + "probe positions." + ) + index = image_index % len(self.images) + knots = ( + self.knots[index] + if corrected + else self._initial_knots[index] + ) + row, column = self.interpolator[index].transform_coordinates(knots) + positions = np.stack([row, column], axis=-1).astype(np.float32) + if strip_padding: + pad_row, pad_column = padding_offset( + (self.shape[1], self.shape[2]), + self.images[0].shape[:2], + ) + positions -= np.asarray([pad_row, pad_column], dtype=np.float32) + result = positions + if plot: + import matplotlib.pyplot as plt + + nominal = self.probe_positions( + index, + corrected=False, + strip_padding=strip_padding, + plot=False, + ) + plt.plot( + nominal[::stride, ::stride, 1], + nominal[::stride, ::stride, 0], + ".", + alpha=0.4, + label="nominal", + ) + plt.plot( + result[::stride, ::stride, 1], + result[::stride, ::stride, 0], + ".", + alpha=0.6, + label="corrected", + ) + plt.gca().invert_yaxis() + plt.gca().set_aspect("equal") + plt.legend() + return result + + +@torch.inference_mode() +def corrected_virtual_images( + self, + image_0, + image_1, + *, + output_frame: str = "scan", +) -> dict[str, np.ndarray]: + """Correct two scalar virtual images like matching 4D-STEM channels. + + Each scalar image is treated as a one-channel dataset with scan axes first, + corrected with the same ``grid_sample`` operator used for every diffraction + pixel, and image 1 is oriented into image 0's display frame before the + average. The returned ``corrected_image`` should therefore match integrating + the same virtual detector from ``corrected_4dstem()`` output, + up to output quantization. + + Parameters + ---------- + image_0, image_1 : array-like + Scalar virtual images from the two 4D-STEM acquisitions. Their scan + shapes must match the images used to solve the correction. + output_frame : {"scan", "canvas"}, default "scan" + Return each corrected image in image 0's scan frame or on the shared + padded correction canvas. Canvas output also includes per-scan and + combined coverage arrays and averages only scans that cover each + output pixel. + + Returns + ------- + dict[str, np.ndarray] + The merged ``corrected_image`` and the separately corrected + ``corrected_image_0`` and ``corrected_image_1``, all in image 0's scan + frame by default. Canvas output also contains ``coverage_image``, + ``coverage_image_0``, and ``coverage_image_1``. + + Raises + ------ + ValueError + If ``output_frame`` is not ``"scan"`` or ``"canvas"``. + + Examples + -------- + >>> images = drift.corrected_virtual_images(vdf_0, vdf_90) + >>> corrected_vdf = images["corrected_image"] + >>> canvas = drift.corrected_virtual_images( + ... vdf_0, vdf_90, output_frame="canvas" + ... ) + """ + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "corrected_virtual_images() requires preprocess() and " + "correct_affine() first." + ) + if len(self.images) != 2: + raise ValueError( + "corrected_virtual_images() expects exactly two scan images" + ) + if output_frame not in {"scan", "canvas"}: + raise ValueError( + "output_frame must be 'scan' or 'canvas'; " + f"got {output_frame!r}" + ) + images = [ + np.asarray(image_0, dtype=np.float32), + np.asarray(image_1, dtype=np.float32), + ] + if ( + images[0].shape != self.images[0].shape + or images[1].shape != self.images[1].shape + ): + raise ValueError( + "virtual image shapes must match the raw scan images used for drift " + f"alignment: got {images[0].shape}, {images[1].shape}; expected " + f"{self.images[0].shape}, {self.images[1].shape}" + ) + + if output_frame == "canvas": + canvas_shape = tuple(int(value) for value in self.shape[-2:]) + components = [] + coverages = [] + for image_index, image in enumerate(images): + corrected, coverage = self.interpolator[image_index].warp_image( + image, + self.knots[image_index], + output_shape=canvas_shape, + kde_sigma=self.kde_sigma, + pad_value=0.0, + ) + components.append(corrected) + coverages.append(coverage) + + valid = [coverage >= 1e-3 for coverage in coverages] + contribution_count = valid[0].astype(np.float32) + valid[1].astype( + np.float32 + ) + merged = np.divide( + components[0] * valid[0] + components[1] * valid[1], + contribution_count, + out=np.zeros(canvas_shape, dtype=np.float32), + where=contribution_count > 0, + ) + return { + "corrected_image": to_numpy(merged, dtype=np.float32), + "corrected_image_0": to_numpy(components[0], dtype=np.float32), + "corrected_image_1": to_numpy(components[1], dtype=np.float32), + "coverage_image": to_numpy( + np.maximum(coverages[0], coverages[1]), + dtype=np.float32, + ), + "coverage_image_0": to_numpy(coverages[0], dtype=np.float32), + "coverage_image_1": to_numpy(coverages[1], dtype=np.float32), + } + + components = [] + for image_index, image in enumerate(images): + corrected = _apply_scan_field( + self, + image[..., None], + image_index=image_index, + chunk_size=1, + output_dtype=np.float32, + )[..., 0] + if image_index == 1: + rot_k = _rot90_to_image0_frame(self, image_index=1) + if rot_k: + corrected = np.rot90(corrected, k=rot_k, axes=(0, 1)).copy() + components.append(corrected) + + merged = (components[0] + components[1]) * 0.5 + + return { + "corrected_image": to_numpy(merged, dtype=np.float32), + "corrected_image_0": to_numpy(components[0], dtype=np.float32), + "corrected_image_1": to_numpy(components[1], dtype=np.float32), + } + + +def regional_diffraction_patterns( + self, + regions: Mapping[str, tuple[float, float]], + *, + radius_px: float = 4.0, + datasets=None, + stages: tuple[str, ...] = ("initial", "corrected"), +) -> dict[str, object]: + """Average diffraction patterns from named specimen regions. + + Region membership is evaluated in the shared scan frame using either the + nominal or drift-corrected probe positions. Detector pixels are never + interpolated: the method averages the original diffraction patterns whose + probe positions fall inside each circular region. This makes before/after + comparisons test spatial indexing without changing diffraction detail. + + Parameters + ---------- + regions : mapping of str to tuple of float + Named region centers in shared ``(row, column)`` scan pixels. + radius_px : float, default 4.0 + Circular region radius in scan pixels. + datasets : sequence of two arrays, optional + Raw scan-axis-leading 4D-STEM datasets. When omitted, use the datasets + retained by ``DriftCorrection.from_4dstem``. Pass this explicitly when + working from a saved correction because serialization intentionally + excludes multi-gigabyte diffraction cubes. + stages : tuple of str, default ("initial", "corrected") + Probe-position stages to compare. Each entry must be ``"initial"`` or + ``"corrected"``. + + Returns + ------- + dict[str, object] + ``patterns`` has shape ``(stage, region, scan, detector_row, + detector_column)``. The result also includes ``sample_counts``, + ``region_names``, ``region_centers_px``, ``radius_px``, ``stages``, and + ``scan_direction_degrees``. + + Examples + -------- + >>> regions = {"Au": (121, 220), "support region": (25, 109)} + >>> comparison = drift.regional_diffraction_patterns(regions, radius_px=4) + >>> comparison["patterns"].shape + (2, 2, 2, 192, 192) + """ + if not isinstance(regions, Mapping) or not regions: + raise ValueError("regions must be a non-empty name-to-(row, column) mapping") + region_names = tuple(regions) + if any(not isinstance(name, str) or not name for name in region_names): + raise ValueError("every region name must be a non-empty string") + region_centers = np.asarray(list(regions.values()), dtype=np.float32) + if region_centers.shape != (len(region_names), 2) or not np.isfinite( + region_centers + ).all(): + raise ValueError( + "region centers must be finite (row, column) pairs; " + f"got shape {region_centers.shape}" + ) + radius = float(radius_px) + if not np.isfinite(radius) or radius <= 0: + raise ValueError(f"radius_px must be positive and finite, got {radius_px!r}") + + requested_stages = tuple(stages) + valid_stages = {"initial", "corrected"} + invalid_stages = [stage for stage in requested_stages if stage not in valid_stages] + if not requested_stages or invalid_stages: + raise ValueError( + "stages must contain 'initial' and/or 'corrected'; " + f"got {requested_stages!r}" + ) + if len(set(requested_stages)) != len(requested_stages): + raise ValueError(f"stages must not contain duplicates, got {requested_stages!r}") + + source_datasets = getattr(self, "_datasets", None) if datasets is None else datasets + if source_datasets is None: + raise RuntimeError( + "regional_diffraction_patterns() needs the two raw 4D-STEM datasets. " + "Pass datasets=(scan_0, scan_90) when using a saved correction." + ) + source_datasets = tuple(data_array(dataset) for dataset in source_datasets) + if len(source_datasets) != 2 or any(dataset is None for dataset in source_datasets): + raise ValueError( + "regional_diffraction_patterns() expects exactly two raw 4D-STEM " + f"datasets, got {len(source_datasets)}" + ) + scan_shapes = [tuple(int(value) for value in data.shape[:2]) for data in source_datasets] + expected_shapes = [ + tuple(int(value) for value in image.shape) for image in self.images + ] + detector_shapes = [tuple(int(value) for value in data.shape[2:]) for data in source_datasets] + if scan_shapes != expected_shapes: + raise ValueError( + "dataset scan shapes must match the images used for correction: " + f"got {scan_shapes}, expected {expected_shapes}" + ) + if len(detector_shapes[0]) != 2 or detector_shapes[0] != detector_shapes[1]: + raise ValueError( + "datasets must share one 2-D detector shape, got " + f"{detector_shapes}" + ) + + patterns = np.empty( + ( + len(requested_stages), + len(region_names), + 2, + *detector_shapes[0], + ), + dtype=np.float32, + ) + sample_counts = np.empty( + (len(requested_stages), len(region_names), 2), + dtype=np.int32, + ) + radius_squared = radius**2 + for stage_index, stage in enumerate(requested_stages): + corrected = stage == "corrected" + for scan_index, dataset in enumerate(source_datasets): + positions = self.probe_positions( + scan_index, + corrected=corrected, + strip_padding=True, + plot=False, + ) + for region_index, (name, center) in enumerate( + zip(region_names, region_centers, strict=True) + ): + mask = ( + (positions[..., 0] - center[0]) ** 2 + + (positions[..., 1] - center[1]) ** 2 + <= radius_squared + ) + count = int(mask.sum()) + if count == 0: + raise ValueError( + f"region {name!r} at ({center[0]:g}, {center[1]:g}) with " + f"radius_px={radius:g} selects no samples for {stage!r} " + f"scan {scan_index}" + ) + coordinates = np.argwhere(mask) + row_start, column_start = coordinates.min(axis=0) + row_stop, column_stop = coordinates.max(axis=0) + 1 + local_mask = mask[ + row_start:row_stop, + column_start:column_stop, + ] + block = dataset[ + int(row_start):int(row_stop), + int(column_start):int(column_stop), + ] + if isinstance(block, torch.Tensor): + # CUDA cannot boolean-index uint16 tensors. Convert only + # this small region, never the full diffraction cube. + local_mask_t = torch.as_tensor(local_mask, device=block.device) + pattern = block.to(torch.float32)[local_mask_t].mean(dim=0) + pattern = to_numpy(pattern, dtype=np.float32) + else: + if hasattr(block, "get"): + block = block.get() + pattern = np.asarray(block, dtype=np.float32)[local_mask].mean( + axis=0, + dtype=np.float32, + ) + patterns[stage_index, region_index, scan_index] = pattern + sample_counts[stage_index, region_index, scan_index] = count + + return { + "patterns": patterns, + "sample_counts": sample_counts, + "region_names": region_names, + "region_centers_px": region_centers, + "radius_px": radius, + "stages": requested_stages, + "scan_direction_degrees": np.asarray( + self.scan_direction_degrees, + dtype=np.float32, + ), + } + + +def corrected_4dstem( + self, + *, + mode: str = "bilinear", + chunk_size: int | None = None, + merge: bool = True, + verbose: bool = True, + output_0: np.ndarray | None = None, + output_1: np.ndarray | None = None, + output_dtype: torch.dtype | np.dtype | str | None = None, + output_device: str | torch.device | None = None, +) -> CorrectionResult: + """Correct and optionally merge a 0/90 4D-STEM acquisition pair. + + Each acquisition receives its learned scan drift before the second scan is + rotated into the first scan's frame. Preallocated NumPy or memmap outputs + keep large detector datasets from requiring another full-size allocation. + + Parameters + ---------- + mode : str, default "bilinear" + Interpolation used along the scan axes. + chunk_size : int or None, default None + Detector channels corrected per batch. ``None`` selects automatically. + merge : bool, default True + Average the two corrected acquisitions in their shared frame. + verbose : bool, default True + Show progress for chunked correction and merging. + output_0, output_1 : numpy.ndarray or None, default None + Preallocated outputs for the two corrected acquisitions. + output_dtype : torch.dtype, numpy dtype, str, or None, default None + Output numeric type. Use ``"same"`` to preserve the input type. + output_device : str, torch.device, or None, default None + Device holding returned arrays when no preallocated output is supplied. + + Returns + ------- + CorrectionResult + Corrected acquisitions and their optional merged dataset. + + Examples + -------- + >>> result = drift.corrected_4dstem(chunk_size=64) + >>> merged = result.corrected_4dstem + """ + if getattr(self, "_datasets", None) is None or self._reference_mode: + raise RuntimeError( + "corrected_4dstem() requires DriftCorrection.from_4dstem(data_0, " + "data_1, ...). For reference-mode EDS/EELS/4D-STEM, use corrected()." + ) + datasets = self._datasets + if self._datasets_consumed: + raise RuntimeError( + "Raw datasets were already released to free device memory " + "during a prior corrected call. Construct a new " + "DriftCorrection to re-correct." + ) + if len(datasets) < 2: + raise ValueError( + f"Need at least 2 datasets for scan collection correction, " + f"got {len(datasets)}" + ) + + # When inputs are device-resident, release each raw dataset as soon as + # its corrected output exists; otherwise we hold four full datasets + # simultaneously, which exceeds device memory for multi-GB scan collections. + inputs_on_device = ( + isinstance(datasets[0], torch.Tensor) and datasets[0].is_cuda + and isinstance(datasets[1], torch.Tensor) and datasets[1].is_cuda + ) + + corrected_4dstem_0 = _apply_scan_field( + self, + datasets[0], + image_index=0, + mode=mode, + chunk_size=chunk_size, + output_dtype=output_dtype, + output=output_0, + ) + if output_device is not None and output_0 is None: + corrected_4dstem_0 = torch.as_tensor(corrected_4dstem_0).to(output_device) + if inputs_on_device: + self._datasets[0] = None + torch.cuda.empty_cache() + corrected_4dstem_1 = _apply_scan_field( + self, + datasets[1], + image_index=1, + mode=mode, + chunk_size=chunk_size, + output_dtype=output_dtype, + output=output_1, + ) + if output_device is not None and output_1 is None: + corrected_4dstem_1 = torch.as_tensor(corrected_4dstem_1).to(output_device) + if inputs_on_device: + self._datasets[1] = None + self._datasets_consumed = True + torch.cuda.empty_cache() + + rot_k = _rot90_to_image0_frame(self, image_index=1) + if rot_k: + if isinstance(corrected_4dstem_1, torch.Tensor): + corrected_4dstem_1 = torch.rot90( + corrected_4dstem_1, k=rot_k, dims=(0, 1), + ) + else: + corrected_4dstem_1 = np.rot90( + corrected_4dstem_1, k=rot_k, axes=(0, 1), + ) + + corrected_4dstem = None + if merge: + if corrected_4dstem_0.shape != corrected_4dstem_1.shape: + raise ValueError( + f"Cannot merge: corrected_4dstem_0 shape {corrected_4dstem_0.shape} " + f"!= corrected_4dstem_1 shape {corrected_4dstem_1.shape}. " + f"Scan collection must have compatible scan dimensions " + f"after correction and scan-angle rotation." + ) + Hm = corrected_4dstem_0.shape[0] + row_block = max(1, min(32, Hm)) + row_starts = range(0, Hm, row_block) + merge_progress = tqdm( + total=Hm, + desc="Merging corrected scans", + unit="row", + disable=not verbose or len(row_starts) <= 1, + ) + if isinstance(corrected_4dstem_0, torch.Tensor): + corrected_4dstem = torch.empty_like( + corrected_4dstem_0, + dtype=torch.float32, + ) + for r0 in row_starts: + r1 = min(r0 + row_block, Hm) + corrected_4dstem[r0:r1] = ( + corrected_4dstem_0[r0:r1].to(torch.float32) + + corrected_4dstem_1[r0:r1].to(torch.float32) + ) * 0.5 + merge_progress.update(r1 - r0) + else: + corrected_4dstem = np.empty_like(corrected_4dstem_0, dtype=np.float32) + for r0 in row_starts: + r1 = min(r0 + row_block, Hm) + np.add( + corrected_4dstem_0[r0:r1], + corrected_4dstem_1[r0:r1], + out=corrected_4dstem[r0:r1], + dtype=np.float32, + ) + corrected_4dstem[r0:r1] *= 0.5 + merge_progress.update(r1 - r0) + merge_progress.close() + + # Extract raw VDFs from the stored alignment images. The scan collection + # reference is the scalar channel correction implied by the learned scan + # drift fields, not an external ground truth. + alignment_vdf_0 = np.asarray(self.images[0].array) + alignment_vdf_1 = np.asarray(self.images[1].array) + scalar_corrected_vdf = corrected_virtual_images( + self, + alignment_vdf_0, + alignment_vdf_1, + )["corrected_image"] + + return CorrectionResult( + corrected_4dstem=corrected_4dstem, + corrected_4dstem_0=corrected_4dstem_0, + corrected_4dstem_1=corrected_4dstem_1, + scalar_corrected_vdf=scalar_corrected_vdf, + ) + + +def to_numpy(array, *, dtype=None): + """Convert a device or host array to NumPy with optional dtype conversion.""" + if isinstance(array, torch.Tensor): + result = ( + array.detach().cpu().numpy() + if array.is_cuda or array.device.type == "mps" + else array.detach().numpy() + ) + elif hasattr(array, "get"): # CuPy ndarray + result = array.get() + else: + result = np.asarray(array) + return result.astype(dtype) if dtype is not None else result diff --git a/tests/imaging/test_drift4d.py b/tests/imaging/test_drift4d.py new file mode 100644 index 00000000..a4de54bc --- /dev/null +++ b/tests/imaging/test_drift4d.py @@ -0,0 +1,425 @@ +"""4D-STEM drift propagation keeps detector coordinates scientifically intact.""" + +import numpy as np +import pytest +import torch +from scipy.ndimage import gaussian_filter + +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.imaging.drift import CorrectionResult, DriftCorrection + + +def _orthogonal_4dstem_pair( + scan_size: int = 24, + detector_shape: tuple[int, int] = (4, 5), +): + """Build two orthogonal scans with fixed per-detector-pixel signatures.""" + rng = np.random.default_rng(14) + image_0 = gaussian_filter( + rng.normal(size=(scan_size, scan_size)).astype(np.float32), + 1.2, + ) + image_0 += np.linspace(0, 2, scan_size, dtype=np.float32)[:, None] + image_1 = np.rot90(image_0, k=-1).copy() + detector_offset = np.arange( + np.prod(detector_shape), + dtype=np.float32, + ).reshape(detector_shape) + cube_0 = image_0[..., None, None] + detector_offset + cube_1 = image_1[..., None, None] + detector_offset + return cube_0, cube_1, detector_offset + + +def _fit_small_pair(cube_0, cube_1): + drift = DriftCorrection.from_4dstem( + cube_0, + cube_1, + scan_direction_degrees=(0.0, 90.0), + scan_sampling=0.2, + scan_units="nm", + device="cpu", + ).preprocess( + pad_fraction=0.25, + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + drift.align_affine( + step=0.01, + num_tests=3, + refine=False, + max_image_shift=8, + chunk_size=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + return drift + + +def _add_known_raw_drift(drift, image_index, row, column): + """Move canvas knots by one requested raw-frame trajectory.""" + scan_rows, scan_cols = drift.images[image_index].shape + aspect = (scan_rows - 1) / (scan_cols - 1) + slow = drift.scan_slow[image_index] + fast = drift.scan_fast[image_index] + drift.knots[image_index][0, :, 0] += slow[0] * row + fast[0] * aspect * column + drift.knots[image_index][1, :, 0] += slow[1] * row + fast[1] * column + + +def test_virtual_detector_matches_numpy_and_torch_integer_inputs(): + """Virtual integration has the same exact integer sum on both backends.""" + data = np.arange(3 * 4 * 2 * 3, dtype=np.uint16).reshape(3, 4, 2, 3) + mask = np.array([[True, False, True], [False, True, False]]) + expected = data[..., mask].sum(axis=-1, dtype=np.uint64).astype(np.float32) + + numpy_image = DriftCorrection.integrate_virtual_detector( + data, + mask, + reduce="sum", + ) + torch_image = DriftCorrection.integrate_virtual_detector( + torch.from_numpy(data), + mask, + reduce="sum", + ) + + np.testing.assert_array_equal(numpy_image, expected) + np.testing.assert_array_equal(torch_image, expected) + + +def test_corrected_4dstem_transforms_scan_axes_not_detector_axes(): + """Every detector pixel receives one shared scan transform.""" + cube_0, cube_1, detector_offset = _orthogonal_4dstem_pair() + drift = _fit_small_pair(cube_0, cube_1) + line_drift = np.linspace(-1.5, 1.5, cube_0.shape[0]) + _add_known_raw_drift(drift, 0, line_drift, 0.5 * line_drift) + + result = drift.corrected_4dstem(chunk_size=5, verbose=False) + + assert isinstance(result, CorrectionResult) + assert result.corrected_4dstem_0.shape == cube_0.shape + assert result.corrected_4dstem_1.shape == cube_1.shape + assert result.corrected_4dstem.shape == cube_0.shape + for corrected in (result.corrected_4dstem_0, result.corrected_4dstem_1): + detector_difference = corrected - corrected[..., :1, :1] + np.testing.assert_allclose( + detector_difference, + np.broadcast_to( + detector_offset - detector_offset[0, 0], + corrected.shape, + ), + atol=2e-5, + ) + + +def test_regional_patterns_average_native_detector_samples(): + """Region membership changes, but diffraction pixels are not interpolated.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = DriftCorrection.from_4dstem( + cube_0, + cube_1, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + _add_known_raw_drift( + drift, + 0, + np.full(16, 2.0), + np.zeros(16), + ) + regions = {"feature": (8.0, 8.0)} + + comparison = drift.regional_diffraction_patterns( + regions, + radius_px=2.0, + stages=("initial", "corrected"), + ) + + assert comparison["patterns"].shape == (2, 1, 2, 4, 5) + for stage_index, corrected in enumerate((False, True)): + for scan_index, cube in enumerate((cube_0, cube_1)): + positions = drift.probe_positions( + scan_index, + corrected=corrected, + strip_padding=True, + plot=False, + ) + mask = ( + (positions[..., 0] - 8.0) ** 2 + + (positions[..., 1] - 8.0) ** 2 + <= 2.0**2 + ) + np.testing.assert_allclose( + comparison["patterns"][stage_index, 0, scan_index], + cube[mask].mean(axis=0, dtype=np.float32), + ) + assert comparison["sample_counts"][stage_index, 0, scan_index] == mask.sum() + assert not np.array_equal( + comparison["patterns"][0, 0, 0], + comparison["patterns"][1, 0, 0], + ) + + +def test_canvas_combination_uses_union_coverage(): + """The combined canvas retains pixels covered by either corrected scan.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = _fit_small_pair(cube_0, cube_1) + image_0 = drift.integrate_virtual_detector(cube_0, np.ones((4, 5), dtype=bool)) + image_1 = drift.integrate_virtual_detector(cube_1, np.ones((4, 5), dtype=bool)) + + result = drift.corrected_virtual_images( + image_0, + image_1, + output_frame="canvas", + ) + + expected_union = np.maximum( + result["coverage_image_0"], + result["coverage_image_1"], + ) + np.testing.assert_allclose(result["coverage_image"], expected_union) + either_scan = expected_union >= 1e-3 + assert np.count_nonzero(result["corrected_image"][either_scan]) > 0 + + +def test_saved_correction_accepts_explicit_4dstem_datasets(): + """Serialized corrections can analyze explicitly reattached raw cubes.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = DriftCorrection.from_4dstem( + cube_0, + cube_1, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + drift._datasets = None + + result = drift.regional_diffraction_patterns( + {"feature": (8.0, 8.0)}, + radius_px=2.0, + datasets=(cube_0, cube_1), + stages=("initial",), + ) + + assert result["patterns"].shape == (1, 1, 2, 4, 5) + + +def test_numpy_cube_can_return_torch_output_on_requested_device(): + """An explicit output device is honored without changing detector layout.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = _fit_small_pair(cube_0, cube_1) + + result = drift.corrected_4dstem( + merge=False, + output_device="cpu", + output_dtype=np.float32, + verbose=False, + ) + + assert isinstance(result.corrected_4dstem_0, torch.Tensor) + assert isinstance(result.corrected_4dstem_1, torch.Tensor) + assert result.corrected_4dstem_0.device.type == "cpu" + assert result.corrected_4dstem_0.shape == cube_0.shape + + +def test_dataset4dstem_metadata_supplies_rotation_and_scan_calibration(): + """QuantEM datasets retain scan metadata while exposing resident data.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + datasets = [] + for cube, angle in ((cube_0, 0.0), (cube_1, 90.0)): + dataset = Dataset4dstem.from_array( + cube, + sampling=(0.2, 0.3, 0.01, 0.01), + units=("nm", "nm", "1/nm", "1/nm"), + ) + dataset.metadata["scan_rotation_deg"] = angle + datasets.append(dataset) + + drift = DriftCorrection.from_4dstem(*datasets, device="cpu") + + np.testing.assert_allclose(drift.scan_direction_degrees, (0.0, 90.0)) + np.testing.assert_allclose(drift.images[0].sampling, (0.2, 0.3)) + assert drift.images[0].units == ["nm", "nm"] + assert drift._datasets[0] is datasets[0].array + + +def test_drift_field_reports_raw_components_for_rotated_scan(): + """A 90-degree scan reports raw row/column drift, not canvas axes.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = DriftCorrection.from_4dstem( + cube_0, + cube_1, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + expected_row = np.linspace(-2.0, 2.0, 16) + expected_column = np.linspace(1.0, -1.0, 16) + _add_known_raw_drift(drift, 1, expected_row, expected_column) + + field = drift.drift_field(1).cpu().numpy() + + np.testing.assert_allclose(field[0], expected_row, atol=2e-6) + np.testing.assert_allclose(field[1], expected_column, atol=2e-6) + + +def test_virtual_detector_integration_commutes_with_scan_warp(): + """Integrating detector pixels before or after correction is equivalent.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = _fit_small_pair(cube_0, cube_1) + line_drift = np.linspace(-1.0, 1.0, 16) + _add_known_raw_drift(drift, 0, line_drift, -0.25 * line_drift) + mask = np.zeros((4, 5), dtype=bool) + mask[1:3, 1:4] = True + virtual_0 = drift.integrate_virtual_detector(cube_0, mask, reduce="sum") + virtual_1 = drift.integrate_virtual_detector(cube_1, mask, reduce="sum") + + corrected_cube = drift.corrected_4dstem(merge=False, verbose=False) + corrected_virtual = drift.corrected_virtual_images(virtual_0, virtual_1) + + np.testing.assert_allclose( + drift.integrate_virtual_detector( + corrected_cube.corrected_4dstem_0, + mask, + reduce="sum", + ), + corrected_virtual["corrected_image_0"], + atol=2e-4, + ) + + +def test_preallocated_output_streams_without_full_device_allocation(monkeypatch): + """A supplied output receives detector chunks without a full device cube.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + drift = _fit_small_pair(cube_0, cube_1) + output_0 = np.empty_like(cube_0, dtype=np.float32) + output_1 = np.empty_like(cube_1, dtype=np.float32) + + def reject_full_allocation(*shape, **kwargs): + requested = ( + tuple(shape[0]) + if len(shape) == 1 and not isinstance(shape[0], int) + else tuple(shape) + ) + if requested in {cube_0.shape, (16, 16, 20)}: + raise AssertionError("attempted full corrected-cube allocation") + return original_empty(*shape, **kwargs) + + original_empty = torch.empty + monkeypatch.setattr(torch, "empty", reject_full_allocation) + corrected = drift.corrected_4dstem( + merge=False, + output_0=output_0, + output_1=output_1, + chunk_size=3, + verbose=False, + ) + + assert corrected.corrected_4dstem_0 is output_0 + assert np.shares_memory(corrected.corrected_4dstem_1, output_1) + assert np.isfinite(output_0).all() + assert np.isfinite(output_1).all() + + +def test_integer_merge_is_float32_and_backend_consistent(): + """Integer acquisitions retain half-counts in one float32 merge policy.""" + cube_0, cube_1, _ = _orthogonal_4dstem_pair(scan_size=16) + minimum = min(float(cube_0.min()), float(cube_1.min())) + cubes_np = [ + np.round((cube - minimum + 1.0) * 100).astype(np.uint16) + for cube in (cube_0, cube_1) + ] + drift_np = _fit_small_pair(*cubes_np) + result_np = drift_np.corrected_4dstem( + output_dtype="same", + verbose=False, + ) + cubes_torch = [torch.from_numpy(cube) for cube in cubes_np] + drift_torch = _fit_small_pair(*cubes_torch) + result_torch = drift_torch.corrected_4dstem( + output_dtype="same", + verbose=False, + ) + + assert result_np.corrected_4dstem.dtype == np.float32 + assert result_torch.corrected_4dstem.dtype == torch.float32 + np.testing.assert_allclose( + result_np.corrected_4dstem, + result_torch.corrected_4dstem.cpu().numpy(), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable") +def test_cuda_matches_cpu_for_uint32_native_detector_frames(): + """CUDA and CPU preserve one native 192-square detector field equally.""" + scan_size = 8 + detector_shape = (192, 192) + row = np.arange(scan_size, dtype=np.uint32)[:, None, None, None] + column = np.arange(scan_size, dtype=np.uint32)[None, :, None, None] + detector = np.arange( + np.prod(detector_shape), + dtype=np.uint32, + ).reshape(1, 1, *detector_shape) + cube_0 = row * 100_000 + column * 10_000 + detector + cube_1 = np.rot90(cube_0, k=-1, axes=(0, 1)).copy() + cpu = DriftCorrection.from_4dstem( + cube_0, + cube_1, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + cuda = DriftCorrection.from_4dstem( + torch.from_numpy(cube_0).cuda(), + torch.from_numpy(cube_1).cuda(), + scan_direction_degrees=(0.0, 90.0), + device="cuda", + ).preprocess( + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + drift_row = np.linspace(-0.75, 0.75, scan_size) + drift_column = np.linspace(0.5, -0.5, scan_size) + _add_known_raw_drift(cpu, 0, drift_row, drift_column) + _add_known_raw_drift(cuda, 0, drift_row, drift_column) + + cpu_result = cpu.corrected_4dstem( + merge=False, + output_dtype=np.float32, + chunk_size=4096, + verbose=False, + ) + cuda_result = cuda.corrected_4dstem( + merge=False, + output_dtype=torch.float32, + chunk_size=4096, + verbose=False, + ) + + np.testing.assert_allclose( + cpu_result.corrected_4dstem_0, + cuda_result.corrected_4dstem_0.cpu().numpy(), + rtol=3e-7, + atol=0.125, + ) From 86fd088fc45979cfd4c62133eca69da1df8a8dc2 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 10:56:51 -0700 Subject: [PATCH 07/11] Accelerate affine drift search without changing calibration (#20) --- src/quantem/imaging/drift/apply.py | 18 +++ src/quantem/imaging/drift/core/affine.py | 170 ++++++++++++++++++++- src/quantem/imaging/drift/preparation.py | 177 ++++++++++++++++++++++ tests/imaging/test_drift_downsampling.py | 183 +++++++++++++++++++++++ 4 files changed, 547 insertions(+), 1 deletion(-) create mode 100644 tests/imaging/test_drift_downsampling.py diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index 3f93f428..cfc5edfe 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -409,6 +409,15 @@ def generate_corrected( sampling=self.images[0].sampling, units=self.images[0].units, ) + image_corr.metadata.update( + { + "downsample": int(getattr(self, "downsample", 1)), + "downsample_method": getattr(self, "downsample_method", "none"), + "downsample_metadata": copy.deepcopy( + getattr(self, "downsample_metadata", {}) + ), + } + ) if show_merged: show_2d(image_corr.array, **kwargs) plt.show() @@ -578,6 +587,15 @@ def generate_corrected_image( sampling=self.images[0].sampling, units=self.images[0].units, ) + image_corr.metadata.update( + { + "downsample": int(getattr(self, "downsample", 1)), + "downsample_method": getattr(self, "downsample_method", "none"), + "downsample_metadata": copy.deepcopy( + getattr(self, "downsample_metadata", {}) + ), + } + ) if show_image: fig, ax = show_2d(image_corr.array, **kwargs) diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index 5724f1db..a4041f4d 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -3,7 +3,7 @@ import torch from torch.fft import fftfreq -from quantem.imaging.drift import diagnostics +from quantem.imaging.drift import diagnostics, preparation from quantem.imaging.drift.core.knots import ( bilinear_kde_batch, transform_coordinates_single_knot, @@ -22,6 +22,7 @@ def align_affine( upsample_factor: int = 8, max_image_shift: float | None = 32, chunk_size: int | None = None, + downsample: int | str = 1, show_merged: bool = True, show_images: bool = False, show_knots: bool = True, @@ -64,6 +65,10 @@ def align_affine( chunk_size : int or None Number of candidates per pass. If None, all candidates at once. Set to a smaller value if you run out of memory. + downsample : int or "auto", default 1 + Exact average-pooling factor used only to find the broad affine + candidate. The accepted rate is transferred back to the native grid + and refined there; corrected output is never downsampled. show_merged : bool Display the merged (averaged) image after alignment. show_images : bool @@ -99,6 +104,26 @@ def align_affine( f"num_tests must be odd (got {num_tests}). Try {num_tests + 1}." ) fixed_reference = getattr(self, "_reference_mode", False) + pyramid_factor = preparation.resolve_downsample( + downsample, + tuple(self.images[0].shape), + ) + if pyramid_factor > 1: + return _align_affine_pyramid( + self, + factor=pyramid_factor, + step=step, + num_tests=num_tests, + refine=refine, + upsample_factor=upsample_factor, + max_image_shift=max_image_shift, + chunk_size=chunk_size, + show_merged=show_merged, + show_images=show_images, + show_knots=show_knots, + verbose=verbose, + kwargs=kwargs, + ) # Build candidate grid with circular mask (~21% fewer than square) grid_axis = np.arange(-(num_tests - 1) / 2, (num_tests + 1) / 2) @@ -191,6 +216,149 @@ def _search_and_apply(candidates, label): return self +def _affine_rate_from_knots(correction, image_index: int) -> np.ndarray: + """Measure one canvas drift-rate vector from a solved one-knot trajectory.""" + delta = ( + correction.knots[image_index] + - correction._initial_knots[image_index] + ) + lines = delta.shape[1] + return ( + delta[:, -1, 0] - delta[:, 0, 0] + ) / max(lines - 1, 1) + + +def _apply_affine_rate(correction, rate: np.ndarray) -> None: + """Apply one canvas drift rate to every non-reference scan.""" + fixed_reference = getattr(correction, "_reference_mode", False) + for image_index in range(correction.shape[0]): + if fixed_reference and image_index == 0: + continue + lines = correction.knots[image_index].shape[1] + offset = np.arange(lines) - (lines - 1) / 2 + correction.knots[image_index][0] += rate[0] * offset[:, None] + correction.knots[image_index][1] += rate[1] * offset[:, None] + + +def _align_affine_pyramid( + self, + *, + factor: int, + step: float, + num_tests: int, + refine: bool, + upsample_factor: int, + max_image_shift: float | None, + chunk_size: int | None, + show_merged: bool, + show_images: bool, + show_knots: bool, + verbose: bool, + kwargs: dict, +): + """Search a pooled grid, transfer its rate, and verify it natively.""" + pooled = [ + preparation.average_downsample_2d(image.array, factor) + for image in self.images + ] + if getattr(self, "_reference_mode", False): + pyramid = type(self).from_reference( + pooled[0], + pooled[1], + scan_direction_degrees=float(self.scan_direction_degrees[1]), + device=self.device, + ) + else: + pyramid = type(self).from_data( + pooled, + scan_direction_degrees=self.scan_direction_degrees, + device=self.device, + ) + pyramid.preprocess( + pad_fraction=self.pad_fraction, + pad_value="median", + kde_sigma=self.kde_sigma, + number_knots=1, + show_merged=False, + show_images=False, + show_knots=False, + ).align_affine( + step=step, + num_tests=num_tests, + refine=refine, + upsample_factor=max(2, upsample_factor // 2), + max_image_shift=( + None + if max_image_shift is None + else max(2.0, float(max_image_shift) / factor) + ), + chunk_size=chunk_size, + downsample=1, + show_merged=False, + show_images=False, + show_knots=False, + verbose=False, + ) + solved_index = 1 if getattr(self, "_reference_mode", False) else 0 + coarse_rate = _affine_rate_from_knots(pyramid, solved_index) + _apply_affine_rate(self, coarse_rate) + + native_rate = coarse_rate.copy() + native_evaluations = 0 + if refine: + residual_step = step / max(num_tests - 1, 1) + axis = np.asarray((-residual_step, 0.0, residual_step)) + row, column = np.meshgrid(axis, axis, indexing="ij") + candidates = np.column_stack((row.ravel(), column.ravel())) + best_index, _ = self._affine_grid_search_batch( + candidates, + upsample_factor, + max_image_shift, + chunk_size, + fixed_reference=getattr(self, "_reference_mode", False), + ) + residual_rate = candidates[best_index] + _apply_affine_rate(self, residual_rate) + native_rate += residual_rate + native_evaluations = len(candidates) + + warped_t = self._warp_and_translate_torch( + max_image_shift, + upsample_factor, + fixed_reference=getattr(self, "_reference_mode", False), + ) + self.calculate_error(1, _warped_t=warped_t) + self.affine_search_info = { + "strategy": "average_pooled_pyramid", + "downsample_factor": factor, + "coarse_rate_px_per_line": coarse_rate.tolist(), + "native_rate_px_per_line": native_rate.tolist(), + "native_candidate_evaluations": native_evaluations, + } + diagnostics._record_stage(self, "affine") + if verbose: + print( + "align_affine: average-pooled pyramid " + f"{factor}x; native rate=({native_rate[0]:+.6f}, " + f"{native_rate[1]:+.6f}) px/line; " + f"{native_evaluations} native candidates" + ) + kwargs.pop("title", None) + if show_merged: + self.plot_merged_images( + show_knots=show_knots, + title="Merged: affine", + **kwargs, + ) + if show_images: + self.plot_transformed_images( + show_knots=show_knots, + title=[f"Image {index}: affine" for index in range(self.shape[0])], + **kwargs, + ) + return self + + @torch.inference_mode() def _affine_grid_search_batch( self, diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py index cbb66594..71c49dcf 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preparation.py @@ -1,7 +1,9 @@ import numpy as np import torch +from numpy.typing import NDArray +from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.dataset3d import Dataset3d from quantem.core.utils.compound_validators import ( validate_pad_value, @@ -17,12 +19,151 @@ ) +def validate_downsample(value: int) -> int: + """Validate one integer computational downsampling factor.""" + factor = int(value) + if factor < 1 or factor != value: + raise ValueError(f"downsample must be a positive integer, got {value!r}.") + return factor + + +def average_downsample_2d(array: NDArray, factor: int) -> NDArray: + """Average-pool a two-dimensional image by an exact integer factor.""" + image = np.asarray(array) + factor = validate_downsample(factor) + if factor == 1: + return np.ascontiguousarray(image) + if image.ndim != 2: + raise ValueError( + f"downsample currently supports 2-D images, got ndim={image.ndim}." + ) + rows, columns = image.shape + if rows % factor or columns % factor: + raise ValueError( + f"downsample={factor} requires image dimensions divisible by " + f"{factor}, got {image.shape}." + ) + downsampled = image.reshape( + rows // factor, + factor, + columns // factor, + factor, + ).mean(axis=(1, 3)) + dtype = image.dtype if np.issubdtype(image.dtype, np.floating) else np.float32 + return np.ascontiguousarray(downsampled.astype(dtype, copy=False)) + + +def automatic_downsample_factor(shape: tuple[int, int]) -> int: + """Return the largest exact average-pooling factor, capped at eight.""" + rows, columns = shape + for factor in (8, 4, 2): + if rows % factor == 0 and columns % factor == 0: + return factor + return 1 + + +def resolve_downsample(value: int | str, shape: tuple[int, int]) -> int: + """Resolve ``"auto"`` or validate one explicit exact divisor.""" + if value == "auto": + return automatic_downsample_factor(shape) + if isinstance(value, str): + raise ValueError( + "downsample must be 'auto' or a positive integer, " + f"got {value!r}." + ) + factor = validate_downsample(value) + if any(size % factor for size in shape): + raise ValueError( + f"downsample={factor} requires both image dimensions to be " + f"divisible by {factor}, got {shape}." + ) + return factor + + +def _scaled_coordinate_metadata( + values, + factor: int, + *, + offset: NDArray | None = None, +) -> NDArray: + """Scale the first two coordinate entries while preserving their length.""" + scaled = np.asarray(values, dtype=float).copy() + if scaled.size < 2: + scaled = np.resize(scaled, 2) + if offset is None: + scaled[:2] *= factor + else: + scaled[:2] += offset[:2] + return scaled + + +def apply_downsample(correction, factor: int) -> None: + """Average-downsample alignment images and preserve physical coordinates.""" + factor = resolve_downsample(factor, tuple(correction.images[0].shape)) + original_records = [] + downsampled_images = [] + for image_index, image in enumerate(correction.images): + original_array = np.asarray(image.array) + original_sampling = np.asarray(image.sampling, dtype=float).copy() + original_origin = np.asarray(image.origin, dtype=float).copy() + original_shape = tuple(int(value) for value in original_array.shape) + downsampled = average_downsample_2d(original_array, factor) + sampling = _scaled_coordinate_metadata(original_sampling, factor) + origin_offset = (factor - 1) * original_sampling[:2] / 2.0 + origin = _scaled_coordinate_metadata( + original_origin, + factor, + offset=origin_offset, + ) + metadata = dict(getattr(image, "metadata", {})) + metadata.update( + { + "downsample": factor, + "downsample_method": "average", + "downsample_original_shape": list(original_shape), + "downsample_original_sampling": original_sampling.tolist(), + "downsample_original_origin": original_origin.tolist(), + "downsampled_shape": list(downsampled.shape), + "sampling_is_downsampled": True, + } + ) + downsampled_image = Dataset2d.from_array( + downsampled, + name=f"{image.name} ({factor}x downsample)", + origin=origin, + sampling=sampling, + units=list(image.units), + signal_units=image.signal_units, + ) + downsampled_image.metadata.update(metadata) + downsampled_images.append(downsampled_image) + original_records.append( + { + "image_index": image_index, + "shape": list(original_shape), + "sampling": original_sampling.tolist(), + "origin": original_origin.tolist(), + "units": list(image.units), + } + ) + correction.images = downsampled_images + correction.downsample = factor + correction.downsample_method = "average" + correction.downsample_metadata = { + "factor": factor, + "method": "average", + "original_images": original_records, + "downsampled_shape": list(correction.images[0].shape), + } + + def preprocess( self, pad_fraction: float = 0.25, pad_value: float | str | list[float] = "median", kde_sigma: float = 0.5, number_knots: int = 1, + downsample: int | str = 1, show_merged: bool = False, show_images: bool = False, show_knots: bool = True, @@ -55,6 +196,10 @@ def preprocess( Number of Bezier knots per scanline. Use ``1`` (recommended) for linear drift correction. Higher values allow per-scanline curvature but are slower and rarely needed. + downsample : int or "auto", default 1 + Exact average-pooling factor for the 2-D alignment images. Sampling + and pixel-center origin are updated so physical coordinates remain + calibrated. This is computational downsampling, not detector binning. show_merged : bool Display the merged (averaged) warped images after preprocessing. show_images : bool @@ -75,6 +220,38 @@ def preprocess( ... images=[im0, im1], scan_direction_degrees=[0, 90]) >>> drift.preprocess(pad_fraction=0.25, kde_sigma=0.5, number_knots=1) """ + factor = resolve_downsample(downsample, tuple(self.images[0].shape)) + if factor > 1: + if getattr(self, "_datasets", None) is not None: + raise ValueError( + "preprocess(downsample>1) is only for 2-D alignment images. " + "Keep multidimensional data native and use the affine search " + "pyramid instead." + ) + if hasattr(self, "_initial_knots"): + raise RuntimeError( + "downsample changes the alignment grid. Create a new " + "DriftCorrection before choosing a different factor." + ) + apply_downsample(self, factor) + else: + self.downsample = 1 + self.downsample_method = "none" + self.downsample_metadata = { + "factor": 1, + "method": "none", + "original_images": [ + { + "image_index": index, + "shape": list(image.shape), + "sampling": np.asarray(image.sampling, dtype=float).tolist(), + "origin": np.asarray(image.origin, dtype=float).tolist(), + "units": list(image.units), + } + for index, image in enumerate(self.images) + ], + "downsampled_shape": list(self.images[0].shape), + } self.pad_fraction = float(pad_fraction) self.pad_value = validate_pad_value(pad_value, self.images) self.kde_sigma = float(kde_sigma) diff --git a/tests/imaging/test_drift_downsampling.py b/tests/imaging/test_drift_downsampling.py new file mode 100644 index 00000000..e343b37f --- /dev/null +++ b/tests/imaging/test_drift_downsampling.py @@ -0,0 +1,183 @@ +"""Calibration and provenance invariants for computational downsampling.""" + +import numpy as np +import pytest +from scipy.ndimage import gaussian_filter + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.imaging.drift import DriftCorrection +from quantem.imaging.drift.core.affine import _affine_rate_from_knots +from quantem.imaging.drift.preparation import ( + average_downsample_2d, + resolve_downsample, +) + + +def _calibrated_pair(size: int = 16): + rows, columns = np.indices((size, size), dtype=np.float32) + image = rows * 10 + columns + metadata = {"source": "synthetic calibrated scan"} + datasets = [] + for array in (image, np.rot90(image, k=-1).copy()): + dataset = Dataset2d.from_array( + array, + origin=(1.0, 2.0), + sampling=(0.2, 0.3), + units=("nm", "nm"), + ) + dataset.metadata.update(metadata) + datasets.append(dataset) + return datasets + + +def test_average_downsample_is_exact_block_mean(): + """Integer count images become float32 block averages, never decimation.""" + image = np.arange(64, dtype=np.uint16).reshape(8, 8) + + result = average_downsample_2d(image, 2) + + expected = image.reshape(4, 2, 4, 2).mean(axis=(1, 3)).astype(np.float32) + assert result.dtype == np.float32 + np.testing.assert_array_equal(result, expected) + + +def test_preprocess_downsample_preserves_pixel_center_calibration(): + """Sampling grows and origin moves to the center of each averaged block.""" + datasets = _calibrated_pair() + drift = DriftCorrection.from_data( + datasets, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + downsample=2, + show_merged=False, + show_images=False, + show_knots=False, + ) + + assert drift.images[0].shape == (8, 8) + np.testing.assert_allclose(drift.images[0].sampling, (0.4, 0.6)) + np.testing.assert_allclose(drift.images[0].origin, (1.1, 2.15)) + assert drift.images[0].units == ["nm", "nm"] + assert drift.images[0].metadata["source"] == "synthetic calibrated scan" + assert drift.downsample_metadata["original_images"][0]["shape"] == [16, 16] + + +def test_corrected_output_records_downsampling_provenance(): + """A corrected image tells readers which computational grid was fitted.""" + drift = DriftCorrection.from_data( + _calibrated_pair(), + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + downsample=2, + show_merged=False, + show_images=False, + show_knots=False, + ) + + corrected = drift.generate_corrected( + upsample_factor=1, + show_merged=False, + ) + + assert corrected.metadata["downsample"] == 2 + assert corrected.metadata["downsample_method"] == "average" + assert corrected.metadata["downsample_metadata"]["factor"] == 2 + np.testing.assert_allclose(corrected.sampling, (0.4, 0.6)) + + +def test_downsample_requires_exact_divisor_and_new_correction(): + """Grid changes fail with a corrective message instead of silent cropping.""" + with pytest.raises(ValueError, match="divisible"): + resolve_downsample(4, (18, 18)) + + drift = DriftCorrection.from_data( + _calibrated_pair(), + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + show_merged=False, + show_images=False, + show_knots=False, + ) + with pytest.raises(RuntimeError, match="Create a new DriftCorrection"): + drift.preprocess( + downsample=2, + show_merged=False, + show_images=False, + show_knots=False, + ) + + +def test_automatic_factor_is_largest_exact_divisor_up_to_eight(): + """Automatic resolution selection is deterministic and shape-safe.""" + assert resolve_downsample("auto", (2048, 2048)) == 8 + assert resolve_downsample("auto", (1026, 1026)) == 2 + assert resolve_downsample("auto", (1025, 1025)) == 1 + + +def test_affine_pyramid_refines_rate_on_native_grid(): + """A pooled broad search retains the native image and affine rate.""" + size = 64 + rng = np.random.default_rng(23) + reference = gaussian_filter(rng.normal(size=(size, size)), 1.5).astype(np.float32) + target = np.empty_like(reference) + columns = np.arange(size, dtype=np.float32) + for row in range(size): + shift = 0.04 * (row - (size - 1) / 2) + target[row] = np.interp( + columns + shift, + columns, + reference[row], + left=float(np.median(reference[row])), + right=float(np.median(reference[row])), + ) + + native = DriftCorrection.from_reference( + reference, + target, + scan_direction_degrees=0.0, + device="cpu", + ).preprocess( + show_merged=False, + show_images=False, + show_knots=False, + ) + pyramid = DriftCorrection.from_reference( + reference, + target, + scan_direction_degrees=0.0, + device="cpu", + ).preprocess( + show_merged=False, + show_images=False, + show_knots=False, + ) + + native.align_affine( + step=0.02, + num_tests=7, + max_image_shift=8, + downsample=1, + show_merged=False, + show_images=False, + show_knots=False, + ) + pyramid.align_affine( + step=0.02, + num_tests=7, + max_image_shift=8, + downsample=2, + show_merged=False, + show_images=False, + show_knots=False, + ) + + assert pyramid.images[0].shape == (size, size) + assert pyramid.affine_search_info["downsample_factor"] == 2 + np.testing.assert_allclose( + _affine_rate_from_knots(pyramid, 1), + _affine_rate_from_knots(native, 1), + atol=0.01, + ) From 9b2ac0804a300a1bf1424c601fac00b263f262f0 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 11:17:59 -0700 Subject: [PATCH 08/11] Reproduce multidimensional manuscript workflows with automatic correction (#21) --- pyproject.toml | 1 + src/quantem/imaging/__init__.py | 1 + src/quantem/imaging/drift/__init__.py | 14 +- src/quantem/imaging/drift/apply.py | 1468 ++++++++------- src/quantem/imaging/drift/core/__init__.py | 7 +- src/quantem/imaging/drift/core/affine.py | 1965 ++++++++++++++++---- src/quantem/imaging/drift/core/knots.py | 534 ++++-- src/quantem/imaging/drift/core/nonrigid.py | 1115 +++++++---- src/quantem/imaging/drift/core/strip.py | 840 ++++++++- src/quantem/imaging/drift/core/warping.py | 733 +++++++- src/quantem/imaging/drift/correction.py | 677 ++++--- src/quantem/imaging/drift/diagnostics.py | 712 ++++--- src/quantem/imaging/drift/fourdstem.py | 298 +-- src/quantem/imaging/drift/io.py | 519 ++++-- src/quantem/imaging/drift/plot.py | 1208 +++++++++--- src/quantem/imaging/drift/preparation.py | 821 +++++--- src/quantem/imaging/drift/report.py | 211 ++- tests/imaging/test_drift.py | 514 ++--- tests/imaging/test_drift3d.py | 33 +- tests/imaging/test_drift4d.py | 76 +- tests/imaging/test_drift_downsampling.py | 98 +- tests/imaging/test_drift_residual.py | 174 ++ tests/imaging/test_drift_utils.py | 29 +- uv.lock | 117 +- 24 files changed, 8678 insertions(+), 3487 deletions(-) create mode 100644 tests/imaging/test_drift_residual.py diff --git a/pyproject.toml b/pyproject.toml index d83c3c51..ed947e33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "colorspacious", "dill", "numpy>2", + "pandas>=2.2", "matplotlib", "scipy", "tqdm", diff --git a/src/quantem/imaging/__init__.py b/src/quantem/imaging/__init__.py index 929f549b..835c3011 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,6 +1,7 @@ from quantem.imaging.drift import ( CorrectionResult as CorrectionResult, DriftCorrection as DriftCorrection, + StripPass as StripPass, ) from quantem.imaging.drift import ( pair_spectrum_image_references as pair_spectrum_image_references, diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index 5e8f720d..2049d242 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -1,15 +1,6 @@ from quantem.imaging.drift.correction import DriftCorrection as DriftCorrection -from quantem.imaging.drift.apply import ( - _bounded_sine_sigmoid_torch as _bounded_sine_sigmoid_torch, - _fourier_crop_torch as _fourier_crop_torch, - bounded_sine_sigmoid as bounded_sine_sigmoid, -) -from quantem.imaging.drift.core.knots import ( - DriftInterpolator as DriftInterpolator, -) -from quantem.imaging.drift.fourdstem import ( - CorrectionResult as CorrectionResult, -) +from quantem.imaging.drift.core.strip import StripPass as StripPass +from quantem.imaging.drift.fourdstem import CorrectionResult as CorrectionResult from quantem.imaging.drift.io import ( pair_spectrum_image_references as pair_spectrum_image_references, read_emd as read_emd, @@ -18,4 +9,3 @@ ) DriftCorrection.__module__ = __name__ -DriftInterpolator.__module__ = __name__ diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index cfc5edfe..efc73d25 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -1,734 +1,926 @@ +"""Apply a solved drift field and return corrected scientific data.""" -import copy -import warnings +from copy import deepcopy -import matplotlib.pyplot as plt import numpy as np import torch -from scipy.ndimage import distance_transform_edt +import torch.nn.functional as F +from numpy.typing import NDArray +from scipy.ndimage import binary_closing as ndi_binary_closing +from tqdm.auto import tqdm from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.dataset3d import Dataset3d -from quantem.core.utils.imaging_utils import ( - fourier_cropping, -) -from quantem.core.visualization import show_2d -from quantem.imaging.drift.core.knots import ( - bilinear_kde_batch, - transform_coordinates_single_knot, +from quantem.imaging.drift.core import knots as drift_knots +from quantem.imaging.drift.core.warping import ( + backward_warp, + ensure_warped_images, + reference_scan_stack, + warp_and_translate, ) -def _raw_frame_displacement(self, image_index: int) -> torch.Tensor: - """Return a one-knot trajectory displacement in raw scan coordinates.""" - index = image_index % len(self.images) - if self.knots[index].shape[-1] != 1: - raise ValueError( - "Multidimensional correction currently requires number_knots=1." - ) - scan_rows, scan_cols = self.images[index].shape - delta_canvas_t = torch.as_tensor( - self.knots[index] - self._initial_knots[index], - dtype=torch.float32, - device=self._device, - )[:, :, 0] - scan_fast_t = torch.as_tensor( - self.scan_fast[index], dtype=torch.float32, device=self._device - ) - scan_slow_t = torch.as_tensor( - self.scan_slow[index], dtype=torch.float32, device=self._device - ) - aspect = float(scan_rows - 1) / float(scan_cols - 1) if scan_cols > 1 else 1.0 - determinant = ( - scan_slow_t[0] * scan_fast_t[1] - - scan_fast_t[0] * aspect * scan_slow_t[1] - ) - drift_row_t = ( - scan_fast_t[1] * delta_canvas_t[0] - - scan_fast_t[0] * aspect * delta_canvas_t[1] - ) / determinant - drift_col_t = ( - -scan_slow_t[1] * delta_canvas_t[0] - + scan_slow_t[0] * delta_canvas_t[1] - ) / determinant - return torch.stack((drift_row_t, drift_col_t)) - - -def _scan_warp_grid(self, image_index: int) -> torch.Tensor: - """Return one raw-frame sampling grid for a solved scan trajectory.""" - index = image_index % len(self.images) - scan_rows, scan_cols = self.images[index].shape - drift_row_t, drift_col_t = _raw_frame_displacement(self, index) - row_t = torch.arange(scan_rows, dtype=torch.float32, device=self._device) - col_t = torch.arange(scan_cols, dtype=torch.float32, device=self._device) - sample_row_t = row_t[:, None] - drift_row_t[:, None] - sample_col_t = col_t[None, :] - drift_col_t[:, None] - return torch.stack( - ( - 2.0 * sample_col_t.expand(scan_rows, scan_cols) / (scan_cols - 1) - 1.0, - 2.0 * sample_row_t.expand(scan_rows, scan_cols) / (scan_rows - 1) - 1.0, - ), - dim=-1, - )[None] - - -@torch.inference_mode() -def _apply_scan_field( - self, - data: np.ndarray | torch.Tensor, - *, - image_index: int, - chunk_size: int | None, - output_dtype: np.dtype | torch.dtype | str | None, - mode: str = "bilinear", - output: np.ndarray | None = None, -) -> np.ndarray | torch.Tensor: - """Warp only the two leading scan axes of an array.""" - index = image_index % len(self.images) - scan_rows, scan_cols = self.images[index].shape - if data.ndim < 2 or tuple(data.shape[:2]) != (scan_rows, scan_cols): - raise ValueError( - "The leading scan axes must match the fitted alignment image; " - f"got {tuple(data.shape[:2])} and {(scan_rows, scan_cols)}." - ) - trailing_shape = tuple(data.shape[2:]) - num_channels = int(np.prod(trailing_shape, dtype=np.int64)) if trailing_shape else 1 - source_is_torch = isinstance(data, torch.Tensor) - source_dtype = data.dtype - if output is not None: - if not isinstance(output, np.ndarray) or output.shape != tuple(data.shape): - raise ValueError("output must be a NumPy array matching the input shape.") - output_flat = output.reshape(scan_rows, scan_cols, num_channels) - else: - output_flat = None - flat = ( - data.reshape(scan_rows, scan_cols, num_channels) - if source_is_torch - else torch.from_numpy( - np.ascontiguousarray(data).reshape(scan_rows, scan_cols, num_channels) - ) - ) - grid_t = _scan_warp_grid(self, index) - chunk_size = min(num_channels, 64 if chunk_size is None else int(chunk_size)) - if chunk_size < 1: - raise ValueError(f"chunk_size must be positive, got {chunk_size}.") - corrected_t = None - if output_flat is None: - corrected_t = torch.empty( - (scan_rows, scan_cols, num_channels), - dtype=torch.float32, - device=self._device, - ) - for start in range(0, num_channels, chunk_size): - stop = min(start + chunk_size, num_channels) - channels_t = flat[:, :, start:stop].to( - device=self._device, - dtype=torch.float32, - ).permute(2, 0, 1)[None] - corrected_block_t = torch.nn.functional.grid_sample( - channels_t, - grid_t, - mode=mode, - padding_mode="border", - align_corners=True, - )[0].permute(1, 2, 0) - if output_flat is None: - corrected_t[:, :, start:stop] = corrected_block_t - else: - output_dtype_np = output_flat.dtype - if np.issubdtype(output_dtype_np, np.integer): - limits = np.iinfo(output_dtype_np) - corrected_block_t = corrected_block_t.round().clamp( - limits.min, - limits.max, - ) - output_flat[:, :, start:stop] = corrected_block_t.cpu().numpy().astype( - output_dtype_np, - copy=False, - ) - if output_flat is not None: - return output - assert corrected_t is not None - corrected_t = corrected_t.reshape((scan_rows, scan_cols, *trailing_shape)) - preserve_dtype = output_dtype == "same" - requested_dtype = source_dtype if preserve_dtype else output_dtype - if requested_dtype is not None: - if source_is_torch or isinstance(requested_dtype, torch.dtype): - torch_dtype = source_dtype if preserve_dtype else requested_dtype - if not isinstance(torch_dtype, torch.dtype): - torch_dtype = torch.from_numpy( - np.empty((), dtype=np.dtype(torch_dtype)) - ).dtype - if not torch_dtype.is_floating_point: - limits = torch.iinfo(torch_dtype) - corrected_t = corrected_t.round().clamp(limits.min, limits.max) - corrected_t = corrected_t.to(torch_dtype) - else: - numpy_dtype = np.dtype(source_dtype if preserve_dtype else requested_dtype) - if np.issubdtype(numpy_dtype, np.integer): - limits = np.iinfo(numpy_dtype) - corrected_t = corrected_t.round().clamp(limits.min, limits.max) - return corrected_t.cpu().numpy().astype(numpy_dtype, copy=False) - if source_is_torch: - return corrected_t.to(data.device) - return corrected_t.cpu().numpy() +def dataset_info(dataset) -> dict[str, object]: + """Copy the calibration and metadata needed for corrected output.""" + if not hasattr(dataset, "array"): + return {} + info = {} + for name in ("name", "origin", "sampling", "units", "signal_units"): + if hasattr(dataset, name): + info[name] = deepcopy(getattr(dataset, name)) + metadata = getattr(dataset, "metadata", None) + if isinstance(metadata, dict): + info["metadata"] = deepcopy(metadata) + return info + + +def _corrected_dataset(array: np.ndarray, info: dict[str, object]): + """Construct corrected data without discarding source calibration.""" + dataset_class = Dataset2d if array.ndim == 2 else Dataset3d + if array.ndim == 4: + from quantem.core.datastructures.dataset4d import Dataset4d + + dataset_class = Dataset4d + kwargs = { + name: deepcopy(info[name]) + for name in ("name", "origin", "sampling", "units", "signal_units") + if name in info + } + result = dataset_class.from_array(array, **kwargs) + if metadata := info.get("metadata"): + result.metadata.update(deepcopy(metadata)) + return result -@torch.inference_mode() -def apply_correction( - self, - data: Dataset3d | np.ndarray | torch.Tensor | None = None, - image_index: int = -1, +def padding_offset( + canvas_shape: tuple[int, int], + scan_shape: tuple[int, int], *, - mode: str = "bilinear", - chunk_size: int | None = None, - output_dtype: np.dtype | torch.dtype | str | None = None, - output_device: str | torch.device | None = None, - output: np.ndarray | None = None, - verbose: bool = True, -): - """Apply one learned spatial field without mixing trailing data axes. + integer: bool = False, +) -> tuple[float, float] | tuple[int, int]: + """Return the ``(row, col)`` offset of a scan in its padded canvas.""" + canvas_h, canvas_w = canvas_shape + scan_h, scan_w = scan_shape + if integer: + return (canvas_h - scan_h) // 2, (canvas_w - scan_w) // 2 + return (canvas_h - scan_h) / 2.0, (canvas_w - scan_w) / 2.0 - Parameters - ---------- - data : Dataset3d, numpy.ndarray, torch.Tensor, or None - Data with leading ``(scan_row, scan_col)`` axes. ``None`` uses a raw - dataset retained by :meth:`DriftCorrection.from_4dstem`. - image_index : int, default -1 - Fitted scan trajectory to apply. - chunk_size : int or None, default None - Number of spectral or detector channels corrected per GPU batch. - output_dtype : numpy dtype, torch dtype, "same", or None - Output dtype. The default returns floating-point interpolation. - output_device : str, torch.device, or None - Device for a tensor result. NumPy inputs return NumPy by default. - output : numpy.ndarray or None - Optional preallocated NumPy destination. - verbose : bool, default True - Reserved for progress reporting in larger workflows. - Returns - ------- - Dataset3d, numpy.ndarray, or torch.Tensor - Corrected data with unchanged shape and trailing-axis order. +def fourier_crop_torch( + fft_array: torch.Tensor, + crop_shape: tuple[int, int], +) -> torch.Tensor: + """Crop a corner-centered FFT tensor to its lowest frequencies.""" + crop_h, crop_w = crop_shape + h1 = crop_h // 2 + h2 = crop_h - h1 + w1 = crop_w // 2 + w2 = crop_w - w1 + result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) + result[:h1, :w1] = fft_array[:h1, :w1] + result[:h1, -w2:] = fft_array[:h1, -w2:] + result[-h2:, :w1] = fft_array[-h2:, :w1] + result[-h2:, -w2:] = fft_array[-h2:, -w2:] + return result - Examples - -------- - >>> corrected = drift.apply_correction(spectrum_image, image_index=1) - """ - if not hasattr(self, "_initial_knots"): - raise RuntimeError( - "No drift field found. Call preprocess() and align_affine() first." - ) - if "affine" not in getattr(self, "_diagnostic_knots", {}): - raise RuntimeError( - "No fitted reference field found. Call align_affine() before " - "applying the correction." + +def largest_rectangle(mask: np.ndarray) -> tuple[int, int, int, int]: + """Return the largest axis-aligned all-true rectangle in ``mask``.""" + bin_factor = max(1, int(np.ceil(max(mask.shape) / 512))) + if bin_factor > 1: + height = mask.shape[0] // bin_factor + width = mask.shape[1] // bin_factor + small = ( + mask[: height * bin_factor, : width * bin_factor] + .reshape(height, bin_factor, width, bin_factor) + .all(axis=(1, 3)) ) - if mode not in {"bilinear", "bicubic"}: - raise ValueError(f"mode must be 'bilinear' or 'bicubic', got {mode!r}.") - if data is None: - datasets = getattr(self, "_datasets", None) - if datasets is None: - raise ValueError("No data supplied and no 4D-STEM dataset is retained.") - data = datasets[image_index % len(datasets)] - dataset = data if isinstance(data, Dataset3d) else None - array = dataset.array if dataset is not None else data - corrected_array = _apply_scan_field( - self, - array, - image_index=image_index, - chunk_size=chunk_size, - output_dtype=output_dtype, - mode=mode, - output=output, - ) - if output is None and output_device is not None: - corrected_array = torch.as_tensor(corrected_array).to(output_device) - if dataset is None: - return corrected_array - corrected = Dataset3d.from_array( - np.asarray(corrected_array), - name=f"drift-corrected {dataset.name}", - origin=dataset.origin.copy(), - sampling=dataset.sampling.copy(), - units=list(dataset.units), - signal_units=dataset.signal_units, + else: + small = mask + heights = np.zeros(small.shape[1], dtype=int) + best = (0, 0, 0, 0, 0) + for row in range(small.shape[0]): + heights = np.where(small[row], heights + 1, 0) + stack: list[tuple[int, int]] = [] + extended = np.append(heights, 0) + for col in range(len(extended)): + start = col + while stack and stack[-1][1] >= extended[col]: + stack_start, height = stack.pop() + if height * (col - stack_start) > best[0]: + best = ( + height * (col - stack_start), + row - height + 1, + row + 1, + stack_start, + col, + ) + start = stack_start + stack.append((start, extended[col])) + _, row_0, row_1, col_0, col_1 = best + row_0, row_1, col_0, col_1 = ( + row_0 * bin_factor, + row_1 * bin_factor, + col_0 * bin_factor, + col_1 * bin_factor, ) - corrected.metadata.update(copy.deepcopy(dataset.metadata)) - return corrected - - -@torch.inference_mode() -def generate_corrected( + row_1 = min(row_1, mask.shape[0]) + col_1 = min(col_1, mask.shape[1]) + while row_0 < row_1 and not mask[row_0, col_0:col_1].all(): + row_0 += 1 + while row_1 > row_0 and not mask[row_1 - 1, col_0:col_1].all(): + row_1 -= 1 + while col_0 < col_1 and not mask[row_0:row_1, col_0].all(): + col_0 += 1 + while col_1 > col_0 and not mask[row_0:row_1, col_1 - 1].all(): + col_1 -= 1 + return row_0, row_1, col_0, col_1 + + +def corrected( self, + *, upsample_factor: int = 2, - output_original_shape: bool = True, + output_original_shape: bool | None = None, strip_padding: bool = False, - mask_output: bool = True, - mask_edge_blend: float = 8.0, - fourier_filter: bool = True, - filter_midpoint: float = 0.5, - kde_sigma: float | None = 0.5, - weight_thresh: float = 0.1, - show_merged: bool = True, - **kwargs, + smoothing_sigma: float | None = 0.5, + stage: str | None = None, + merge: bool = True, + verbose: bool = True, ): - """Generate the final drift-corrected image on GPU using torch. + """Return corrected scientific data in the natural type for the acquisition. + + A single endpoint keeps HAADF, spectrum-image, and 4D-STEM workflows from + requiring separate output APIs while preserving each dataset's calibration + and axis order. Parameters ---------- upsample_factor : int, default 2 - Factor to upsample the output image for enhanced interpolation accuracy. - output_original_shape : bool, default True - If True, crop the output image back to the original input dimensions. + Sampling multiplier for the corrected image. + output_original_shape : bool or None, default None + ``None`` keeps the natural frame: the padded solver canvas for paired + 2-D scans and the native input frame for reference datasets. ``True`` + returns the original image shape; ``False`` keeps the solver canvas. strip_padding : bool, default False - If True (and output_original_shape is True), also strip the scan padding - to return only the original scan-area pixels. - mask_output : bool, default True - If True, mask the output using the probe position weights. - mask_edge_blend : float, default 8.0 - Pixels over which the mask edge is blended. - fourier_filter : bool, default True - Whether to apply Fourier-based directional filtering to merge corrected images. - filter_midpoint : float, default 0.5 - Midpoint for the sigmoid-based Fourier weighting filter. - kde_sigma : float or None, default 0.5 - Standard deviation for kernel density estimation. Uses object's kde_sigma if None. - weight_thresh : float, default 0.1 - Threshold for masking outputs. - show_merged : bool, default True - Whether to display the final corrected image. - **kwargs - Additional keyword arguments passed to the plotting function. + Remove pixels that do not share measured coverage across scans. + smoothing_sigma : float or None, default 0.5 + Gaussian smoothing applied during scanline interpolation. + stage : {"initial", "affine", "strip", None}, optional + Saved correction stage to apply. ``None`` uses the current solution. + merge : bool, default True + Average corrected image pairs. ``False`` returns each scan separately. + verbose : bool, default True + Show progress for a multi-channel reference dataset. Returns ------- - image_corr : Dataset2d - The final drift-corrected output image. + Dataset2d, Dataset3d, or list of Dataset2d + Corrected data with the source calibration and metadata. + + Examples + -------- + >>> corrected = drift.corrected() + >>> scans = drift.corrected(merge=False) """ - if not hasattr(self, "knots"): + automatic_output_frame = output_original_shape is None + output_original_shape = ( + self._reference_mode + if automatic_output_frame + else bool(output_original_shape) + ) + + if self._reference_mode: + if ( + upsample_factor != 2 + or not output_original_shape + or strip_padding + or smoothing_sigma != 0.5 + or not merge + ): + raise ValueError( + "Reference-mode corrected() returns the corrected source dataset " + "at its native shape. Use apply_correction() for interpolation " + "controls and crop() to remove unsupported borders." + ) + drifted = self._datasets[1] + corrected = apply_correction( + self, + drifted, + image_index=1, + stage=stage, + verbose=verbose, + ) + if isinstance(corrected, torch.Tensor): + corrected = corrected.cpu().numpy() + corrected = corrected.astype(np.float32, copy=False) + return _corrected_dataset( + corrected, + getattr(self, "_reference_dataset_info", {}), + ) + if getattr(self, "_datasets", None) is not None: raise RuntimeError( - "No knots found. Call .preprocess() before generating the corrected image." + "4D-STEM collection correction uses the explicit " + "corrected_4dstem() API. Use " + "DriftCorrection.from_4dstem(data_0, data_1, ...)" + ".preprocess().correct_affine().corrected_4dstem()." ) + if not merge: + if ( + upsample_factor != 2 + or (not automatic_output_frame and not output_original_shape) + or strip_padding + or smoothing_sigma != 0.5 + ): + raise ValueError( + "corrected(merge=False) returns the individual scans on the " + "solver canvas. Resampling, output-shape, smoothing, and crop " + "controls apply only to the merged image." + ) + panels = comparison_panels(self, stage) + return [ + Dataset2d.from_array( + np.asarray(array), + name=f"drift corrected scan {index}", + origin=self.imgs[0].origin, + sampling=self.imgs[0].sampling, + units=self.imgs[0].units, + ) + for index, array in enumerate(panels["corrected_scans"]) + ] + device = self._device dtype = self._dtype up_h = round(self.shape[1] * upsample_factor) up_w = round(self.shape[2] * upsample_factor) - canvas_shape = (up_h, up_w) - - if kde_sigma is None: - kde_sigma = self.kde_sigma + canvas_up = (up_h, up_w) + if smoothing_sigma is None: + smoothing_sigma = self.kde_sigma stack_corr = torch.zeros(self.shape[0], up_h, up_w, dtype=dtype, device=device) - weight_corr = torch.zeros_like(stack_corr) - for img_idx in range(self.shape[0]): - knots_t = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_t, - self.scan_fast_t[img_idx], - self.images[img_idx].shape, - ) - warped, weights = bilinear_kde_batch( - row_t[None] * upsample_factor, - col_t[None] * upsample_factor, - self.images_t[img_idx], - canvas_shape, - kde_sigma * upsample_factor, - self.pad_value[img_idx], - ) - stack_corr[img_idx] = warped[0] - weight_corr[img_idx] = weights[0] - - if fourier_filter: - freq_row = torch.fft.fftfreq(up_h, dtype=dtype, device=device)[:, None] - freq_col = torch.fft.fftfreq(up_w, dtype=dtype, device=device)[None, :] - freq_angle = torch.atan2(freq_col, freq_row) - stack_fft = torch.fft.fft2(stack_corr) - weights = torch.zeros_like(stack_corr) - for img_idx in range(self.shape[0]): - weights[img_idx] = torch.abs( - torch.remainder( - (freq_angle - self.scan_direction[img_idx]) / np.pi + 0.5, - 1.0, - ) - 0.5 - ) / 0.5 - weights[img_idx, 0, 0] = 1.0 - weights[img_idx] = _bounded_sine_sigmoid_torch( - weights[img_idx], - midpoint=filter_midpoint, - ) - stack_fft[img_idx] *= weights[img_idx] - weights_sum = weights.sum(0) - fft_sum = stack_fft.sum(0) - image_corr_fft = torch.where( - weights_sum > 0.0, - fft_sum / weights_sum.clamp(min=1e-8), - torch.zeros_like(fft_sum), - ) - else: - image_corr_fft = torch.fft.fft2(stack_corr.mean(0)) - - if mask_output: - weight_np = weight_corr.cpu().numpy() - mask_edge = np.prod(weight_np >= (weight_thresh / upsample_factor**2), axis=0) - mask_edge[:, 0] = False - mask_edge[:, -1] = False - mask_edge[0, :] = False - mask_edge[-1, :] = False - mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend - mask_np = ( - np.cos( - (np.pi / 2) - * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) - ) - ** 2 - ) - mask_t = torch.as_tensor(mask_np, dtype=dtype, device=device) - pad_value_mean = float(np.mean(self.pad_value)) - image_corr_fft = torch.fft.fft2( - torch.fft.ifft2(image_corr_fft).real * mask_t + pad_value_mean * (1 - mask_t) + knots = drift_knots.stage_knots(self, stage) + for image_index in range(self.shape[0]): + warped, _ = drift_knots.interpolator( + self, image_index, knots[image_index] + ).warp_to_canvas( + self.imgs_t[image_index], + canvas_up, + smoothing_sigma * upsample_factor, + self.pad_value[image_index], + upsample_factor=upsample_factor, ) + stack_corr[image_index] = warped + image_corr_fft = torch.fft.fft2(stack_corr.mean(0)) - if output_original_shape: - image_corr_fft = _fourier_crop_torch(image_corr_fft, self.shape[-2:]) / upsample_factor**2 - - corr_np = torch.fft.ifft2(image_corr_fft).real.cpu().numpy() + output_shape = ( + tuple(int(value) for value in self.imgs[0].shape[:2]) + if output_original_shape + else tuple(int(value) for value in self.shape[-2:]) + ) + image_corr_fft = fourier_crop_torch( + image_corr_fft, + output_shape, + ) / upsample_factor**2 + corrected_array = torch.fft.ifft2(image_corr_fft).real.cpu().numpy() if strip_padding and output_original_shape: - scan_h, scan_w = self.images[0].shape[:2] - canvas_h, canvas_w = corr_np.shape[:2] - pad_h = (canvas_h - scan_h) // 2 - pad_w = (canvas_w - scan_w) // 2 - corr_np = corr_np[pad_h:pad_h + scan_h, pad_w:pad_w + scan_w] + scan_h, scan_w = self.imgs[0].shape[:2] + pad_h, pad_w = padding_offset(corrected_array.shape[:2], (scan_h, scan_w), integer=True) + corrected_array = corrected_array[pad_h : pad_h + scan_h, pad_w : pad_w + scan_w] image_corr = Dataset2d.from_array( - corr_np, + corrected_array, name="drift corrected image", - origin=self.images[0].origin, - sampling=self.images[0].sampling, - units=self.images[0].units, + origin=self.imgs[0].origin, + sampling=self.imgs[0].sampling, + units=self.imgs[0].units, ) image_corr.metadata.update( { "downsample": int(getattr(self, "downsample", 1)), "downsample_method": getattr(self, "downsample_method", "none"), - "downsample_metadata": copy.deepcopy( - getattr(self, "downsample_metadata", {}) + "downsample_metadata": getattr( + self, + "downsample_metadata", + {"factor": 1, "method": "none"}, ), + "downsample_sampling": np.asarray(self.imgs[0].sampling, dtype=float).tolist(), + "downsample_units": list(self.imgs[0].units), } ) - if show_merged: - show_2d(image_corr.array, **kwargs) - plt.show() return image_corr -def generate_corrected_image( - self, - upsample_factor: int = 2, - output_original_shape: bool = True, - mask_output: bool = True, - mask_edge_blend: float = 8.0, - fourier_filter: bool = True, - filter_midpoint: float = 0.5, - kde_sigma: float = 0.5, - weight_thresh=0.1, - show_image: bool = True, - **kwargs, -): +def apply_correction_to_dataset( + correction, + ds_4d: torch.Tensor | np.ndarray | None = None, + *, + image_index: int = -1, + mode: str = "bilinear", + chunk_size: int | None = None, + output_dtype: torch.dtype | np.dtype | str | None = None, + output_device: str | torch.device | None = None, + output: np.ndarray | None = None, + verbose: bool = False, + progress_desc: str = "Applying drift correction", + stage: str | None = None, +) -> torch.Tensor | np.ndarray: + """Apply drift correction to a ≥3-D dataset with scan axes leading. + + Internal worker for the 4D-STEM / spectral path of + :meth:`DriftCorrection.apply_correction`. It selects single-shot or chunked + processing from available device memory and supports preallocated + ``output=`` for zero-copy memmap workflows. """ - Generate the final drift-corrected image after aligning a stack of input images. + # Resolve dataset from stored data when not provided explicitly + if ds_4d is None: + datasets = correction._datasets + if datasets is None: + raise ValueError( + "No dataset provided and none stored. Pass ds_4d " + "explicitly, or build this instance with " + "DriftCorrection(ds_a, ds_b, ...)." + ) + if image_index < 0: + image_index = len(datasets) + image_index + if image_index < 0 or image_index >= len(datasets): + raise IndexError( + f"image_index={image_index} out of range for " + f"{len(datasets)} stored datasets" + ) + ds_4d = datasets[image_index] + + is_numpy = isinstance(ds_4d, np.ndarray) + original_shape = ds_4d.shape if is_numpy else tuple(ds_4d.shape) + input_np_dtype = ds_4d.dtype if is_numpy else None + use_external_output = output is not None + + if use_external_output: + if not isinstance(output, np.ndarray): + raise TypeError( + "output must be a numpy ndarray (or np.memmap), " + f"got {type(output).__name__}" + ) + if tuple(output.shape) != tuple(original_shape): + raise ValueError( + f"output shape {output.shape} does not match " + f"ds_4d shape {original_shape}" + ) - Parameters - ---------- - upsample_factor : int, default 2 - Factor to upsample the output image for enhanced interpolation accuracy. - output_original_shape : bool, default True - If True, crop the output image back to the original input dimensions after processing. - mask_output : bool, default True - If true, mask the output using the probe position weights - mask_edge_blend : float, default 8.0 - Value in pixels to blend from the edge of the mask (where we have data) - fourier_filter : bool, default True - Whether to apply Fourier-based directional filtering to merge corrected images. - filter_midpoint : float, default 0.5 - Midpoint for the sigmoid-based Fourier weighting filter, determining transition smoothness. - Setting this to a low value close to 0 will include more signal but also more slow scan artifacts. - If using 2 images at 0 and 90 degrees scan angles, any value >0.75 will be unstable. - Only use larger values (close to 1.0) if multiple images covering many scan angles are used. - kde_sigma : float, default 0.5 - Standard deviation for kernel density estimation used during image interpolation. Defaults - to the object's stored kde_sigma if set to None. - weight_thresh: float, default 0.1 - This value sets the threshold for masking the outputs. - For very large jitter artifacts this value can be lowered. - show_image : bool, default True - Whether to display the final corrected image after processing. - **kwargs : dict - Additional keyword arguments passed to the plotting function when displaying the image. + ndim = len(original_shape) + if ndim < 3: + raise ValueError( + f"ds_4d must be at least 3D, got shape {original_shape}" + ) - Returns - ------- - image_corr : Dataset2d - The final drift-corrected output image encapsulated in a Dataset2d object. - - Notes - ----- - - The function applies per-frame warping using knot-based interpolation and optionally - performs directional Fourier filtering to blend multiple warped images. - - The Fourier filter suppresses directional artifacts by weighting image contributions based - on their scan angles, utilizing a bounded sine sigmoid for smooth transition. - - Upsampling enhances interpolation precision but may increase computational cost. - """ + scan_h, scan_w = original_shape[0], original_shape[1] + n_channels = 1 + for d in range(2, ndim): + n_channels *= original_shape[d] - # init - stack_corr = np.zeros( - ( - self.shape[0], - np.round(self.shape[1] * upsample_factor).astype("int"), - np.round(self.shape[2] * upsample_factor).astype("int"), - ) + device = torch.device(correction._device) + + # ── Drift from knots (canvas → raw frame) ── + idx = image_index % len(correction.knots) + # Validates preprocess+align ran. + knots = drift_knots.stage_knots(correction, stage)[idx] + knot_h = knots.shape[1] + if knot_h != scan_h: + raise ValueError( + f"Drift grid has {knot_h} rows but ds_4d has " + f"{scan_h} scan rows. Ensure reference image and ds_4d " + f"have matching scan dimensions (check padding / resize).") + + drift = drift_knots.interpolator(correction, idx, knots).drift_raw( + correction._initial_knots[idx] + ).to( + device=device, dtype=torch.float32 ) - weight_corr = np.zeros( - ( - self.shape[0], - np.round(self.shape[1] * upsample_factor).astype("int"), - np.round(self.shape[2] * upsample_factor).astype("int"), - ) + # K=1 → drift is (2, H), broadcast across columns. + # K>=2 → drift is (2, H, W), varies along fast axis. + if drift.ndim == 2: + drift_row = drift[0][:, None] + drift_col = drift[1][:, None] + else: + drift_row = drift[0] + drift_col = drift[1] + row_coords = torch.arange(scan_h, device=device, dtype=torch.float32) + col_coords = torch.arange(scan_w, device=device, dtype=torch.float32) + sample_row = row_coords[:, None].expand(scan_h, scan_w) - drift_row + sample_col = col_coords[None, :].expand(scan_h, scan_w) - drift_col + # ── Pre-compute warp grid ONCE (tiny: 1×H×W×2 f32) ── + warp_grid = torch.stack([ + 2.0 * sample_col / (scan_w - 1) - 1.0, + 2.0 * sample_row / (scan_h - 1) - 1.0, + ], dim=-1)[None] # (1, H, W, 2) + + # ── Flatten input to (H, W, C) view ── + flat = ( + torch.from_numpy(ds_4d.reshape(scan_h, scan_w, n_channels)) + if is_numpy + else ds_4d.reshape(scan_h, scan_w, n_channels) ) - if kde_sigma is None: - kde_sigma = self.kde_sigma - - # Update images - for ind in range(self.shape[0]): - stack_corr[ind], weight_corr[ind] = self.interpolator[ind].warp_image( - self.images[ind].array, - self.knots[ind], - kde_sigma=kde_sigma, - upsample_factor=upsample_factor, - ) + # Output dtype for device intermediates. + out_dt = torch.float32 + if output_dtype == "same": + if is_numpy and input_np_dtype is not None: + out_dt = torch.from_numpy( + np.empty(0, dtype=input_np_dtype) + ).dtype + elif not is_numpy: + out_dt = ds_4d.dtype + elif isinstance(output_dtype, torch.dtype): + out_dt = output_dtype + + # ── Target device ── + # Default the output to the input's device so the pipeline stays + # in place; explicit ``output_device`` overrides. + if use_external_output: + target = torch.device("cpu") + elif output_device is not None: + target = torch.device(output_device) + if target.type == "cuda": + target = device + elif ( + isinstance(ds_4d, torch.Tensor) + and (ds_4d.is_cuda or ds_4d.device.type == "mps") + ): + target = device + else: + target = torch.device("cpu") + return_numpy = ( + use_external_output + or (is_numpy and output_device is None) + ) - if fourier_filter: - # Apply fourier filtering - kx = np.fft.fftfreq(stack_corr.shape[1])[:, None] - ky = np.fft.fftfreq(stack_corr.shape[2])[None, :] - kt = np.arctan2(ky, kx) - - stack_fft = np.fft.fft2(stack_corr) - weights = np.zeros_like(stack_corr) - - for ind in range(stack_corr.shape[0]): - # Calculate weights as a function of angle - weights[ind] = np.abs( - np.mod((kt - self.scan_direction[ind]) / np.pi + 0.5, 1.0) - 0.5 - ) / (1 / 2) - weights[ind][0, 0] = 1.0 - - # Apply sigmoid to weighting function - weights[ind] = bounded_sine_sigmoid( - weights[ind], - midpoint=filter_midpoint, + # Choose a chunk size from available device memory. + if chunk_size is None: + bytes_per_ch = scan_h * scan_w * 4 + if device.type == "cuda": + try: + free_bytes, _ = torch.cuda.mem_get_info(device) + except RuntimeError: + free_bytes = 0 + else: + free_bytes = 0 + if target.type == "cuda" and not use_external_output: + out_elem = torch.tensor([], dtype=out_dt).element_size() + free_bytes = max( + 0, + free_bytes - n_channels * scan_h * scan_w * out_elem, ) + if device.type == "mps": + chunk_size = min(n_channels, 64) + elif device.type == "cuda": + chunk_size = min( + n_channels, + max(1, int(free_bytes * 0.7 / (bytes_per_ch * 2))), + ) + else: + chunk_size = min(n_channels, 64) - # Weight the fourier transformed images - stack_fft[ind] *= weights[ind] - - weights_sum = np.sum(weights, axis=0) - image_corr_fft = np.zeros_like(weights_sum, dtype=complex) - np.divide( - np.sum(stack_fft, axis=0), - weights_sum, - where=weights_sum > 0.0, - out=image_corr_fft, - ) - + # ── Allocate output ── + if use_external_output: + out_flat = output.reshape(scan_h, scan_w, n_channels) else: - image_corr_fft = np.fft.fft2(np.mean(stack_corr, axis=0)) - - if mask_output: - # Note that we compute 2 boolean masks to round off the corners of image blending - - # calculate mask from product of individual image masks - # scale weights by upsample factor to normalize to mean value of 1.0 - mask_edge = np.prod(weight_corr >= (weight_thresh / upsample_factor**2), axis=0) - # Set outermost pixels to False to define the boundary for edge blending - mask_edge[:, 0] = False - mask_edge[:, -1] = False - mask_edge[0, :] = False - mask_edge[-1, :] = False - # Find inner boundary mask - mask_inner = distance_transform_edt(mask_edge) <= mask_edge_blend - # compute mask using edge blending value - mask = ( - np.cos( - (np.pi / 2) - * np.clip(distance_transform_edt(mask_inner) / mask_edge_blend, 0.0, 1.0) - ) - ** 2 + internal_output = torch.empty( + scan_h, scan_w, n_channels, dtype=out_dt, device=target, ) - # Mean pad value - pad_value_mean = np.mean([ind.pad_value for ind in self.interpolator]) - # apply mask - image_corr_fft = np.fft.fft2( - np.fft.ifft2(image_corr_fft) * mask + pad_value_mean * (1 - mask) - ) - - if output_original_shape: - image_corr_fft = fourier_cropping(image_corr_fft, self.shape[-2:]) / upsample_factor**2 - # TODO - adjust origin / sampling if output sampling is different from input - # i.e. if output_original_shape is False, and upsample_factor > 1 - image_corr = Dataset2d.from_array( - np.real(np.fft.ifft2(image_corr_fft)), - name="drift corrected image", - origin=self.images[0].origin, - sampling=self.images[0].sampling, - units=self.images[0].units, - ) - image_corr.metadata.update( - { - "downsample": int(getattr(self, "downsample", 1)), - "downsample_method": getattr(self, "downsample_method", "none"), - "downsample_metadata": copy.deepcopy( - getattr(self, "downsample_metadata", {}) - ), - } + # ── Numpy dtype for external output conversion ── + if use_external_output: + _out_np_dtype = output.dtype + + # ── Vectorized grid_sample with pre-computed grid ── + chunks = range(0, n_channels, chunk_size) + num_chunks = len(chunks) + correction_progress = tqdm( + total=n_channels, + desc=progress_desc, + unit="channel", + disable=not verbose or num_chunks <= 1, ) + for start in chunks: + end = min(start + chunk_size, n_channels) + warped = F.grid_sample( + flat[:, :, start:end].permute(2, 0, 1).contiguous() + .to(device=device, dtype=torch.float32)[None], + warp_grid, + mode=mode, align_corners=True, padding_mode="border", + )[0].permute(1, 2, 0) - if show_image: - fig, ax = show_2d(image_corr.array, **kwargs) - # Force a render whether we're drawing into a provided Axes or a fresh Figure - ax_to_draw = kwargs.get("ax", ax) - try: - ax_to_draw.figure.canvas.draw_idle() - # If we're not drawing into a caller-provided Axes, also pop the window - if "ax" not in kwargs: - plt.show() - except Exception: - # Fallback: if backend is odd, try a blocking show - plt.show() - return image_corr + # Integer casts truncate. Round first or low-count detector pixels + # collapse to zero after interpolation. + is_int = isinstance(out_dt, torch.dtype) and not out_dt.is_floating_point + if is_int: + warped_cast = warped.round().clamp_( + torch.iinfo(out_dt).min, torch.iinfo(out_dt).max) + else: + warped_cast = warped + if use_external_output: + out_flat[:, :, start:end] = ( + warped_cast.cpu().numpy().astype(_out_np_dtype) + ) + else: + internal_output[:, :, start:end] = warped_cast.to( + device=target, dtype=out_dt, + ) + correction_progress.update(end - start) + correction_progress.close() + + if use_external_output: + return output + + result = internal_output.reshape(original_shape) + if return_numpy: + return result.detach().cpu().numpy() + return result -def calculate_error( +def apply_correction( self, - mode: int, - _warped_t: torch.Tensor | None = None, -): - """Compute per-image MAE against the mean and append to error history. + data: Dataset2d | Dataset3d | torch.Tensor | np.ndarray | None = None, + image_index: int = -1, + *, + stage: str | None = None, + mode: str = "bilinear", + chunk_size: int | None = None, + output_dtype: torch.dtype | np.dtype | str | None = None, + output_device: str | torch.device | None = None, + output: np.ndarray | None = None, + verbose: bool = True, +) -> torch.Tensor | np.ndarray: + """Apply the learned drift correction to an image or dataset. - Measures how well the warped images agree by computing the mean - absolute difference of each image from the stack mean. Without - error tracking, there is no way to verify that alignment steps - are actually improving the result. + Image collections use the last two axes as scan coordinates. Reference + and 4D-STEM datasets use the first two axes, so spectra and diffraction + patterns remain attached to their corrected probe positions. Large + datasets are processed in chunks and may be written directly into a + preallocated array. Parameters ---------- - mode : int - Stage identifier (0=preprocess, 1=affine, 2=nonrigid). - _warped_t : torch.Tensor or None - If provided, compute error from this tensor directly, - avoiding a GPU-to-CPU round-trip. + data : Dataset2d, Dataset3d, ndarray, or torch.Tensor, optional + Data to correct. If omitted, use the stored alignment image or + 4D-STEM dataset. + image_index : int, optional + Scan trajectory to apply. Default is the last scan. + stage : {"initial", "affine", "strip", None}, optional + Saved correction stage to apply. ``None`` uses the current solution. + mode : str, optional + Interpolation kernel, either ``"bilinear"`` or ``"bicubic"``. + Default is ``"bilinear"``. + chunk_size : int, optional + Number of detector or spectral channels corrected together. The + default selects a size that fits the available device memory. + output_dtype : torch.dtype, numpy dtype, or str, optional + Output precision for dataset correction. By default, preserve the + input precision. + output_device : str or torch.device, optional + Device for the returned tensor. NumPy inputs return NumPy arrays by + default. + output : ndarray, optional + Preallocated destination, such as a memory-mapped array. + verbose : bool, optional + Show progress for multi-chunk datasets. Default is ``True``. + + Returns + ------- + Dataset2d, Dataset3d, torch.Tensor, or ndarray + Corrected data with the same shape and axis order as the input. + QuantEM datasets retain calibration, units, and metadata. + + Examples + -------- + >>> dc = DriftCorrection(reference, moving, scan_direction_degrees=(0, 90)) + >>> dc.correct_affine(show_combined=False) + >>> corrected = dc.apply_correction(image_index=1) + """ - if _warped_t is not None: - images_mean = _warped_t.mean(dim=0) - sig_diff = torch.mean( - torch.abs(_warped_t - images_mean[None]), dim=(1, 2) - ).cpu().numpy() - else: - # Lazy refresh: align_nonrigid defers the warped→numpy sync until - # someone reads it, so calculate_error must trigger the refresh. - self._ensure_warped_images() - images_mean = np.mean(self.images_warped.array, axis=0) - sig_diff = np.mean( - np.abs(self.images_warped.array - images_mean[None, :, :]), axis=(1, 2) + if isinstance(data, (Dataset2d, Dataset3d)): + source = data + corrected = apply_correction( + self, + source.array, + image_index=image_index, + stage=stage, + mode=mode, + chunk_size=chunk_size, + output_dtype=output_dtype, + output_device=output_device, + output=output, + verbose=verbose, + ) + if isinstance(corrected, torch.Tensor): + corrected = corrected.detach().cpu().numpy() + return _corrected_dataset( + np.asarray(corrected), + dataset_info(source), ) - # Error vector - error_current = np.hstack((mode, np.mean(sig_diff), sig_diff)) + if not hasattr(self, "knots") or not hasattr(self, "_initial_knots"): + raise RuntimeError( + "apply_correction() requires preprocess() and correct_affine() " + "first. Run dc.preprocess().correct_affine() (and optionally " + ".correct_nonrigid()) before apply_correction()." + ) + valid_modes = {"bilinear", "bicubic"} + if mode not in valid_modes: + raise ValueError(f"mode must be one of {valid_modes}, got {mode!r}") + index = image_index % len(self.knots) + drift_knots.knot_delta_canvas(self, index) + is_4dstem_mode = getattr(self, "_datasets", None) is not None and not self._reference_mode + scan_h = self.imgs[index].shape[0] + scan_w = self.imgs[index].shape[1] + dataset_layout = data is None and is_4dstem_mode + if data is None: + if ( + getattr(self, "_built_from_datasets", False) + and getattr(self, "_datasets", None) is None + ): + raise ValueError( + "apply_correction() has no data: this instance was built " + "from a 4D-STEM / reference dataset, but save() dropped the " + "dataset (too large to serialize) and it was not re-attached " + "after load. Pass the dataset explicitly, e.g. " + "dc.apply_correction(data=my_4dstem_array), or re-attach it " + "before calling apply_correction()." + ) + if self._reference_mode: + data = self._datasets[1] + if data is not None: + ndim = data.ndim + shape = tuple(data.shape) + dataset_layout = ndim >= 4 + if ndim == 3: + cube_layout = shape[0] == scan_h and shape[1] == scan_w + batch_layout = shape[-2] == scan_h and shape[-1] == scan_w and not cube_layout + dataset_layout = cube_layout and (not batch_layout or is_4dstem_mode) + + if dataset_layout: + return apply_correction_to_dataset( + self, + data, + image_index=image_index, + stage=stage, + mode=mode, + chunk_size=chunk_size, + output_dtype=output_dtype, + output_device=output_device, + output=output, + verbose=verbose, + ) - # Initialize or append to error tracking array - if not hasattr(self, "error_track"): - self.error_track = error_current[None, :] # initialize with first row + if data is None: + data_t = self.imgs_t[index] + elif isinstance(data, np.ndarray): + data_t = torch.tensor(data, dtype=self._dtype, device=self._device) else: - self.error_track = np.vstack((self.error_track, error_current)) + data_t = data.to(device=self._device, dtype=self._dtype) + image_height = data_t.shape[-2] + knots = drift_knots.stage_knots(self, stage)[index] + knot_height = knots.shape[1] + if image_height != knot_height: + raise ValueError( + f"Input scan-row axis ({image_height}) does not match knot grid " + f"height ({knot_height}). For 4D-STEM mode the leading axis is " + "the scan row; for image collection mode the trailing-2 axes " + "are scan." + ) + drift = drift_knots.interpolator(self, index, knots).drift_raw( + self._initial_knots[index] + ) + return backward_warp(data_t, drift=drift, mode=mode) -def bounded_sine_sigmoid(x, midpoint=0.5, width=1.0): - """ - Piecewise bounded sigmoid: zero, raised sine squared, one. +def crop(self, image: NDArray, *, shape: str = "square") -> NDArray: + """Crop an image to the field measured by every corrected scan. Parameters ---------- - x : array-like, shape (...,) - Input values in [0, 1]. - midpoint : float - Center of the sigmoid transition. - width : float - Width of the sigmoid (range over which it ramps from 0 to 1). + image : numpy.ndarray + Corrected image or scan-axis-leading dataset. + shape : {"square", "rectangle"}, default "square" + Keep the largest centered square or the full common rectangle. + Returns ------- - y : array-like - Output in [0, 1], same shape as x. + numpy.ndarray + Cropped data with trailing channel or detector axes unchanged. + + Examples + -------- + >>> cropped = drift.crop(corrected.array, shape="rectangle") """ - x = np.asarray(x) - # Truncate width if midpoint too close to edge - left_max = midpoint - width / 2 - right_min = midpoint + width / 2 - if left_max < 0: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * midpoint}.", - RuntimeWarning, - ) - width = 2 * midpoint + rows, cols = crop_slices(self) + if shape == "square": + height = rows.stop - rows.start + width = cols.stop - cols.start + side = min(height, width) + row_start = rows.start + (height - side) // 2 + col_start = cols.start + (width - side) // 2 + rows = slice(row_start, row_start + side) + cols = slice(col_start, col_start + side) + elif shape != "rectangle": + raise ValueError(f'shape must be "rectangle" or "square", got {shape!r}') + # Scan axes are the LEADING two: an EDS cube is (row, col, channel), so + # trailing-axis indexing would slice width and channels instead of the + # scan field. 4D-STEM mode shares the same (row, col, ...) layout. + return np.asarray(image)[rows, cols, ...] + + +def crop_slices(self) -> tuple[slice, slice]: + """Return row and column slices for the common measured field of view. + + Use these slices when several related arrays must receive exactly the same + crop as the corrected image. + + Returns + ------- + tuple of slice + Row and column slices in ``(row, col)`` order. - if right_min > 1: - warnings.warn( - f"width={width} is too large for midpoint={midpoint}, " - f"clamping width to {2 * (1 - midpoint)}.", - RuntimeWarning, + Examples + -------- + >>> rows, cols = drift.crop_slices() + >>> cropped_spectrum = spectrum[rows, cols, :] + """ + if getattr(self, "_reference_mode", False) and hasattr(self, "_initial_knots"): + scan_h, scan_w = np.asarray(self.imgs[0].array).shape[:2] + pad = 4 + field = self.drift_field(1).detach().cpu().numpy() + row_drift = field[0].ravel() + col_drift = field[1].ravel() + top = max(0.0, float(row_drift.max())) + bottom = max(0.0, float(-row_drift.min())) + left = max(0.0, float(col_drift.max())) + right = max(0.0, float(-col_drift.min())) + row = slice( + int(np.ceil(top)) + pad, + scan_h - int(np.ceil(bottom)) - pad, ) - width = 2 * (1 - midpoint) - # Recalculate edges - left = midpoint - width / 2 - right = midpoint + width / 2 - - y = np.zeros_like(x, dtype=float) - in_band = (x >= left) & (x <= right) - # Map [left, right] to [0, pi/2] - t = (x[in_band] - left) / width # goes from 0 to 1 - y[in_band] = np.sin(t * np.pi / 2) ** 2 - y[x > right] = 1.0 - return y - - -def _bounded_sine_sigmoid_torch( - x: torch.Tensor, - midpoint: float = 0.5, - width: float = 1.0, -) -> torch.Tensor: - width = min(width, 2 * midpoint, 2 * (1 - midpoint)) - left = midpoint - width / 2 - right = midpoint + width / 2 - t = ((x - left) / width).clamp(0.0, 1.0) - return torch.where(x > right, torch.ones_like(x), torch.sin(t * (np.pi / 2)) ** 2) + col = slice( + int(np.ceil(left)) + pad, + scan_w - int(np.ceil(right)) - pad, + ) + return row, col + mask = coverage_mask(self) + scan_h, scan_w = mask.shape + pad = 4 + row_0, row_1, col_0, col_1 = largest_rectangle(mask) + row = slice(min(row_0 + pad, scan_h), max(row_1 - pad, 0)) + col = slice(min(col_0 + pad, scan_w), max(col_1 - pad, 0)) + return row, col -def _fourier_crop_torch( - fft_array: torch.Tensor, - crop_shape: tuple[int, int], -) -> torch.Tensor: - crop_h, crop_w = crop_shape - h1 = crop_h // 2 - h2 = crop_h - h1 - w1 = crop_w // 2 - w2 = crop_w - w1 - result = torch.zeros(crop_shape, dtype=fft_array.dtype, device=fft_array.device) - result[:h1, :w1] = fft_array[:h1, :w1] - result[:h1, -w2:] = fft_array[:h1, -w2:] - result[-h2:, :w1] = fft_array[-h2:, :w1] - result[-h2:, -w2:] = fft_array[-h2:, -w2:] - return result +def coverage_mask(self) -> np.ndarray: + """Identify pixels supported by every corrected scan. + + The mask separates measured overlap from padded canvas pixels, making NCC + and other comparisons use the same physical field of view. + + Returns + ------- + numpy.ndarray + Boolean mask in the original scan frame. + + Examples + -------- + >>> common_pixels = drift.coverage_mask() + >>> ncc = compare(first[common_pixels], second[common_pixels]) + """ + canvas_h, canvas_w = self.imgs_warped.array.shape[-2:] + knot_counts = {int(knots.shape[2]) for knots in self.knots} + if knot_counts != {1}: + knots = torch.stack( + [value.detach() for value in self.knots] + ).to(device=self._device, dtype=self._dtype) + _, weights = warp_and_translate( + self, + max_image_shift=None, + knots_batch=knots, + solve_translation=False, + return_weights=True, + ) + common = (weights >= 1e-3).all(dim=0).cpu().numpy() + scan_h, scan_w = self.imgs[0].shape[:2] + offset_row = (canvas_h - scan_h) // 2 + offset_col = (canvas_w - scan_w) // 2 + return common[ + offset_row : offset_row + scan_h, + offset_col : offset_col + scan_w, + ] + + common = np.ones((canvas_h, canvas_w), dtype=bool) + for index in range(len(self.knots)): + knots_full = self.knots[index].detach().cpu().numpy() + knots = knots_full[:, :, 0] + fast = np.asarray(self.scan_fast[index], dtype=float) + width = int(self.imgs[index].shape[1]) + position = np.arange(width, dtype=float) + rows = np.round(knots[0][:, None] + fast[0] * position[None, :]).astype(int) + cols = np.round(knots[1][:, None] + fast[1] * position[None, :]).astype(int) + footprint = np.zeros((canvas_h, canvas_w), dtype=bool) + inside = (rows >= 0) & (rows < canvas_h) & (cols >= 0) & (cols < canvas_w) + footprint[rows[inside], cols[inside]] = True + footprint = ndi_binary_closing(footprint, structure=np.ones((3, 3), dtype=bool)) + common &= footprint + scan_h = int(self.imgs[0].shape[0]) + scan_w = int(self.imgs[0].shape[1]) + offset_row = (canvas_h - scan_h) // 2 + offset_col = (canvas_w - scan_w) // 2 + return common[ + offset_row : offset_row + scan_h, + offset_col : offset_col + scan_w, + ] + + +def warped_stack(correction, stage: str | None = None) -> np.ndarray: + """Return aligned scans at one checkpoint without changing the solved object. + + This is the shared data path behind stage-aware figures, reports, and + ``corrected(merge=False)``. Historical checkpoints are rendered directly + from their saved knots, so asking for an affine result after non-rigid + refinement cannot alter the final correction or its cache. + """ + if stage in (None, "nonrigid", "non-rigid"): + ensure_warped_images(correction) + return np.asarray(correction.imgs_warped.array, dtype=np.float32) + + knots = drift_knots.stage_knots(correction, stage) + if not correction._reference_mode: + warped = warp_and_translate( + correction, + max_image_shift=None, + upsample_factor=8, + knots_batch=torch.stack([k.detach() for k in knots]), + solve_translation=False, + ) + return warped.detach().cpu().numpy().astype(np.float32, copy=False) + + stack = np.stack(reference_scan_stack(correction, knots)).astype( + np.float32, copy=False + ) + canvas = np.empty((2, *correction.shape[1:]), dtype=np.float32) + row = (correction.shape[1] - stack.shape[1]) // 2 + col = (correction.shape[2] - stack.shape[2]) // 2 + for index in range(2): + canvas[index].fill(float(correction.pad_value[index])) + canvas[ + index, + row : row + stack.shape[1], + col : col + stack.shape[2], + ] = stack[index] + return canvas + + +def comparison_panels(correction, stage: str | None = None) -> dict: + """Compare raw and corrected scans in the first acquisition's frame.""" + num_scans = correction.shape[0] + angles = np.asarray(correction.scan_direction_degrees, dtype=float) + quarter = [int(round(-angle / 90.0)) % 4 for angle in angles] + reference = quarter[0] + raw = [ + np.rot90( + np.asarray(correction.imgs[index].array), + (quarter[index] - reference) % 4, + ) + for index in range(num_scans) + ] + corrected = np.rot90( + warped_stack(correction, stage), + (-reference) % 4, + axes=(1, 2), + ) + raw_combined = sum(image.astype(np.float32) for image in raw) / num_scans + relative_angles = [ + ((float(angle) - float(angles[0]) + 180.0) % 360.0) - 180.0 + for angle in angles + ] + names = ["0deg"] + [ + f"{int(round(relative_angles[index]))}deg" for index in range(1, num_scans) + ] + raw_labels = [f"{names[0]} scan"] + [ + f"{names[index]} scan -> 0deg frame" for index in range(1, num_scans) + ] + images = raw + [raw_combined] + list(corrected) + [corrected.mean(0)] + labels = ( + raw_labels + + ["combined scan"] + + [f"corrected {name}" for name in names] + + ["corrected combined scan"] + ) + sampling = np.asarray(correction.imgs[0].sampling, dtype=float) + unit = correction.imgs[0].units[0] if getattr(correction.imgs[0], "units", None) else "pixels" + return { + "images": images, + "labels": labels, + "ncols": num_scans + 1, + "raw_scans": raw, + "raw_combined": raw_combined, + "corrected_scans": list(corrected), + "corrected_combined": corrected.mean(0), + "pixel_size": float(sampling[0]), + "pixel_unit": unit, + } diff --git a/src/quantem/imaging/drift/core/__init__.py b/src/quantem/imaging/drift/core/__init__.py index b9253d00..d0e9e8b6 100644 --- a/src/quantem/imaging/drift/core/__init__.py +++ b/src/quantem/imaging/drift/core/__init__.py @@ -1 +1,6 @@ -"""Numerical stages for drift correction.""" +"""Numerical core for drift-correction alignment stages. + +The public workflow remains :class:`quantem.imaging.drift.DriftCorrection`. +This package groups affine, strip, and non-rigid implementations with their +shared warping and scanline-knot mathematics. +""" diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index a4041f4d..502f9a90 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -1,88 +1,375 @@ +"""Affine drift-rate search and scanline-knot updates.""" + +import time import numpy as np import torch from torch.fft import fftfreq +from tqdm import tqdm -from quantem.imaging.drift import diagnostics, preparation -from quantem.imaging.drift.core.knots import ( - bilinear_kde_batch, - transform_coordinates_single_knot, -) +import quantem.imaging.drift.apply as drift_apply +import quantem.imaging.drift.plot as drift_plot +import quantem.imaging.drift.preparation as preparation +import quantem.imaging.drift.report as report +from quantem.imaging.drift.core import knots as drift_knots from quantem.imaging.drift.core.warping import ( + backward_warp_grid_search, cross_corr_batch, - translate_align, + fixed_overlap_ncc, + translate_align_pair_batch, + warp_and_translate, ) -def align_affine( +def drift_rate(correction) -> tuple[float, float]: + """Return affine drift rate ``(row, col)`` in pixels per scanline. + + The affine checkpoint is used when later strip or non-rigid refinement has + modified the live knots, keeping this measurement specific to the linear + drift model. + """ + if not hasattr(correction, "_initial_knots"): + raise RuntimeError("Call preprocess() then correct_affine() first.") + index = len(correction.knots) - 1 + if hasattr(correction, "_knots_after_affine"): + delta = ( + correction._knots_after_affine[index] + - correction._initial_knots[index] + ) + else: + delta = drift_knots.knot_delta_canvas(correction, index) + lines = delta.shape[1] + row = float( + (delta[0, -1, 0] - delta[0, 0, 0]) / max(lines - 1, 1) + ) + column = float( + (delta[1, -1, 0] - delta[1, 0, 0]) / max(lines - 1, 1) + ) + return row, column + + +def validate_num_rates(num_rates: int) -> int: + """Validate the number of candidate rates sampled along each drift axis.""" + width = int(num_rates) + if width < 3 or width % 2 == 0: + raise ValueError( + "num_rates must be an odd integer >= 3 so the search includes " + f"zero drift; got {num_rates!r}." + ) + return width + + +def candidate_grid( + center: np.ndarray, + radius: float, + width: int, + *, + circular: bool, +) -> np.ndarray: + """Build a small 2-D drift-rate grid around ``center``.""" + axis = np.linspace(-radius, radius, width, dtype=np.float64) + row, col = np.meshgrid(axis, axis, indexing="ij") + keep = row**2 + col**2 <= (radius * 1.001) ** 2 if circular else np.ones_like(row, dtype=bool) + return center[None, :] + np.column_stack((row[keep], col[keep])) + + +def _apply_affine_rate(correction, rate, fixed_set: frozenset[int]) -> None: + """Apply one ``(row, col)`` drift rate to every free scan.""" + rate = torch.as_tensor( + rate, + dtype=correction.knots[0].dtype, + device=correction.knots[0].device, + ) + for image_index in range(correction.shape[0]): + if image_index not in fixed_set: + drift_knots.interpolator(correction, image_index).apply_affine_shift(rate) + + +def _scan_disagreement(warped: torch.Tensor, scan_axis: int) -> torch.Tensor: + """Return absolute differences from the mean corrected scan.""" + mean = warped.mean(dim=scan_axis, keepdim=True) + return torch.abs(warped - mean) + + +def _downsampled_correction(correction, factor): + """Build the same correction problem on average-pooled scan images.""" + images = [ + preparation.average_downsample_2d(np.asarray(image.array), factor) + for image in correction.imgs + ] + if correction._reference_mode: + pyramid = type(correction).from_reference( + images[0], + images[1], + scan_direction_degrees=float(correction.scan_direction_degrees[1]), + ) + else: + pyramid = type(correction).from_images( + *images, + scan_direction_degrees=tuple(correction.scan_direction_degrees), + ) + pyramid.preprocess( + padding_fraction=correction.pad_fraction, + padding_value="median", + smoothing_sigma=correction.kde_sigma, + num_knots=1, + normalize=False, + verbose=False, + ) + return pyramid + + +def _delivered_candidate( + correction, + rate, + starting_knots, + fixed_set, + max_image_shift, + upsample_factor, + bridge=None, +): + """Score one affine rate after the translation applied to the final image.""" + correction.knots = [knot.clone() for knot in starting_knots] + if bridge is not None: + _apply_affine_rate(correction, bridge, fixed_set) + warp_and_translate( + correction, + max_image_shift, + upsample_factor, + fixed_indices=fixed_set, + ) + rate = rate - bridge + _apply_affine_rate(correction, rate, fixed_set) + warped = warp_and_translate( + correction, + max_image_shift, + upsample_factor, + fixed_indices=fixed_set, + ) + return ( + float(_scan_disagreement(warped, 0).mean().cpu()), + [knot.clone() for knot in correction.knots], + warped.clone(), + None, + ) + + +def _delivered_candidates( + correction, + rates, + starting_knots, + fixed_set, + max_image_shift, + upsample_factor, +): + """Score a two-image affine neighborhood together when memory permits.""" + def sequential(): + return [ + _delivered_candidate( + correction, + rate, + starting_knots, + fixed_set, + max_image_shift, + upsample_factor, + ) + for rate in rates + ] + + if correction.shape[0] != 2 or torch.device(correction._device).type != "cuda": + return sequential(), set() + + free_bytes, _ = torch.cuda.mem_get_info(correction._device) + bytes_per_element = torch.finfo(correction._dtype).bits // 8 + estimated_bytes = ( + len(rates) + * correction.shape[1] + * correction.shape[2] + * bytes_per_element + * 64 + ) + if estimated_bytes > free_bytes * 0.4: + return sequential(), {"memory"} + + correction.knots = [knot.clone() for knot in starting_knots] + rates_t = torch.as_tensor( + rates, + dtype=correction._dtype, + device=correction._device, + ) + canvas_shape = (correction.shape[1], correction.shape[2]) + first_warps = [] + candidate_coordinates = [] + for image_index in range(correction.shape[0]): + row_base, col_base, scanline_offset = drift_knots.interpolator( + correction, image_index + ).affine_candidate_base() + row_candidates = ( + row_base[None] + + rates_t[:, 0, None, None] * scanline_offset[None, :, None] + ) + col_candidates = ( + col_base[None] + + rates_t[:, 1, None, None] * scanline_offset[None, :, None] + ) + candidate_coordinates.append((row_candidates, col_candidates)) + candidate_warped, _ = drift_knots.bilinear_kde_batch( + row_candidates, + col_candidates, + correction.imgs_t[image_index], + canvas_shape, + correction.kde_sigma, + correction.pad_value[image_index], + ) + first_warps.append(candidate_warped) + shifts = translate_align_pair_batch( + torch.stack(first_warps, dim=1), + upsample_factor, + max_image_shift, + ) + final_warps = [] + for image_index, (row_candidates, col_candidates) in enumerate( + candidate_coordinates + ): + candidate_warped, _ = drift_knots.bilinear_kde_batch( + row_candidates + shifts[:, image_index, 0, None, None], + col_candidates + shifts[:, image_index, 1, None, None], + correction.imgs_t[image_index], + canvas_shape, + correction.kde_sigma, + correction.pad_value[image_index], + ) + final_warps.append(candidate_warped) + warped_batch = torch.stack(final_warps, dim=1) + costs = _scan_disagreement(warped_batch, 1).mean(dim=(1, 2, 3)) + ranked = torch.sort(costs).values + if len(rates) > 1 and float(ranked[1] - ranked[0]) <= max( + 1e-7, abs(float(ranked[0])) * 1e-6 + ): + return sequential(), {"batch", "tie"} + return ( + [ + ( + float(costs[index].cpu()), + None, + warped_batch[index].clone(), + shifts[index].clone(), + ) + for index in range(len(rates)) + ], + {"batch"}, + ) + + +def correct_affine( self, - step: float = 0.01, - num_tests: int = 9, + *, + max_drift_rate: float | None = None, + num_rates: int | None = None, refine: bool = True, - upsample_factor: int = 8, - max_image_shift: float | None = 32, - chunk_size: int | None = None, - downsample: int | str = 1, - show_merged: bool = True, - show_images: bool = False, + max_image_shift: float | None | str = "auto", + fixed_scans: list[int] | None = None, + region: str | tuple[int, int, int, int] | None = None, + region_smoothing_sigma: float = 4.0, + show_combined: bool = True, + show_scans: bool = False, show_knots: bool = True, - verbose: bool = False, - **kwargs, + show_knot_plot: bool = False, + show_report: bool = False, + verbose: bool = True, + downsample: int | str = "auto", + chunk_size: int | None = None, ): - """Correct affine drift between scan pairs using a batched grid search. + """Automatically correct the dominant linear drift between scans. + + With no search parameters, QuantEM prepares the scanline-knot canvas, + chooses the minimum safe padding and translation bound, selects a + broad-search downsampling factor, expands the drift-rate range when + needed, and refines the winning candidate at full coordinate scale. + It then updates the scanline knots; the corrected output itself is not + downsampled. - Builds a grid of candidate linear-drift vectors, warps both images - for each candidate, and picks the one with the lowest cross-correlation - cost. An optional refinement pass subdivides the winning cell for - sub-step accuracy. Without affine correction, per-scanline drift - causes shear distortion that translation alignment alone cannot fix. + Advanced users may set ``max_drift_rate`` and ``num_rates`` together + to reproduce an explicit candidate grid. Numerical controls such as + ``downsample`` and ``chunk_size`` normally remain automatic. Affine + correction is needed because translation alone cannot remove the + shear caused by a changing position from one scanline to the next. Parameters ---------- - step : float - Search resolution in pixels per scan line. The grid search - tests drift rates from ``-step * num_tests/2`` to - ``+step * num_tests/2`` px/line. For example, ``step=0.02`` - with ``num_tests=11`` searches drifts from -0.10 to +0.10 - px/line. Smaller values detect subtler drift but test more - candidates. - num_tests : int - Number of drift rates to test along each axis. Must be odd - so the grid is centered on zero drift. Total candidates - ≈ ``π/4 * num_tests²``: ``num_tests=5`` → 21, - ``num_tests=9`` → 61, ``num_tests=11`` → 97. + max_drift_rate : float or None, default None + Search scale for the row and column drift-rate components, in + pixels per scanline. QuantEM first samples each component from + ``-max_drift_rate`` to ``+max_drift_rate``, then removes the corner + combinations outside a circular 2-D search region. Increase this + when the best candidate reaches the search boundary. Leave this + and ``num_rates`` unset for the automatic image-pyramid search. + num_rates : int or None, default None + Number of candidate drift rates sampled along each of the row and + column axes, including both bounds and zero. Must be odd. This is + a per-axis count, not the final number of 2-D candidates: QuantEM + first forms all ``num_rates ** 2`` ``(row_rate, column_rate)`` + combinations, then removes the square grid's corners with a + circular mask. The mask avoids evaluating diagonal candidates + whose combined drift magnitude is much larger than the requested + per-axis search scale. For example, ``num_rates=11`` forms + ``11 * 11 = 121`` combinations; the circular mask removes 24 + corner vectors, leaving 97 candidates to evaluate. Leave this and + ``max_drift_rate`` unset for automatic range expansion and local + refinement. refine : bool - If True, run a second pass at ``step / (num_tests - 1)`` - resolution, centered on the coarse winner. - upsample_factor : int - Sub-pixel precision for measuring the translational shift - between warped image pairs. 8 means 1/8-pixel precision. - Higher values are more accurate but slower. - max_image_shift : float or None + If True, run a second, finer search centered on the coarse winner. + max_image_shift : float, None, or "auto", default "auto" Maximum allowed translational shift in pixels. Cross-correlation peaks beyond this radius are masked to reject spurious matches - from noise or periodic artifacts. Set to None to allow any shift. + from noise or periodic artifacts. ``"auto"`` derives the bound + from image geometry. Set to None to allow any shift. chunk_size : int or None Number of candidates per pass. If None, all candidates at once. Set to a smaller value if you run out of memory. - downsample : int or "auto", default 1 - Exact average-pooling factor used only to find the broad affine - candidate. The accepted rate is transferred back to the native grid - and refined there; corrected output is never downsampled. - show_merged : bool - Display the merged (averaged) image after alignment. - show_images : bool + downsample : int or "auto", default "auto" + Average-pooling factor used only for the automatic broad affine + search. ``"auto"`` selects the largest exact divisor up to 8. + Set ``4``, ``2``, or ``1`` to retain more native detail when + validating a difficult or highly periodic specimen. This does not + downsample the corrected output or change the knot coordinates. + fixed_scans : list[int] or None + Indices of images whose knots should never be modified. + Use ``fixed_scans=[0]`` for single-sided alignment where + image 0 is a fixed reference (e.g. a merged HAADF) and only + the remaining images are optimized. When ``None`` (default), + all images receive the affine drift correction - the standard + behavior for 0°/90° scan pairs. + region : str, tuple of int, or None + Region used to estimate the affine drift. The fitted affine model is + still applied to the complete scans. Leave as None for the standard + whole-image search. For periodic lattices, use ``diagnose_affine()`` + to identify a region containing distinctive defects, then pass its + name (``"top_left"``, ``"top_right"``, ``"bottom_left"``, or + ``"bottom_right"``) or + ``(row_start, row_stop, column_start, column_stop)`` bounds. + region_smoothing_sigma : float, default 4.0 + Gaussian smoothing used only when ``region`` is set. Smoothing helps + the regional search follow distinctive defect structure instead of + selecting a neighboring periodic lattice peak. + show_combined : bool + Display the combined RGB comparison after alignment. + show_scans : bool Display each individual warped image after alignment. - show_knots : bool - Overlay knot positions on the displayed images. + show_knots : bool, default True + Show knot positions on top of the combined/per-scan plots + (cheap, useful diagnostic). + show_knot_plot : bool, default False + Render the standalone 2-panel knot trajectory + per-row delta + chart via ``dc.plot_knots()`` after this step. + show_report : bool, default False + Print a screenshot-friendly common/top/middle/bottom NCC table + comparing the before and affine checkpoints. verbose : bool - If True, print the top 5 candidate drift vectors with their - cost and direction after the grid search. Useful for - diagnosing ambiguous alignments or verifying the winning - candidate has a clear margin over runner-ups. - **kwargs - Additional keyword arguments passed to the plotting functions. - + If True, show candidate progress when the search requires + multiple batches, then print the top 5 drift vectors with their + cost and direction. Useful for diagnosing ambiguous alignments or + verifying the winning candidate has a clear margin over runner-ups. Returns ------- DriftCorrection @@ -90,46 +377,162 @@ def align_affine( Examples -------- - >>> drift = DriftCorrection.from_data( - ... images=[im0, im1], scan_direction_degrees=[0, 90]) - >>> drift.preprocess().align_affine(step=0.02, num_tests=11) + >>> drift = DriftCorrection( + ... im0, im1, scan_direction_degrees=[0, 90]) + >>> drift.correct_affine() + + Exact reproduction of a historical explicit grid: + + >>> drift.correct_affine(max_drift_rate=0.10, num_rates=11) + + Diagnose a periodic lattice, then anchor the affine fit to a region with + distinctive structure: + + >>> figure, regions = drift.diagnose_affine(stage="initial") + >>> drift.correct_affine(region="top_left") + + Single-sided alignment (4D-STEM VDF against a fixed HAADF reference): + + >>> drift = DriftCorrection( + ... haadf_ref, vdf, scan_direction_degrees=[0, 0]) + >>> drift.correct_affine(fixed_scans=[0]) """ - if self.shape[0] < 2: + # Reference-mode auto-anchors the reference image (index 0) so the + # user doesn't repeat what they declared via the constructor reference mode. + if fixed_scans is None and self._reference_mode: + fixed_scans = [0] + fixed_set = frozenset(fixed_scans) if fixed_scans is not None else frozenset() + automatic = max_drift_rate is None and num_rates is None + if (max_drift_rate is None) != (num_rates is None): raise ValueError( - f"align_affine requires at least 2 images (got {self.shape[0]}). " - f"Provide image pairs with different scan directions." + "Set both max_drift_rate and num_rates for an explicit grid, " + "or leave both unset for automatic affine alignment." ) - if num_tests % 2 == 0: + if not automatic and downsample != "auto": raise ValueError( - f"num_tests must be odd (got {num_tests}). Try {num_tests + 1}." + "downsample only applies when max_drift_rate and " + "num_rates are left unset for automatic affine alignment." ) - fixed_reference = getattr(self, "_reference_mode", False) - pyramid_factor = preparation.resolve_downsample( - downsample, - tuple(self.images[0].shape), - ) - if pyramid_factor > 1: - return _align_affine_pyramid( + if region is None and region_smoothing_sigma != 4.0: + raise ValueError( + "region_smoothing_sigma only applies when region is set. " + "Pass region='top_left' or custom row/column bounds, or leave " + "region_smoothing_sigma at its default." + ) + if region is not None: + _correct_affine_region( self, - factor=pyramid_factor, - step=step, - num_tests=num_tests, + region=region, + smoothing_sigma=region_smoothing_sigma, + max_drift_rate=max_drift_rate, + num_rates=num_rates, refine=refine, - upsample_factor=upsample_factor, max_image_shift=max_image_shift, + fixed_scans=fixed_scans, + verbose=verbose, + downsample=downsample, chunk_size=chunk_size, - show_merged=show_merged, - show_images=show_images, + ) + drift_plot.show_after_step( + self, + "affine", + show_combined=show_combined, + show_scans=show_scans, show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + if show_report: + print(self.report().to_string()) + return self + if not hasattr(self, "_initial_knots"): + preparation_start = time.perf_counter() + planned_rate = 0.25 if automatic else abs(float(max_drift_rate)) + translation_margin = ( + min(self.imgs[0].shape[:2]) * 0.125 + if self._built_from_datasets and not self._reference_mode + else 0.0 + ) + if self._reference_mode: + normalize = True + normalization_reason = "reference_mode" + elif self._built_from_datasets: + normalize = False + normalization_reason = "4dstem_collection" + else: + normalize, normalization_reason = preparation.automatic_alignment_normalization( + self.imgs + ) + padding = preparation.minimum_affine_padding_fraction( + tuple(int(value) for value in self.imgs[0].shape[:2]), + self.scan_direction_degrees, + planned_rate, + translation_margin, + ) + self.preprocess( + padding_fraction=padding, + normalize=normalize, + verbose=False, + show_combined=False, + show_scans=False, + show_knots=False, + ) + self._implicit_preprocess_seconds = time.perf_counter() - preparation_start + self.preprocess_info.update( + { + "padding_mode": "implicit_auto", + "planned_max_drift_rate": planned_rate, + "translation_margin": translation_margin, + "normalization_mode": "implicit_auto", + "normalization_reason": normalization_reason, + "seconds": self._implicit_preprocess_seconds, + } + ) + if verbose: + print( + "correct_affine: prepared scanline knots automatically; " + f"padding={padding:.4g}, canvas={self.shape[1:]}, " + f"{self._implicit_preprocess_seconds:.2f} s" + ) + else: + self._implicit_preprocess_seconds = 0.0 + if self.shape[0] < 2: + raise ValueError( + f"correct_affine requires at least 2 images (got {self.shape[0]}). " + f"Provide image pairs with different scan directions." + ) + if automatic: + if not refine: + raise ValueError("refine=False is only available for an explicit affine grid.") + return automatic_affine_search( + self, + fixed_set=fixed_set, + max_image_shift=max_image_shift, + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + show_knot_plot=show_knot_plot, + show_report=show_report, verbose=verbose, - kwargs=kwargs, + upsample_factor=8, + chunk_size=chunk_size, + pyramid_downsample=downsample, ) - + # Translation-peak refinement is an internal numerical detail. Eight + # is accurate (~0.01 px after the final parabolic fit) and is not a + # meaningful microscope control, so the public API does not expose it. + upsample_factor = 8 + num_tests = validate_num_rates(num_rates) + if max_image_shift == "auto": + max_image_shift = 256.0 # Build candidate grid with circular mask (~21% fewer than square) grid_axis = np.arange(-(num_tests - 1) / 2, (num_tests + 1) / 2) row_grid, col_grid = np.meshgrid(grid_axis, grid_axis, indexing="ij") circular_mask = row_grid**2 + col_grid**2 <= (num_tests / 2) ** 2 - drift_vectors = np.vstack((row_grid[circular_mask], col_grid[circular_mask])).T * step + drift_rate_step = max_drift_rate / ((num_tests - 1) / 2) + drift_vectors = ( + np.vstack((row_grid[circular_mask], col_grid[circular_mask])).T * drift_rate_step + ) def _print_top_candidates(label, candidates, costs_tensor): costs_np = costs_tensor.cpu().numpy() @@ -141,232 +544,956 @@ def _print_top_candidates(label, candidates, costs_tensor): drift_row, drift_col = candidates[idx] magnitude = np.sqrt(drift_row**2 + drift_col**2) gap = (costs_np[idx] - best_cost) / best_cost * 100 if rank > 0 else 0 - print(f" drift=({drift_row:+.4f}, {drift_col:+.4f}) px/line " - f"({magnitude:.4f} magnitude), cost={costs_np[idx]:.4f}" - f"{f' (+{gap:.1f}%)' if rank > 0 else ' (best)'}") - - def _apply_drift(drift_vec): - for img_idx in range(self.shape[0]): - if fixed_reference and img_idx == 0: - continue - scanline_offset = np.arange(self.knots[img_idx].shape[1]) - (self.knots[img_idx].shape[1] - 1) / 2 - self.knots[img_idx][0] += drift_vec[0] * scanline_offset[:, None] - self.knots[img_idx][1] += drift_vec[1] * scanline_offset[:, None] - - def _search_and_apply(candidates, label): - best_idx, costs = self._affine_grid_search_batch( - candidates, + print( + f" drift=({drift_row:+.4f}, {drift_col:+.4f}) px/line " + f"({magnitude:.4f} magnitude), cost={costs_np[idx]:.4f}" + f"{f' (+{gap:.1f}%)' if rank > 0 else ' (best)'}" + ) + + def _search_and_apply(candidates, label, accumulated_drift=None): + # When fixed_indices is set, backward_warp_grid_search scores + # absolute drift rates on the raw images (not canvas-warped). + # After the coarse pass, _apply_drift bakes the coarse drift into + # the knots, but the raw images are unchanged - so the refine + # candidates (small deltas) must be offset by the accumulated + # drift so that backward_warp_grid_search tests the correct total + # drift rates. + search_candidates = candidates + if fixed_set and accumulated_drift is not None: + search_candidates = candidates + accumulated_drift[None, :] + best_idx, costs = grid_search_batch( + self, + search_candidates, upsample_factor, max_image_shift, chunk_size, - fixed_reference=fixed_reference, + fixed_indices=fixed_set, + progress_desc=f"Affine {label.lower()}" if verbose else None, ) - _apply_drift(candidates[best_idx]) + _apply_affine_rate(self, candidates[best_idx], fixed_set) if verbose: _print_top_candidates(label, candidates, costs) - warped_t = self._warp_and_translate_torch( - max_image_shift, - upsample_factor, - fixed_reference=fixed_reference, + warped_t = warp_and_translate( + self, + max_image_shift, upsample_factor, fixed_indices=fixed_set ) - self.calculate_error(1, _warped_t=warped_t) - return candidates[best_idx] + report.record_error(self, 1, warped_t) + + # Confidence: cost gap between best and runner-up (%) + costs_np = costs.cpu().numpy() + ranked = np.argsort(costs_np) + best_cost = costs_np[ranked[0]] + runner_up = costs_np[ranked[1]] if len(ranked) > 1 else best_cost + margin = (runner_up - best_cost) / (best_cost + 1e-12) * 100 + return candidates[best_idx], margin - drift_total = _search_and_apply(drift_vectors, "Coarse search") + drift_total, coarse_margin = _search_and_apply(drift_vectors, "Coarse search") if refine: drift_fine = drift_vectors / (num_tests - 1) - drift_total = drift_total + _search_and_apply(drift_fine, "Refine search") + dt, refine_margin = _search_and_apply( + drift_fine, "Refine search", accumulated_drift=drift_total + ) + drift_total = drift_total + dt + self.affine_confidence_margin = refine_margin + else: + self.affine_confidence_margin = coarse_margin if verbose: - num_rows = self.images[0].shape[0] + num_rows = self.imgs[0].shape[0] drift_rate = np.sqrt(drift_total[0] ** 2 + drift_total[1] ** 2) total_shift = drift_rate * num_rows angle_deg = np.degrees(np.arctan2(drift_total[1], drift_total[0])) - print(f"align_affine: step={step}, num_tests={num_tests} " - f"({len(drift_vectors)} candidates), refine={refine}, " - f"max_image_shift={max_image_shift}") - msg = (f"Drift: ({drift_total[0]:+.4f}, {drift_total[1]:+.4f}) px/line, " - f"{drift_rate:.4f} magnitude, {angle_deg:.1f} deg, " - f"{total_shift:.1f} px total over {num_rows} lines") - if self.images[0].sampling is not None: - px_size = self.images[0].sampling[0] - unit = self.images[0].units[0] if self.images[0].units else "px" + print( + f"correct_affine: max_drift_rate={max_drift_rate:g}, " + f"num_rates={num_tests} per axis; " + f"{len(drift_vectors)} drift vectors evaluated; refine={refine}, " + f"max_image_shift={max_image_shift}" + ) + msg = ( + f"Drift: ({drift_total[0]:+.4f}, {drift_total[1]:+.4f}) px/line, " + f"{drift_rate:.4f} magnitude, {angle_deg:.1f}°, " + f"{total_shift:.1f} px total over {num_rows} lines" + ) + if self.imgs[0].sampling is not None: + px_size = self.imgs[0].sampling[0] + unit = self.imgs[0].units[0] if self.imgs[0].units else "px" msg += f" = {total_shift * px_size:.2f} {unit}" print(msg) err = self.error_track - print(f"Error: {err[0, 1]:.2f} -> {err[-1, 1]:.2f} " - f"({(err[0, 1] - err[-1, 1]) / err[0, 1] * 100:+.1f}%)") - - diagnostics._record_stage(self, "affine") - - # Plots - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images( - show_knots=show_knots, - title="Merged: affine", - **kwargs, - ) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: affine" for i in range(self.shape[0])], - **kwargs, + print( + f"Error: {err[0, 1]:.2f} -> {err[-1, 1]:.2f} " + f"({(err[0, 1] - err[-1, 1]) / err[0, 1] * 100:+.1f}%)" ) + margin = self.affine_confidence_margin + confidence = "high" if margin > 5 else "low" if margin < 2 else "moderate" + print(f"Confidence: {margin:.1f}% cost margin to runner-up ({confidence})") + drift_plot.show_after_step( + self, + "affine", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + self._knots_after_affine = [k.clone() for k in self.knots] + # Knots moved, so the cached warped stack no longer matches them. + # correct_strip and correct_nonrigid mark this too; without it every + # display after an affine-only solve draws the pre-alignment state. + self._images_warped_stale = True + if show_report: + print(self.report().to_string()) return self -def _affine_rate_from_knots(correction, image_index: int) -> np.ndarray: - """Measure one canvas drift-rate vector from a solved one-knot trajectory.""" - delta = ( - correction.knots[image_index] - - correction._initial_knots[image_index] +def _correct_affine_region( + self, + *, + region: str | tuple[int, int, int, int], + smoothing_sigma: float = 4.0, + max_drift_rate: float | None = None, + num_rates: int | None = None, + refine: bool = True, + max_image_shift: float | None | str = "auto", + fixed_scans: list[int] | None = None, + verbose: bool = True, + downsample: int | str = "auto", + chunk_size: int | None = None, +): + """Fit one affine model from a trusted region and apply it to both scans.""" + image_rows, image_columns = self.imgs[0].shape[:2] + middle_row = image_rows // 2 + middle_column = image_columns // 2 + quadrants = { + "top_left": (0, middle_row, 0, middle_column), + "top_right": (0, middle_row, middle_column, image_columns), + "bottom_left": (middle_row, image_rows, 0, middle_column), + "bottom_right": ( + middle_row, + image_rows, + middle_column, + image_columns, + ), + } + if isinstance(region, str): + if region not in quadrants: + raise ValueError( + f"You entered region={region!r}. Choose from " + f"{sorted(quadrants)} or provide four pixel bounds." + ) + bounds = quadrants[region] + region_name = region + else: + bounds = tuple(int(value) for value in region) + if len(bounds) != 4: + raise ValueError( + "region needs four bounds: " + "(row_start, row_stop, column_start, column_stop)." + ) + region_name = "custom" + row_start, row_stop, column_start, column_stop = bounds + if not ( + 0 <= row_start < row_stop <= image_rows + and 0 <= column_start < column_stop <= image_columns + ): + raise ValueError( + f"region bounds {bounds} are outside the image shape " + f"{(image_rows, image_columns)}." + ) + region_slice = ( + slice(row_start, row_stop), + slice(column_start, column_stop), ) - lines = delta.shape[1] - return ( - delta[:, -1, 0] - delta[:, 0, 0] - ) / max(lines - 1, 1) + self.preprocess( + padding_fraction=0.25, + smoothing_sigma=smoothing_sigma, + num_knots=1, + normalize=True, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + panels = drift_apply.comparison_panels(self, stage="initial") + angles = np.asarray(self.scan_direction_degrees, dtype=float) + relative_angles = (angles - angles[0] + 180.0) % 360.0 - 180.0 + quarter_turns = np.rint(-relative_angles / 90.0).astype(int) + regional_images = [ + np.ascontiguousarray( + np.rot90(image[region_slice], -quarter_turns[index]) + ) + for index, image in enumerate(panels["raw_scans"]) + ] + regional = type(self).from_images( + *regional_images, + scan_direction_degrees=tuple(relative_angles), + device=self.device, + ) + regional.preprocess( + padding_fraction=0.25, + smoothing_sigma=smoothing_sigma, + num_knots=1, + normalize=False, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + regional.correct_affine( + max_drift_rate=max_drift_rate, + num_rates=num_rates, + refine=refine, + max_image_shift=max_image_shift, + fixed_scans=fixed_scans, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=verbose, + downsample=downsample, + chunk_size=chunk_size, + ) -def _apply_affine_rate(correction, rate: np.ndarray) -> None: - """Apply one canvas drift rate to every non-reference scan.""" - fixed_reference = getattr(correction, "_reference_mode", False) - for image_index in range(correction.shape[0]): - if fixed_reference and image_index == 0: + solved_knots = [knots.clone() for knots in regional.knots] + regional.knots = [knots.clone() for knots in regional._initial_knots] + _apply_affine_rate(regional, regional.drift_rate, frozenset()) + translations = [ + (solved - rate_only).mean(dim=(1, 2)) + for solved, rate_only in zip( + solved_knots, + regional.knots, + strict=True, + ) + ] + fixed_set = frozenset(fixed_scans) if fixed_scans is not None else frozenset() + _apply_affine_rate(self, regional.drift_rate, fixed_set) + for index, (knots, translation) in enumerate( + zip(self.knots, translations, strict=True) + ): + if index in fixed_set: continue - lines = correction.knots[image_index].shape[1] - offset = np.arange(lines) - (lines - 1) / 2 - correction.knots[image_index][0] += rate[0] * offset[:, None] - correction.knots[image_index][1] += rate[1] * offset[:, None] + knots[0] += translation[0] + knots[1] += translation[1] + self._knots_after_affine = [knots.clone() for knots in self.knots] + self._images_warped_stale = True + self.affine_search_info = { + **regional.affine_search_info, + "strategy": "trusted_region", + "trusted_region": region_name, + "trusted_region_bounds_row_column": list(bounds), + "full_image_num_knots": 1, + } + return self -def _align_affine_pyramid( +@torch.inference_mode() +def automatic_affine_search( self, *, - factor: int, - step: float, - num_tests: int, - refine: bool, - upsample_factor: int, - max_image_shift: float | None, - chunk_size: int | None, - show_merged: bool, - show_images: bool, + fixed_set: frozenset[int], + max_image_shift: float | None | str, + show_combined: bool, + show_scans: bool, show_knots: bool, + show_knot_plot: bool, + show_report: bool, verbose: bool, - kwargs: dict, + upsample_factor: int, + chunk_size: int | None, + pyramid_downsample: int | str, ): - """Search a pooled grid, transfer its rate, and verify it natively.""" - pooled = [ - preparation.average_downsample_2d(image.array, factor) - for image in self.images - ] - if getattr(self, "_reference_mode", False): - pyramid = type(self).from_reference( - pooled[0], - pooled[1], - scan_direction_degrees=float(self.scan_direction_degrees[1]), - device=self.device, - ) + """Find the affine drift basin cheaply, then verify the delivered result. + + Broad rates are searched on an exact average-pooled pyramid. Small local + grids recover native-coordinate precision, and final candidates are ranked + after rigid translation because that is the corrected image users receive. + """ + # The four phases stay together because each one narrows or verifies the + # result of the phase before it. Read from top to bottom, the solver moves + # from a cheap broad search to the final image delivered to the scientist. + # Splitting this sequence across a one-use class would hide that numerical + # story behind temporary object state. + + # Preserve the untouched state so every fallback starts from the same scan + # geometry instead of accumulating shifts from a rejected candidate. + start_time = time.perf_counter() + starting_knots = [knot.clone() for knot in self.knots] + initial_error_row = np.asarray(self.error_track[-1], dtype=np.float64).copy() + image_shape = tuple(int(value) for value in self.imgs[0].shape[:2]) + factor = preparation.resolve_downsample( + pyramid_downsample, + image_shape, + ) + + coarse = _downsampled_correction(self, factor) + + # Pooling both scans preserves a mutual alignment objective. A fixed + # reference is more sensitive to pooling bias, so it refines natively. + reference_search = bool(fixed_set) + if reference_search or any(size % 2 for size in image_shape): + refine_factor = 1 + elif max(image_shape) <= 256: + refine_factor = 2 else: - pyramid = type(self).from_data( - pooled, - scan_direction_degrees=self.scan_direction_degrees, - device=self.device, + refine_factor = 4 + refine_object = ( + self if refine_factor == 1 else _downsampled_correction(self, refine_factor) + ) + if max_image_shift == "auto": + native_shift = max( + 16.0, + min(image_shape) * (0.0625 if reference_search else 0.25), ) - pyramid.preprocess( - pad_fraction=self.pad_fraction, - pad_value="median", - kde_sigma=self.kde_sigma, - number_knots=1, - show_merged=False, - show_images=False, - show_knots=False, - ).align_affine( - step=step, - num_tests=num_tests, - refine=refine, - upsample_factor=max(2, upsample_factor // 2), - max_image_shift=( - None - if max_image_shift is None - else max(2.0, float(max_image_shift) / factor) - ), - chunk_size=chunk_size, - downsample=1, - show_merged=False, - show_images=False, - show_knots=False, - verbose=False, + else: + native_shift = max_image_shift + coarse_shift = None if native_shift is None else max(2.0, float(native_shift) / factor) + refine_shift = ( + None if native_shift is None else max(2.0, float(native_shift) / refine_factor) ) - solved_index = 1 if getattr(self, "_reference_mode", False) else 0 - coarse_rate = _affine_rate_from_knots(pyramid, solved_index) - _apply_affine_rate(self, coarse_rate) - native_rate = coarse_rate.copy() + # ------------------------------------------------------------------------- + # Phase 1 of 4 — Find the broad drift-rate basin cheaply + # Search nested circular grids on the pooled scans. Stop expanding when the + # winner is safely inside the tested physical drift-rate radius. + # ------------------------------------------------------------------------- + center = np.zeros(2, dtype=np.float64) + best_center = center.copy() + best_cost = np.inf + coarse_step = 0.0 + evaluations = 0 + history = [] + coarse_cost_cache: dict[tuple[float, float], float] = {} + broad_candidate_requests = 0 + if reference_search: + radii = (0.10, 0.20, 0.40, 0.80, 1.60) + elif self.shape[0] == 2 and max(image_shape) <= 256: + # On small mutual scans, kernel-launch overhead dominates and the + # local 2x grids supply the fine resolution. Start broad; evaluate + # 0.25 only when the 0.20 winner approaches the boundary. Figure 5 + # stays interior and drops from 203 to 105 total evaluations. + radii = (0.20, 0.25) + else: + radii = (0.05, 0.10, 0.20, 0.25) + for radius in radii: + candidates = candidate_grid( + np.zeros(2), + radius, + 9, + circular=True, + ) + broad_candidate_requests += len(candidates) + # Absolute broad grids are nested: the 0.10 grid repeats 13 + # vectors from 0.05, and each doubled radius does the same. Cache + # only identical coordinates; stage ordering and costs stay exact. + candidate_keys = [tuple(np.round(candidate, 12)) for candidate in candidates] + new_indices = [ + index for index, key in enumerate(candidate_keys) if key not in coarse_cost_cache + ] + if new_indices: + _, new_costs = grid_search_batch( + coarse, + candidates[new_indices], + upsample_factor=max(2, upsample_factor // 2), + max_image_shift=coarse_shift, + chunk_size=chunk_size, + fixed_indices=fixed_set, + # A second top-1 score only guards the broad basin. Once the + # basin is chosen, the established objective owns refinement. + fixed_overlap_check=True, + ) + evaluations += len(new_indices) + for index, cost in zip( + new_indices, + new_costs.cpu().tolist(), + strict=True, + ): + coarse_cost_cache[candidate_keys[index]] = float(cost) + costs = np.asarray( + [coarse_cost_cache[key] for key in candidate_keys], + dtype=np.float64, + ) + best_index = int(np.argmin(costs)) + stage_cost = float(costs[best_index]) + stage_center = candidates[int(best_index)] + previous_best = best_cost + if stage_cost < best_cost: + best_cost = stage_cost + best_center = stage_center.copy() + coarse_step = 2.0 * radius / 8.0 + boundary_fraction = float(np.linalg.norm(stage_center) / radius) + history.append( + { + "stage": "pyramid", + "radius": float(radius), + "best_rate": stage_center.tolist(), + "cost": stage_cost, + } + ) + if not reference_search and boundary_fraction < 0.72: + break + if not reference_search and stage_cost > previous_best * 1.001: + break + + # ------------------------------------------------------------------------- + # Phase 2 of 4 — Refine the broad winner in native coordinates + # Repeatedly halve a local 3x3 grid. The images may remain pooled for speed, + # but the drift rate and radius always use native pixels per scanline. + # ------------------------------------------------------------------------- + center = best_center + radius = max(coarse_step * 1.5, 0.002) + bridge_center = None + native_levels = 8 if reference_search else 5 + final_costs = None + final_index = 0 + refinement_evaluations = 0 native_evaluations = 0 - if refine: - residual_step = step / max(num_tests - 1, 1) - axis = np.asarray((-residual_step, 0.0, residual_step)) - row, column = np.meshgrid(axis, axis, indexing="ij") - candidates = np.column_stack((row.ravel(), column.ravel())) - best_index, _ = self._affine_grid_search_batch( + level = 0 + while level < native_levels: + previous_center = center.copy() + candidates = candidate_grid( + center, + radius, + 3, + circular=False, + ) + final_index, final_costs = grid_search_batch( + refine_object, candidates, + upsample_factor=upsample_factor, + max_image_shift=refine_shift, + chunk_size=chunk_size, + fixed_indices=fixed_set, + ) + evaluations += len(candidates) + refinement_evaluations += len(candidates) + if refine_factor == 1: + native_evaluations += len(candidates) + center = candidates[int(final_index)] + if ( + level == 0 + and refine_factor == 4 + and np.max(np.abs(center - previous_center)) >= radius * 0.999 + ): + # A boundary winner on the first 4x local grid says the + # coarser pyramid changed the correlation basin. Retry at + # 2x automatically; periodic WS2 needs this, while the other + # real workflows retain the faster 4x refinement. + history.append( + { + "stage": "refine_4x_probe_rejected", + "radius": float(radius), + "best_rate": center.tolist(), + "cost": float(final_costs[int(final_index)].cpu()), + } + ) + refine_factor = 2 + refine_object = _downsampled_correction(self, refine_factor) + refine_shift = ( + None if native_shift is None else max(2.0, float(native_shift) / refine_factor) + ) + center = best_center.copy() + radius = max(coarse_step * 1.5, 0.002) + continue + history.append( + { + "stage": ( + f"native_{level + 1}" + if refine_factor == 1 + else f"refine_{refine_factor}x_{level + 1}" + ), + "radius": float(radius), + "best_rate": center.tolist(), + "cost": float(final_costs[int(final_index)].cpu()), + } + ) + if bridge_center is None and radius * image_shape[0] <= 40.0: + bridge_center = center.copy() + radius *= 0.5 + level += 1 + + # ------------------------------------------------------------------------- + # Phase 3 of 4 — Remove the small drift-rate bias from pooling + # Test one native-resolution sentinel grid before accepting the pyramid + # result, while retaining nearly all of the speed from downsampling. + # ------------------------------------------------------------------------- + if not reference_search and refine_factor > 1: + candidates = candidate_grid( + center, + radius * 2.0, + 3, + circular=True, + ) + final_index, final_costs = grid_search_batch( + self, + candidates, + upsample_factor=upsample_factor, + max_image_shift=native_shift, + chunk_size=chunk_size, + fixed_indices=fixed_set, + ) + evaluations += len(candidates) + native_evaluations += len(candidates) + center = candidates[int(final_index)] + history.append( + { + "stage": "native_sentinel", + "radius": float(radius * 2.0), + "best_rate": center.tolist(), + "cost": float(final_costs[int(final_index)].cpu()), + } + ) + + # ------------------------------------------------------------------------- + # Phase 4 of 4 — Verify the image that will actually be delivered + # Rank the final neighborhood after translation correction. This prevents + # periodic specimens from selecting an adjacent, translation-equivalent + # basin that scores well during the cheaper pyramid search. + # ------------------------------------------------------------------------- + validation_costs = None + validation_evaluations = 0 + validation_status = set() + translation_verification_shift = np.zeros(2, dtype=np.float64) + translation_verification_gain = 0.0 + if not reference_search: + # The batched forward-scatter score is excellent for locating the + # drift-rate basin, but the public result is judged *after* the + # translation solve. Polish that delivered objective with five + # geometry-scaled trials. This also avoids committing to a + # neighboring lattice-translation basin on periodic specimens. + validation_radius = 2.0 / image_shape[0] + cache = {} + + row_rates = np.stack( + [ + center + np.asarray((row_delta, 0.0)) + for row_delta in ( + -validation_radius, + 0.0, + validation_radius, + ) + ] + ) + row_results, status = _delivered_candidates( + self, + row_rates, + starting_knots, + fixed_set, + native_shift, upsample_factor, - max_image_shift, - chunk_size, - fixed_reference=getattr(self, "_reference_mode", False), ) - residual_rate = candidates[best_index] - _apply_affine_rate(self, residual_rate) - native_rate += residual_rate - native_evaluations = len(candidates) + validation_status.update(status) + row_trials = [] + for rate, result in zip(row_rates, row_results, strict=True): + cache[tuple(rate)] = result + row_trials.append((result[0], rate)) + _, row_center = min(row_trials, key=lambda item: item[0]) + col_trials = [] + new_col_rates = [] + for col_delta in ( + -validation_radius, + 0.0, + validation_radius, + ): + rate = row_center + np.asarray((0.0, col_delta)) + key = tuple(rate) + if key not in cache: + new_col_rates.append(rate) + if new_col_rates: + new_col_rates_array = np.stack(new_col_rates) + new_col_results, status = _delivered_candidates( + self, + new_col_rates_array, + starting_knots, + fixed_set, + native_shift, + upsample_factor, + ) + validation_status.update(status) + for rate, result in zip( + new_col_rates_array, + new_col_results, + strict=True, + ): + cache[tuple(rate)] = result + for col_delta in ( + -validation_radius, + 0.0, + validation_radius, + ): + rate = row_center + np.asarray((0.0, col_delta)) + key = tuple(rate) + col_trials.append((cache[key][0], rate)) + _, polished_center = min(col_trials, key=lambda item: item[0]) - warped_t = self._warp_and_translate_torch( - max_image_shift, - upsample_factor, - fixed_reference=getattr(self, "_reference_mode", False), + # One extra two-stage candidate protects periodic images where a + # translation-only solve before the rate correction establishes + # the physically correct correlation basin. + zero_bridge = np.zeros(2, dtype=np.float64) + zero_result = _delivered_candidate( + self, + polished_center, + starting_knots, + fixed_set, + native_shift, + upsample_factor, + zero_bridge, + ) + cache[("zero_bridge",)] = zero_result + best_key, best_result = min(cache.items(), key=lambda item: item[1][0]) + if best_key == ("zero_bridge",): + center = polished_center + bridge_center = zero_bridge + else: + center = np.asarray(best_key, dtype=np.float64) + bridge_center = center.copy() + if best_result[1] is None: + self.knots = [knot.clone() for knot in starting_knots] + _apply_affine_rate(self, center, fixed_set) + for image_index in range(self.shape[0]): + self.knots[image_index][0] += best_result[3][image_index, 0] + self.knots[image_index][1] += best_result[3][image_index, 1] + else: + self.knots = [knot.clone() for knot in best_result[1]] + warped = best_result[2] + residual_limit = min(float(native_shift or 64.0), 64.0) + _, residual, residual_gain = fixed_overlap_ncc( + warped[:1], + warped[1:], + image_shape, + residual_limit, + ) + if residual_gain[0] >= 0.01: + residual = residual[0] + translation_verification_shift = residual.cpu().numpy() + translation_verification_gain = float(residual_gain[0].cpu()) + pair_shifts = torch.stack((-residual / 2, residual / 2)) + for image_index in range(2): + self.knots[image_index][0] += pair_shifts[image_index, 0] + self.knots[image_index][1] += pair_shifts[image_index, 1] + warped = warp_and_translate( + self, + native_shift, + solve_translation=False, + ) + self.imgs_warped.array[:] = warped.cpu().numpy() + validation_costs = sorted(result[0] for result in cache.values()) + validation_evaluations = len(cache) + evaluations += validation_evaluations + history.append( + { + "stage": "delivered_objective", + "radius": float(validation_radius), + "best_rate": center.tolist(), + "translation_bridge": bridge_center.tolist(), + "cost": float(best_result[0]), + "evaluations": validation_evaluations, + } + ) + else: + if bridge_center is None: + bridge_center = center.copy() + _apply_affine_rate(self, bridge_center, fixed_set) + warped = warp_and_translate( + self, + native_shift, + upsample_factor, + fixed_indices=fixed_set, + ) + final_delta = center - bridge_center + if np.any(final_delta != 0): + _apply_affine_rate(self, final_delta, fixed_set) + warped = warp_and_translate( + self, + native_shift, + upsample_factor, + fixed_indices=fixed_set, + ) + report.record_error(self, 1, warped) + + fallback_reason = None + suspicious_reference_basin = ( + reference_search + and np.linalg.norm(center) > 0.25 + and ( + self.error_track[-1, 1] > initial_error_row[1] * 0.8 + or best_cost < float(history[0]["cost"]) * 0.75 + ) ) - self.calculate_error(1, _warped_t=warped_t) + if reference_search and ( + self.error_track[-1, 1] > initial_error_row[1] or suspicious_reference_basin + ): + # Backward-warp screening can still become border-dominated for + # extreme rates. When the delivered result is weak, worse than + # the input, or implausibly better in a wide-rate basin, switch + # to a slower multi-start pattern search that scores the actual + # post-translation output. Ordinary reference workflows + # (including XEDS) do not pay this fallback cost. + fallback_reason = ( + "screened solution entered a suspicious wide-rate basin" + if suspicious_reference_basin + else "screened solution worsened delivered error" + ) + attempted_center = center.copy() + best_cost = float(initial_error_row[1]) + best_rate = np.zeros(2, dtype=np.float64) + best_knots = [knot.clone() for knot in starting_knots] + best_warped = None + fallback_costs = [best_cost] + fallback_evaluations = 0 + + # Rank the winners from every broad-search radius with the + # delivered post-translation objective. A border-dominated + # backward-warp minimum at the widest radius must not erase a + # physically correct seed found one level earlier. + seed_records = [ + (np.zeros(2, dtype=np.float64), 4.0 / image_shape[0]), + (attempted_center, coarse_step * 1.5), + ] + seed_records.extend( + ( + np.asarray(item["best_rate"], dtype=np.float64), + 1.5 * (2.0 * float(item["radius"]) / 8.0), + ) + for item in history + if item["stage"] == "pyramid" + ) + unique_seeds = {} + for seed_rate, seed_radius in seed_records: + unique_seeds.setdefault( + tuple(seed_rate), + (seed_rate, seed_radius), + ) + seed_results = [] + for seed_rate, seed_radius in unique_seeds.values(): + result = _delivered_candidate( + self, + seed_rate, + starting_knots, + fixed_set, + native_shift, + upsample_factor, + ) + seed_results.append( + ( + result[0], + seed_rate.copy(), + seed_radius, + result[1], + result[2], + ) + ) + fallback_evaluations += len(seed_results) + fallback_costs.extend(result[0] for result in seed_results) + seed_best = min(seed_results, key=lambda result: result[0]) + if seed_best[0] < best_cost: + ( + best_cost, + best_rate, + _, + best_knots, + best_warped, + ) = seed_best + + # Refine the three best delivered seeds independently. The + # single-point seed ranking can still favor a nearby alias before + # either basin reaches its optimum; multi-start keeps the robust + # basin without returning to a dense global native grid. + selected_seeds = sorted( + seed_results, + key=lambda result: result[0], + )[:3] + for required_rate in (attempted_center, best_center): + required = next( + result for result in seed_results if np.array_equal(result[1], required_rate) + ) + if not any(np.array_equal(result[1], required[1]) for result in selected_seeds): + selected_seeds.append(required) + for seed_result in selected_seeds: + fallback_center = seed_result[1].copy() + fallback_radius = max( + float(seed_result[2]), + 4.0 / image_shape[0], + ) + while True: + candidates = candidate_grid( + fallback_center, + fallback_radius, + 3, + circular=False, + ) + local_results = [] + for rate in candidates: + result = _delivered_candidate( + self, + rate, + starting_knots, + fixed_set, + native_shift, + upsample_factor, + ) + local_results.append( + (result[0], rate.copy(), result[1], result[2]) + ) + fallback_evaluations += len(candidates) + local_best = min(local_results, key=lambda result: result[0]) + fallback_costs.extend(result[0] for result in local_results) + fallback_center = local_best[1].copy() + if local_best[0] < best_cost: + ( + best_cost, + best_rate, + best_knots, + best_warped, + ) = local_best + if fallback_radius * image_shape[0] <= 0.5: + break + fallback_radius *= 0.5 + + evaluations += fallback_evaluations + validation_evaluations += fallback_evaluations + validation_costs = sorted(fallback_costs) + center = best_rate.copy() + bridge_center = center.copy() + self.knots = [knot.clone() for knot in best_knots] + if best_warped is None: + warped = warp_and_translate( + self, + native_shift, + upsample_factor, + solve_translation=False, + fixed_indices=fixed_set, + ) + final_error_row = initial_error_row.copy() + final_error_row[0] = 1.0 + else: + warped = best_warped + self.imgs_warped.array[:] = warped.cpu().numpy() + per_image = ( + _scan_disagreement(warped, 0).mean(dim=(1, 2)) + .cpu() + .numpy() + ) + final_error_row = np.hstack((1.0, np.mean(per_image), per_image)) + self.error_track[-1] = final_error_row + history.append( + { + "stage": "delivered_reference_fallback", + "attempted_rate": attempted_center.tolist(), + "best_rate": center.tolist(), + "cost": float(self.error_track[-1, 1]), + "evaluations": fallback_evaluations, + } + ) + + if self.error_track[-1, 1] > initial_error_row[1]: + # Automatic correction may find no improvement, but it must never + # return a worse alignment than the untouched input. + fallback_reason = fallback_reason or "delivered error worsened" + self.knots = [knot.clone() for knot in starting_knots] + warped = warp_and_translate( + self, + native_shift, + upsample_factor, + solve_translation=False, + fixed_indices=fixed_set, + ) + final_error_row = initial_error_row.copy() + final_error_row[0] = 1.0 + self.error_track[-1] = final_error_row + center = np.zeros(2, dtype=np.float64) + bridge_center = center.copy() + history.append( + { + "stage": "safe_noop_fallback", + "best_rate": center.tolist(), + "cost": float(initial_error_row[1]), + } + ) + + if validation_costs is not None and len(validation_costs) > 1: + self.affine_confidence_margin = ( + (validation_costs[1] - validation_costs[0]) / (validation_costs[0] + 1e-12) * 100 + ) + elif final_costs is not None and len(final_costs) > 1: + ranked = torch.sort(final_costs).values + self.affine_confidence_margin = float( + ((ranked[1] - ranked[0]) / (ranked[0] + 1e-12) * 100).cpu() + ) + else: + self.affine_confidence_margin = 0.0 + elapsed = time.perf_counter() - start_time self.affine_search_info = { - "strategy": "average_pooled_pyramid", + "strategy": "automatic_pyramid", "downsample_factor": factor, - "coarse_rate_px_per_line": coarse_rate.tolist(), - "native_rate_px_per_line": native_rate.tolist(), + "refine_downsample_factor": refine_factor, + "candidate_evaluations": evaluations, + "broad_candidate_evaluations": len(coarse_cost_cache), + "broad_candidate_reuses": (broad_candidate_requests - len(coarse_cost_cache)), + "refinement_candidate_evaluations": refinement_evaluations, "native_candidate_evaluations": native_evaluations, + "delivered_objective_evaluations": validation_evaluations, + "delivered_objective_batched": "batch" in validation_status, + "delivered_objective_tie_fallback": "tie" in validation_status, + "delivered_objective_memory_fallback": "memory" in validation_status, + "drift_rate_row_col": center.tolist(), + "translation_bridge_row_col": bridge_center.tolist(), + "translation_verification_shift_row_col": ( + translation_verification_shift.tolist() + ), + "translation_verification_ncc_gain": translation_verification_gain, + "max_image_shift": native_shift, + "fallback_reason": fallback_reason, + "seconds": elapsed, + "preprocess_seconds": self._implicit_preprocess_seconds, + "total_seconds": elapsed + self._implicit_preprocess_seconds, + "preprocess": dict(self.preprocess_info), + "history": history, } - diagnostics._record_stage(self, "affine") if verbose: print( - "align_affine: average-pooled pyramid " - f"{factor}x; native rate=({native_rate[0]:+.6f}, " - f"{native_rate[1]:+.6f}) px/line; " - f"{native_evaluations} native candidates" - ) - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images( - show_knots=show_knots, - title="Merged: affine", - **kwargs, + "correct_affine: automatic pyramid " + f"{factor}x, {native_evaluations} native candidates, " + f"{elapsed:.2f} s" ) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {index}: affine" for index in range(self.shape[0])], - **kwargs, + print( + "Drift: " + f"({center[0]:+.5f}, {center[1]:+.5f}) px/line; " + f"translation bound {native_shift} px" ) + + # Fallback candidates are created inside this inference-mode method. + # Materialize ordinary tensors before handing knots to subsequent + # affine/strip/nonrigid calls, which legitimately update them in + # place outside inference mode. + with torch.inference_mode(False): + self.knots = [knot.detach().clone() for knot in self.knots] + drift_plot.show_after_step( + self, + "affine", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + with torch.inference_mode(False): + self._knots_after_affine = [knot.detach().clone() for knot in self.knots] + self._images_warped_stale = True + if show_report: + print(self.report().to_string()) return self @torch.inference_mode() -def _affine_grid_search_batch( +def grid_search_batch( self, drift_vectors, upsample_factor, max_image_shift, chunk_size=None, - fixed_reference=False, + fixed_indices=None, + progress_desc=None, + fixed_overlap_check=False, ): """Evaluate all candidate drift vectors in parallel. @@ -375,6 +1502,11 @@ def _affine_grid_search_batch( batching, each candidate would be a separate Python iteration - this is the key operation that enables the 300x speedup. + When ``fixed_indices`` is provided, images at those indices are + warped once with their current knots (no candidate drift) and + reused across all chunks. Only non-fixed images receive the + candidate drift offsets. + Parameters ---------- drift_vectors : ndarray, shape (N, 2) @@ -385,6 +1517,15 @@ def _affine_grid_search_batch( Maximum allowed shift for cross-correlation peak search. chunk_size : int or None Number of candidates per pass. If None, all at once. + fixed_indices : frozenset[int] or None + Indices of images whose knots should not receive candidate + drift. These images are warped once and reused. + progress_desc : str or None + Description for a progress bar shown only when candidate + evaluation requires multiple chunks. ``None`` disables it. + fixed_overlap_check : bool + Guard the automatic search's broad basin against a translation peak + selected mostly by padding. Refinement keeps the established cost. Returns ------- @@ -395,191 +1536,205 @@ def _affine_grid_search_batch( """ device = self._device dtype = self._dtype + fixed_set = fixed_indices if fixed_indices else frozenset() num_candidates = drift_vectors.shape[0] drift_vectors_t = torch.tensor(drift_vectors, dtype=dtype, device=device) + + # When fixed_indices is set, use backward-warp scoring to avoid + # KDE forward-scatter bias. The periodic wrapping in + # bilinear_kde_batch creates geometry-dependent seam artifacts + # that differ between the fixed reference and drift-shifted + # moving image, making the MAE minimum diverge from the true + # drift. Backward-warp scoring works at original resolution + # with grid_sample (no canvas, no KDE). + if fixed_set: + fixed_idx = sorted(fixed_set)[0] + moving_indices = [i for i in range(len(self.imgs_t)) if i not in fixed_set] + if not moving_indices: + raise ValueError("All images are fixed - nothing to optimize.") + total_costs = None + for mov_idx in moving_indices: + desc = progress_desc + if desc is not None and len(moving_indices) > 1: + desc = f"{desc} (scan {mov_idx})" + _, costs = backward_warp_grid_search( + self.imgs_t[fixed_idx], + self.imgs_t[mov_idx], + drift_vectors_t, + upsample_factor, + max_image_shift, + chunk_size, + progress_desc=desc, + ) + total_costs = costs if total_costs is None else total_costs + costs + return torch.argmin(total_costs).item(), total_costs + canvas_shape = (self.shape[1], self.shape[2]) + n_images = len(self.imgs_t) # Base coordinates shared across all candidates base_data = [] - for img_idx in range(2): - knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) - row_base, col_base = transform_coordinates_single_knot( - knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) - num_rows = self.knots[img_idx].shape[1] - scanline_offset = (torch.arange(num_rows, dtype=dtype, device=device) - - (num_rows - 1) / 2) - base_data.append((self.images_t[img_idx], row_base, col_base, scanline_offset)) + for img_idx in range(n_images): + row_base, col_base, scanline_offset = drift_knots.interpolator( + self, img_idx + ).affine_candidate_base() + base_data.append((self.imgs_t[img_idx], row_base, col_base, scanline_offset)) # Precompute shift mask and frequency grids (shared across chunks) shift_mask = None if max_image_shift is not None: canvas_rows, canvas_cols = canvas_shape freq_row = fftfreq(canvas_rows, 1.0 / canvas_rows, device=device, dtype=dtype) freq_col = fftfreq(canvas_cols, 1.0 / canvas_cols, device=device, dtype=dtype) - shift_mask = freq_row[:, None] ** 2 + freq_col[None, :] ** 2 >= max_image_shift ** 2 + shift_mask = freq_row[:, None] ** 2 + freq_col[None, :] ** 2 >= max_image_shift**2 freq_grids = ( fftfreq(canvas_shape[0], device=device, dtype=dtype)[:, None], fftfreq(canvas_shape[1], device=device, dtype=dtype)[None, :], ) + device_type = torch.device(device).type if chunk_size is None: - chunk_size = self._auto_chunk_size(num_candidates, canvas_shape, dtype, device) - on_cuda = torch.device(device).type == "cuda" - chunked = on_cuda and chunk_size < num_candidates + chunk_size = automatic_chunk_size(num_candidates, canvas_shape, dtype, device) + chunked = chunk_size < num_candidates all_costs = [] + all_overlap_costs = [] chunk_start = 0 chunk_idx = 0 - while chunk_start < num_candidates: - chunk_end = min(chunk_start + chunk_size, num_candidates) - drift_chunk = drift_vectors_t[chunk_start:chunk_end] - if chunk_idx == 0 and chunked: - torch.cuda.reset_peak_memory_stats(device) - warped_pair = [] - for img_idx in range(2): - image_t, row_base, col_base, scanline_offset = base_data[img_idx] - if fixed_reference and img_idx == 0: - batch_size = drift_chunk.shape[0] - row_candidates = row_base[None].expand(batch_size, -1, -1) - col_candidates = col_base[None].expand(batch_size, -1, -1) - else: - row_candidates = row_base[None] + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] - col_candidates = col_base[None] + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] - warped, _ = bilinear_kde_batch( - row_candidates, col_candidates, image_t, - canvas_shape, self.kde_sigma, - self.pad_value[img_idx]) - warped_pair.append(warped) - all_costs.append(cross_corr_batch( - warped_pair[0], warped_pair[1], - upsample_factor, - max_shift_mask=shift_mask, - freq_grids=freq_grids)) - # After chunk 0, replace the conservative static estimate with the - # actual measured per-candidate cost and print one summary line so - # the user can see how the chunking adapted to their GPU state. - if chunk_idx == 0 and chunked: - per_candidate_actual = torch.cuda.max_memory_allocated(device) / chunk_size - free_bytes, total_bytes = torch.cuda.mem_get_info(device) - tuned_chunk_size = max(1, int(free_bytes * 0.5 / per_candidate_actual)) - tuned_chunk_size = min(tuned_chunk_size, num_candidates) - if tuned_chunk_size > chunk_size: - chunk_size = tuned_chunk_size - num_chunks_final = 1 + (num_candidates - chunk_end + chunk_size - 1) // chunk_size - print( - f" affine grid: {num_candidates} cand × {canvas_shape[0]}×{canvas_shape[1]}, " - f"{per_candidate_actual / 1e9:.2f} GB/cand → {chunk_size}/chunk × {num_chunks_final} passes " - f"({free_bytes / 1e9:.0f}/{total_bytes / 1e9:.0f} GB free)" - ) - chunk_start = chunk_end - chunk_idx += 1 + pbar = tqdm( + total=num_candidates, + desc=progress_desc, + unit="candidate", + disable=progress_desc is None or not chunked, + ) + try: + while chunk_start < num_candidates: + chunk_end = min(chunk_start + chunk_size, num_candidates) + drift_chunk = drift_vectors_t[chunk_start:chunk_end] + if chunk_idx == 0 and chunked and device_type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + # Warp each image (fixed → expand once, moving → drift-shifted) + warped_images = [] + for img_idx in range(n_images): + image_t, row_base, col_base, scanline_offset = base_data[img_idx] + row_candidates = ( + row_base[None] + + drift_chunk[:, 0, None, None] * scanline_offset[None, :, None] + ) + col_candidates = ( + col_base[None] + + drift_chunk[:, 1, None, None] * scanline_offset[None, :, None] + ) + warped, _ = drift_knots.bilinear_kde_batch( + row_candidates, + col_candidates, + image_t, + canvas_shape, + self.kde_sigma, + self.pad_value[img_idx], + ) + warped_images.append(warped) + # Score all unique pairs and sum costs + chunk_cost = torch.zeros(chunk_end - chunk_start, dtype=dtype, device=device) + overlap_cost = torch.zeros_like(chunk_cost) + for i in range(n_images): + for j in range(i + 1, n_images): + chunk_cost += cross_corr_batch( + warped_images[i], + warped_images[j], + upsample_factor, + max_shift_mask=shift_mask, + freq_grids=freq_grids, + ) + if fixed_overlap_check: + pair_cost, _, _ = fixed_overlap_ncc( + warped_images[i], + warped_images[j], + tuple(int(value) for value in self.imgs[i].shape[:2]), + max_image_shift, + ) + overlap_cost += pair_cost + all_costs.append(chunk_cost) + if fixed_overlap_check: + all_overlap_costs.append(overlap_cost) + # After chunk 0, replace the conservative static estimate with the + # actual measured per-candidate cost and print one summary line so + # the user can see how the chunking adapted to their GPU state. + if chunk_idx == 0 and chunked and device_type in {"cuda", "mps"}: + if device_type == "cuda": + per_candidate_actual = torch.cuda.max_memory_allocated(device) / chunk_size + free_bytes, total_bytes = torch.cuda.mem_get_info(device) + tuned_chunk_size = max(1, int(free_bytes * 0.5 / per_candidate_actual)) + tuned_chunk_size = min(tuned_chunk_size, num_candidates) + if tuned_chunk_size > chunk_size: + chunk_size = tuned_chunk_size + memory_text = f"{free_bytes / 1e9:.0f}/{total_bytes / 1e9:.0f} GB free" + else: + total_bytes = torch.mps.recommended_max_memory() + live_bytes = max( + torch.mps.current_allocated_memory(), + torch.mps.driver_allocated_memory(), + ) + free_bytes = max(total_bytes - live_bytes, 0) + memory_text = ( + f"{free_bytes / 1e9:.0f}/{total_bytes / 1e9:.0f} GB MPS headroom" + ) + num_chunks_final = ( + 1 + (num_candidates - chunk_end + chunk_size - 1) // chunk_size + ) + if progress_desc is not None: + pbar.write( + f" affine grid: {num_candidates} drift vectors × " + f"{canvas_shape[0]}×{canvas_shape[1]}, " + f"auto chunk {chunk_size}/chunk × " + f"{num_chunks_final} passes ({memory_text})" + ) + pbar.update(chunk_end - chunk_start) + chunk_start = chunk_end + chunk_idx += 1 + finally: + pbar.close() all_costs = torch.cat(all_costs) - return torch.argmin(all_costs).item(), all_costs + if not fixed_overlap_check: + return torch.argmin(all_costs).item(), all_costs + overlap_costs = torch.cat(all_overlap_costs) + legacy_index = torch.argmin(all_costs) + overlap_index = torch.argmin(overlap_costs) + overlap_gain = float(overlap_costs[legacy_index] - overlap_costs[overlap_index]) + if overlap_gain >= 0.20: + return overlap_index.item(), overlap_costs + return legacy_index.item(), all_costs -def _auto_chunk_size(num_candidates, canvas_shape, dtype, device): +def automatic_chunk_size(num_candidates, canvas_shape, dtype, device): """Pick a candidate-batch size that fits in current free GPU memory. Empirical per-candidate peak (measured at 4096×4096): bilinear KDE scatter buffers, gaussian smoothing temporaries, then cross-correlation FFT pairs (complex64) - together about ``32 × canvas_pixels`` - ``× dtype_bytes`` at peak. We sample free memory at call time, divide - by that estimate with a 0.4 safety factor, and cap the result at - ``num_candidates`` (no point splitting if it all fits). - On CPU we just process all candidates at once - no VRAM constraint. + ``× dtype_bytes`` at peak. CUDA uses a 0.4 safety factor and then + retunes after the first chunk from measured peak memory. + + MPS (Apple unified memory) needs a more conservative static factor: + Metal shares memory with the OS and the bilinear KDE + smoothing + kernels can become memory-pressure limited well before the reported + recommended maximum. Use the larger of current and driver allocation + as live memory, then spend only a small fraction of the remaining + headroom. On CPU we process all candidates at once - no separate + device pool to overflow. """ device = torch.device(device) - if device.type != "cuda": - return num_candidates bytes_per_element = torch.finfo(dtype).bits // 8 per_candidate_bytes = canvas_shape[0] * canvas_shape[1] * bytes_per_element * 32 - free_bytes, _ = torch.cuda.mem_get_info(device) - chunk_size = max(1, int(free_bytes * 0.4 / per_candidate_bytes)) - return min(chunk_size, num_candidates) - - -@torch.inference_mode() -def _warp_and_translate_torch( - self, - max_image_shift: float | None, - upsample_factor: int = 8, - knots_batch: torch.Tensor | None = None, - solve_translation: bool = True, - fixed_reference: bool = False, -) -> torch.Tensor: - """Regenerate warped images and solve translation on GPU. - - Three phases: warp → solve translation → re-warp. When ``knots_batch`` - is provided, reads/writes a single batched torch tensor (zero numpy - crossings). Without it, falls back to ``self.knots`` (numpy) for - compatibility with ``align_affine``. - - Set ``solve_translation=False`` to only warp and sync without - re-solving translation - used after the nonrigid loop to populate - ``self.images_warped`` from final knots. - - Parameters - ---------- - max_image_shift : float or None - Maximum allowed translational shift in pixels. - upsample_factor : int - Sub-pixel precision for cross-correlation (1/N pixel). - knots_batch : torch.Tensor or None - If provided, batched ``(N, 2, num_rows)`` torch tensor on GPU. - Translation shifts are applied in-place. Skips numpy sync. - solve_translation : bool - If False, skip translation alignment (Phase 2+3). Only warp - once using current knots and sync to CPU. - fixed_reference : bool, default False - Keep image 0 fixed while translating later images. - - Returns - ------- - torch.Tensor - Warped images on GPU, shape ``(num_images, H, W)``. - """ - device = self._device - dtype = self._dtype - num_images = self.shape[0] - canvas_shape = (self.shape[1], self.shape[2]) - - def _warp_all(warped_t, weights_t): - """Warp all images onto the canvas using current knots.""" - for img_idx in range(num_images): - if knots_batch is not None: - # transform_coordinates_single_knot expects (2, N, 1) - knots_img = knots_batch[img_idx].detach()[:, :, None] - else: - knots_img = torch.as_tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_img, self.scan_fast_t[img_idx], self.images[img_idx].shape) - warped, weights = bilinear_kde_batch( - row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, - self.kde_sigma, self.pad_value[img_idx]) - warped_t[img_idx] = warped[0] - weights_t[img_idx] = weights[0] - - warped_t = torch.zeros(num_images, *canvas_shape, dtype=dtype, device=device) - weights_t = torch.zeros_like(warped_t) - _warp_all(warped_t, weights_t) - if not solve_translation: - self.images_warped.array[:] = warped_t.cpu().numpy() - self.weights_warped.array[:] = weights_t.cpu().numpy() - return warped_t - # Solve translation shifts and apply to knots - shifts_t = translate_align( - warped_t, - upsample_factor, - max_image_shift, - anchor_first=fixed_reference, - ) - if knots_batch is not None: - knots_batch[:, 0] += shifts_t[:, 0:1] - knots_batch[:, 1] += shifts_t[:, 1:2] + if device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(device) + safety_factor = 0.4 + elif device.type == "mps": + # recommended_max is Metal's working-set ceiling; subtract the + # larger live allocation estimate to get practical headroom. + live_bytes = max( + torch.mps.current_allocated_memory(), + torch.mps.driver_allocated_memory(), + ) + free_bytes = max(torch.mps.recommended_max_memory() - live_bytes, 0) + safety_factor = 0.075 else: - shifts_np = shifts_t.cpu().numpy() - for img_idx in range(num_images): - self.knots[img_idx][0] += shifts_np[img_idx, 0] - self.knots[img_idx][1] += shifts_np[img_idx, 1] - # Re-warp with corrected knots - _warp_all(warped_t, weights_t) - if knots_batch is None: - self.images_warped.array[:] = warped_t.cpu().numpy() - self.weights_warped.array[:] = weights_t.cpu().numpy() - return warped_t + return num_candidates + chunk_size = max(1, int(free_bytes * safety_factor / per_candidate_bytes)) + return min(chunk_size, num_candidates) diff --git a/src/quantem/imaging/drift/core/knots.py b/src/quantem/imaging/drift/core/knots.py index 4e1ddb52..3c39e592 100644 --- a/src/quantem/imaging/drift/core/knots.py +++ b/src/quantem/imaging/drift/core/knots.py @@ -1,122 +1,211 @@ -"""Knot interpolation and batched scan-coordinate projection.""" - +"""Scanline-knot geometry and forward warping. + +:class:`DriftKnot` keeps the single- and multi-knot coordinate models consistent. +It provides per-pixel canvas coordinates (:meth:`to_canvas`), raw-frame drift +(:meth:`drift_raw`), in-place affine slope (:meth:`apply_affine_shift`), +and forward scatter onto the canvas (:meth:`warp_to_canvas`). + +The forward-warp kernels live here too because forward scatter is what +the knot *does* to a source image: + +* :func:`bilinear_kde_batch` — batched bilinear forward scatter. +* :func:`gaussian_smooth_batch` / :func:`gaussian_smooth_1d` — separable + Gaussian smoothing reused by KDE normalization and (1-D) by the + knot regularizer in :mod:`nonrigid`. +* :func:`initialize_scanline_knots` — initial knot grid for ``preprocess``. +* Private :func:`_transform_coordinates_single_knot` / + :func:`_transform_coordinates_multi_knot` — the K-specific coordinate + formulas the class dispatches between. + +Backward warping, cross-correlation, and translation alignment are a +separate concerns and live in :mod:`warping`. +""" import numpy as np import torch from numpy.typing import NDArray -from scipy.interpolate import interp1d -from quantem.core.utils.imaging_utils import bilinear_kde +def interpolator(correction, image_index: int, knots: torch.Tensor | None = None): + """Map one scan's knot positions into its padded correction canvas. -class DriftInterpolator: - def __init__( - self, - input_shape, - output_shape, - scan_fast, - scan_slow, - pad_value, - kde_sigma, - ): - self.input_shape = input_shape - self.output_shape = output_shape - self.scan_fast = scan_fast - self.scan_slow = scan_slow - self.pad_value = pad_value - self.kde_sigma = kde_sigma - - self.rows_input = np.arange(input_shape[0]) - self.cols_input = np.arange(input_shape[1]) - self.u = np.linspace(0, 1, input_shape[1]) - - def transform_rows( - self, - knots_row: NDArray, - ): - num_knots = knots_row.shape[-1] - basis = np.linspace(0, 1, num_knots) - - if num_knots == 1: - xa = knots_row[0] + self.u[None, :] * self.scan_fast[0] * (self.input_shape[0] - 1) - ya = knots_row[1] + self.u[None, :] * self.scan_fast[1] * (self.input_shape[1] - 1) - elif num_knots == 2: - xa = interp1d(basis, knots_row[0], kind="linear", assume_sorted=True)(self.u) - ya = interp1d(basis, knots_row[1], kind="linear", assume_sorted=True)(self.u) - else: - kind = "quadratic" if num_knots == 3 else "cubic" - xa = interp1d( - basis, - knots_row[0], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - ya = interp1d( - basis, - knots_row[1], - kind=kind, - fill_value="extrapolate", - assume_sorted=True, - )(self.u) - - return xa, ya - - def transform_coordinates( - self, - knots: NDArray, - ): - num_knots = knots.shape[-1] - - if num_knots == 1: - # vectorized version for speed - xa, ya = self.transform_rows(knots) - else: - xa = np.zeros(self.input_shape) - ya = np.zeros(self.input_shape) - for i in range(self.input_shape[0]): - xa[i], ya[i] = self.transform_rows(knots[:, i]) + The knot geometry is shared by affine fitting, non-rigid warping, corrected + products, and probe-position recovery. Keeping that geometry here gives + every stage the same row/column coordinate transformation. + """ + if knots is None: + knots = correction.knots[image_index] + return DriftKnot( + knots, + correction.scan_fast_t[image_index], + correction.scan_slow_t[image_index], + correction.imgs[image_index].shape, + ) - return xa, ya - def warp_image( - self, - image: NDArray, - knots: NDArray, # shape: (2, rows, num_knots) - kde_sigma=None, - output_shape=None, - pad_value=None, - upsample_factor=None, - ) -> NDArray: - xa, ya = self.transform_coordinates( - knots, +def knot_delta_canvas(correction, image_index: int) -> torch.Tensor: + """Return one scan's fitted knot displacement on the correction canvas.""" + if not hasattr(correction, "_initial_knots"): + raise RuntimeError( + "apply_correction() requires preprocess() and correct_affine() " + "first. Run dc.preprocess().correct_affine() (and optionally " + ".correct_nonrigid()) before apply_correction()." ) + return correction.knots[image_index] - correction._initial_knots[image_index] + + +def stage_knots(correction, stage: str | None) -> list[torch.Tensor]: + """Select the saved knot field for a correction stage comparison.""" + if stage in (None, "nonrigid", "non-rigid"): + return correction.knots + attribute = { + "initial": "_initial_knots", + "raw": "_initial_knots", + "affine": "_knots_after_affine", + "strip": "_knots_after_strip", + }.get(stage) + if attribute is None: + raise ValueError( + f"Unknown correction stage {stage!r}. Choose initial, affine, " + "strip, nonrigid, or None for the current result." + ) + knots = getattr(correction, attribute, None) + if knots is None: + required = "preprocess" if stage in ("initial", "raw") else f"correct_{stage}" + raise ValueError(f"stage={stage!r} needs a prior {required} call.") + return knots - if kde_sigma is None: - kde_sigma = self.kde_sigma - - if output_shape is None: - output_shape = self.output_shape - if pad_value is None: - pad_value = self.pad_value +def initialize_scanline_knots( + input_shape: tuple[int, int], + output_shape: tuple[int, int], + scan_fast: NDArray, + scan_slow: NDArray, + number_knots: int, +) -> NDArray: + """Build the initial knot grid used by ``DriftCorrection.preprocess``. + + The knot anchors define where each scanline starts on the padded canvas + before any affine or non-rigid optimization. For ``number_knots == 1``, + this is a vertical line of anchors at the fast-scan start edge of the + centered footprint. The full scanline width is then added later by + ``_transform_coordinates_single_knot``. - if upsample_factor is None: - upsample_factor = 1.0 + Parameters + ---------- + input_shape : tuple[int, int] + Raw image shape ``(num_rows, num_cols)``. + output_shape : tuple[int, int] + Padded canvas shape ``(num_rows, num_cols)``. + scan_fast : NDArray + Unit vector of the fast scan direction in ``(row, col)`` order. + scan_slow : NDArray + Unit vector of the slow scan direction in ``(row, col)`` order. + number_knots : int + Number of control knots per scanline. - image_interp, weight_interp = bilinear_kde( - xa=xa * upsample_factor, # rows - ya=ya * upsample_factor, # cols - values=image, - output_shape=np.round(np.array(output_shape) * upsample_factor).astype("int"), - kde_sigma=kde_sigma * upsample_factor, - pad_value=pad_value, - return_pix_count=True, - ) + Returns + ------- + NDArray + Initial knot array with shape ``(2, input_rows, number_knots)``. + """ + v_slow = np.linspace(-(input_shape[0] - 1) / 2, (input_shape[0] - 1) / 2, input_shape[0]) + u_fast = np.linspace(-(input_shape[1] - 1) / 2, (input_shape[1] - 1) / 2, number_knots) + row_knots = ((output_shape[0] - 1) / 2 + + u_fast[None, :] * scan_fast[0] + + v_slow[:, None] * scan_slow[0]) + col_knots = ((output_shape[1] - 1) / 2 + + u_fast[None, :] * scan_fast[1] + + v_slow[:, None] * scan_slow[1]) + return np.stack([row_knots, col_knots], axis=0) + + +def resize_scanline_knots(correction, num_knots: int): + """Change fast-scan knot density without changing the fitted drift field + + Affine and strip correction establish a displacement field before a + scientist decides how much fast-scan flexibility the non-rigid stage + needs. Resampling the displacement at a new knot density lets + ``correct_nonrigid(num_knots=...)`` retain that corrected geometry instead + of repeating affine correction or discarding its result. - return image_interp, weight_interp + Parameters + ---------- + correction : DriftCorrection + Prepared correction containing the current and initial knot fields. + num_knots : int + New number of knots along every fast-scan line. + Returns + ------- + DriftCorrection + The same correction with every saved checkpoint represented at the + requested knot density. + """ + count = int(num_knots) + if count < 1: + raise ValueError(f"num_knots must be >= 1, got {num_knots!r}.") + current = {int(value.shape[2]) for value in correction.knots} + if current == {count}: + return correction + if len(current) != 1: + raise ValueError( + "All scans must use the same knot count before resizing; " + f"got {sorted(current)}." + ) -def transform_coordinates_single_knot( + old_initial = correction._initial_knots + new_initial = [ + torch.as_tensor( + initialize_scanline_knots( + input_shape=correction.imgs[index].shape, + output_shape=correction.shape[1:], + scan_fast=correction.scan_fast[index], + scan_slow=correction.scan_slow[index], + number_knots=count, + ), + dtype=correction._dtype, + device=correction._device, + ) + for index in range(correction.shape[0]) + ] + + def resize_checkpoint(checkpoint): + resized = [] + for value, initial, target in zip( + checkpoint, + old_initial, + new_initial, + strict=True, + ): + displacement = value - initial + if displacement.shape[2] == 1: + displacement = displacement.expand(-1, -1, count) + else: + rows = displacement.shape[1] + displacement = torch.nn.functional.interpolate( + displacement.reshape(1, 2 * rows, -1), + size=count, + mode="linear", + align_corners=True, + ).reshape(2, rows, count) + resized.append(target + displacement) + return resized + + checkpoints = { + name: resize_checkpoint(getattr(correction, name)) + for name in ("knots", "_knots_after_affine", "_knots_after_strip") + if hasattr(correction, name) + } + correction._initial_knots = new_initial + for name, values in checkpoints.items(): + setattr(correction, name, values) + correction.number_knots = count + correction.preprocess_info["num_knots"] = count + correction._images_warped_stale = True + return correction + + +def _transform_coordinates_single_knot( knots: torch.Tensor, scan_fast: torch.Tensor, input_shape: tuple[int, int], @@ -125,14 +214,9 @@ def transform_coordinates_single_knot( **Single-knot only.** Each scanline has exactly one (row, col) anchor; the fast-scan-direction position is filled in by linear interpolation - along the scanline. Multi-knot Bezier interpolation is intentionally - not supported here - that's the scipy backend's job. The pytorch path - optimizes for the common single-knot case (≥95% of real STEM workflows). - - Called by ``preprocess``, ``_affine_grid_search_batch``, and - ``_warp_and_translate_torch`` to map source image pixels onto the - padded output canvas. Without this, the warped images would have - no spatial mapping and the grid search couldn't score test drifts. + along the scanline. Multi-knot input is handled by + :func:`_transform_coordinates_multi_knot`; :class:`DriftKnot` dispatches + on ``knots.shape[2]``. Each input row maps to a line on the canvas: ``row = knot_row + fraction * scan_fast[0] * (num_rows - 1)`` @@ -143,8 +227,6 @@ def transform_coordinates_single_knot( ---------- knots : torch.Tensor Knot positions, shape ``(2, num_rows, 1)``. First dim is (row, col). - The trailing 1 is the single-knot dimension; multi-knot inputs are - rejected by the caller before reaching this function. scan_fast : torch.Tensor Fast scan direction vector, shape ``(2,)``. input_shape : tuple[int, int] @@ -156,14 +238,6 @@ def transform_coordinates_single_knot( Row coordinates on canvas, shape ``(num_rows, num_cols)``. col_coords : torch.Tensor Column coordinates on canvas, shape ``(num_rows, num_cols)``. - - Examples - -------- - >>> knots = torch.zeros(2, 64, 1) - >>> scan_fast = torch.tensor([0.0, 1.0]) - >>> r, c = transform_coordinates_single_knot(knots, scan_fast, (64, 64)) - >>> r.shape - torch.Size([64, 64]) """ num_rows, num_cols = input_shape fast_fraction = torch.linspace(0, 1, num_cols, dtype=knots.dtype, device=knots.device) @@ -172,6 +246,214 @@ def transform_coordinates_single_knot( return row_coords, col_coords +def _transform_coordinates_multi_knot( + knots: torch.Tensor, + input_shape: tuple[int, int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Multi-knot path: linearly interpolate K knot anchors per scanline. + + For ``K`` knots per scanline (``K >= 2``), knot ``k`` sits at fast-axis + fraction ``k / (K - 1)`` of the scanline. Per-pixel canvas coordinates + are obtained by piecewise-linear interpolation between the two adjacent + knots. The K=1 case is intentionally not handled here — it has no + end-knot to interpolate toward, so the caller dispatches to + :func:`_transform_coordinates_single_knot` (which walks along + ``scan_fast`` instead). + + With ``K = 2`` and the initial knot grid (knots at the start and end + of each scanline along ``scan_fast``), this reduces to the same + per-pixel canvas positions as the K=1 path — verified in tests. + + Parameters + ---------- + knots : torch.Tensor + Knot positions, shape ``(2, num_rows, K)`` with ``K >= 2``. + First axis is ``(row, col)``. + input_shape : tuple[int, int] + Original image shape ``(num_rows, num_cols)``. Only ``num_cols`` + is consulted (``num_rows`` is implicit in ``knots.shape[1]``). + + Returns + ------- + row_coords : torch.Tensor + Row coordinates on canvas, shape ``(num_rows, num_cols)``. + col_coords : torch.Tensor + Column coordinates on canvas, shape ``(num_rows, num_cols)``. + """ + _, num_cols = input_shape + K = knots.shape[2] + t = torch.linspace(0, 1, num_cols, dtype=knots.dtype, device=knots.device) * (K - 1) + seg = torch.clamp(t.long(), max=K - 2) + frac = (t - seg.to(knots.dtype))[None, :] + row_lo = knots[0, :, seg] + row_hi = knots[0, :, seg + 1] + col_lo = knots[1, :, seg] + col_hi = knots[1, :, seg + 1] + row_coords = row_lo + (row_hi - row_lo) * frac + col_coords = col_lo + (col_hi - col_lo) * frac + return row_coords, col_coords + + +class DriftKnot: + """Maps K knot anchors per scanline to canvas geometry, drift, and warps. + + This is the single dispatch point for K=1 versus K>=2 geometry. + K=1 walks along ``scan_fast`` (one anchor per row, scanline geometry + implicit); K>=2 piecewise-linearly interpolates the K anchors along + the fast axis. :func:`interpolator` builds the geometry for each scan. + + Attributes + ---------- + knots : torch.Tensor + Knot anchors, shape ``(2, H, K)`` (row/col × scanline × knot). + scan_fast, scan_slow : torch.Tensor + Fast / slow scan unit vectors ``(2,)``. ``scan_fast`` is only + consulted when ``K == 1`` (walk); ``scan_slow`` enters + :meth:`drift_raw` to invert the canvas Jacobian. + input_shape : tuple[int, int] + Source image shape ``(H, W)``. + K : int + Number of knots per scanline (cached from ``knots.shape[2]``). + """ + + def __init__( + self, + knots: torch.Tensor, + scan_fast: torch.Tensor, + scan_slow: torch.Tensor, + input_shape: tuple[int, int], + ): + self.knots = knots + self.scan_fast = scan_fast + self.scan_slow = scan_slow + self.input_shape = input_shape + self.K = knots.shape[2] + + def to_canvas(self) -> tuple[torch.Tensor, torch.Tensor]: + """Per-pixel canvas (row, col) coordinates for warping. + + Returns ``(row_coords, col_coords)`` each ``(H, W)``. K=1 uses the + ``scan_fast`` walk fast path; K>=2 uses the multi-knot lerp. + """ + if self.K == 1: + return _transform_coordinates_single_knot( + self.knots, self.scan_fast, self.input_shape) + return _transform_coordinates_multi_knot(self.knots, self.input_shape) + + def _drift_canvas(self, initial_knots: torch.Tensor) -> torch.Tensor: + """Drift in canvas coordinates relative to ``initial_knots``. + + Internal step used by :meth:`drift_raw`. Returns ``(2, H)`` for + K=1 (per-row shift) or ``(2, H, W)`` for K>=2 (per-pixel via lerp). + Callers want raw-frame drift, not canvas-frame, so this method + stays private. + """ + delta = self.knots - initial_knots + if self.K == 1: + return delta[:, :, 0] + _, num_cols = self.input_shape + t = torch.linspace(0, 1, num_cols, dtype=delta.dtype, device=delta.device) * (self.K - 1) + seg = torch.clamp(t.long(), max=self.K - 2) + frac = (t - seg.to(delta.dtype))[None, :] + delta_row = delta[0, :, seg] + (delta[0, :, seg + 1] - delta[0, :, seg]) * frac + delta_col = delta[1, :, seg] + (delta[1, :, seg + 1] - delta[1, :, seg]) * frac + return torch.stack([delta_row, delta_col]) + + def drift_raw(self, initial_knots: torch.Tensor) -> torch.Tensor: + """Drift relative to ``initial_knots`` in raw-frame coordinates. + + Returns ``(2, H)`` for K=1 (per-row shift) or ``(2, H, W)`` for K>=2 + (per-pixel). Inverts the canvas Jacobian so callers feed the + result straight into ``backward_warp``. For square images this + reduces to a rotation by the scan angle; the ``alpha`` factor + handles non-square scans. + """ + delta_canvas = self._drift_canvas(initial_knots) + scan_h, scan_w = self.input_shape + alpha = float(scan_h - 1) / float(scan_w - 1) if scan_w > 1 else 1.0 + det = (self.scan_slow[0] * self.scan_fast[1] + - self.scan_fast[0] * alpha * self.scan_slow[1]) + drift_row = ( + self.scan_fast[1] * delta_canvas[0] + - self.scan_fast[0] * alpha * delta_canvas[1] + ) / det + drift_col = ( + -self.scan_slow[1] * delta_canvas[0] + + self.scan_slow[0] * delta_canvas[1] + ) / det + return torch.stack([drift_row, drift_col]) + + def affine_candidate_base(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-pixel canvas coords plus the scanline-centered offset axis. + + Returns ``(row_base, col_base, scanline_offset)``. The first two + are :meth:`to_canvas`'s output; the third is + ``arange(H) - (H-1)/2`` shaped ``(H,)`` so callers can broadcast a + candidate ``drift_vec`` over scanlines as + ``row_base + drift_vec[0] * scanline_offset[:, None]``. + + Used by ``affine.grid_search_batch`` so candidate-broadcast geometry + construction lives on the geometry class instead of the orchestrator. + """ + row_base, col_base = self.to_canvas() + H = self.knots.shape[1] + scanline_offset = ( + torch.arange(H, dtype=self.knots.dtype, device=self.knots.device) + - (H - 1) / 2 + ) + return row_base, col_base, scanline_offset + + def apply_affine_shift(self, drift_vec: torch.Tensor) -> None: + """Add an affine drift slope to every knot, in place. + + For each scanline ``i``, shifts the knots by + ``drift_vec * (i - (H - 1) / 2)`` so the slow-axis-centered slope + accumulates linearly across the image. Used by ``correct_affine`` + to bake a per-row drift candidate into the knot grid. + """ + H = self.knots.shape[1] + scanline_offset = ( + torch.arange(H, dtype=self.knots.dtype, device=self.knots.device) + - (H - 1) / 2 + )[:, None] + self.knots[0] += drift_vec[0] * scanline_offset + self.knots[1] += drift_vec[1] * scanline_offset + + def warp_to_canvas( + self, + source_image: torch.Tensor, + canvas_shape: tuple[int, int], + kde_sigma: float, + pad_value, + upsample_factor: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Forward-scatter ``source_image`` onto the canvas using the knot grid. + + Wraps :meth:`to_canvas` + :func:`bilinear_kde_batch` so callers + don't repeat the (knots → coords → scatter) idiom. Returns + ``(warped, weights)``, each ``canvas_shape``. + + ``upsample_factor`` scales the canvas coordinates and KDE sigma in + lockstep, so ``corrected`` can scatter at a finer grid + without recomputing the interpolation. + """ + row_t, col_t = self.to_canvas() + if upsample_factor != 1: + row_t = row_t * upsample_factor + col_t = col_t * upsample_factor + warped, weights = bilinear_kde_batch( + row_t[None], col_t[None], source_image, canvas_shape, + kde_sigma, pad_value) + return warped[0], weights[0] + + +# --------------------------------------------------------------------------- +# Forward-warp kernels: scatter source pixels onto the canvas, smooth, normalize. +# Used by :meth:`DriftKnot.warp_to_canvas` (and ``affine.grid_search_batch``'s +# candidate broadcast which calls bilinear_kde_batch directly). +# --------------------------------------------------------------------------- + + def bilinear_kde_batch( row_coords: torch.Tensor, col_coords: torch.Tensor, @@ -276,8 +558,6 @@ def bilinear_kde_batch( sum_values / torch.clamp(sum_weights, min=1e-8) ) return warped_images, sum_weights - - def gaussian_smooth_batch( field_stack: torch.Tensor, sigma: float, @@ -336,14 +616,16 @@ def gaussian_smooth_1d( kernel, radius = _gaussian_kernel_1d(sigma, signal.dtype, signal.device) signal_padded = _symmetric_pad_1d(signal[:, None], radius) return torch.nn.functional.conv1d(signal_padded, kernel[None, None, :])[:, 0] - - def _symmetric_pad_1d(signal: torch.Tensor, pad: int) -> torch.Tensor: """Symmetric 1D padding matching scipy's reflect mode. Same edge-repeat semantics as ``_symmetric_pad`` but for 1D signals. Used by ``gaussian_smooth_1d`` for regularization of knot vectors. """ + if pad <= 0: + # signal[:, :, -0:] is the whole signal, not an empty slice, so a naive + # tail slice would double the length. No padding needed when pad == 0. + return signal left = signal[:, :, :pad].flip(-1) right = signal[:, :, -pad:].flip(-1) return torch.cat([left, signal, right], dim=-1) @@ -394,7 +676,7 @@ def _symmetric_pad( return field_stack -def _gaussian_kernel_1d(sigma, dtype, device, _cache={}): +def _gaussian_kernel_1d(sigma: float, dtype: torch.dtype, device: torch.device, _cache: dict = {}) -> torch.Tensor: """Normalized 1D Gaussian ``exp(-0.5*(x/sigma)^2)``, radius ``4*sigma``. Cached via mutable default arg - the grid search calls this ~800 times diff --git a/src/quantem/imaging/drift/core/nonrigid.py b/src/quantem/imaging/drift/core/nonrigid.py index 7620af48..4dbf2997 100644 --- a/src/quantem/imaging/drift/core/nonrigid.py +++ b/src/quantem/imaging/drift/core/nonrigid.py @@ -1,384 +1,246 @@ - +"""Per-scanline non-rigid drift optimization.""" import numpy as np import torch -from scipy.ndimage import gaussian_filter -from scipy.optimize import minimize +import torch.nn.functional as F from tqdm import tqdm -from quantem.imaging.drift import diagnostics -from quantem.imaging.drift.core.knots import gaussian_smooth_1d +import quantem.imaging.drift.plot as drift_plot +from quantem.imaging.drift.core import knots as drift_knots +from quantem.imaging.drift.core.warping import warp_and_translate -def align_nonrigid( - self, - backend: str = "pytorch", - optimizer_name: str = "adam", - num_iterations: int = 8, - regularization_sigma_px: float = 16.0, - regularization_update_step_size: float | None = 0.8, - regularization_poly_order: int = 1, - max_image_shift: float | None = 32.0, - adam_steps: int = 30, - lr: float | None = None, - lbfgs_max_iter: int = 20, - max_optimize_iterations: int = 10, - regularization_max_image_shift_px: float | None = None, - solve_individual_rows: bool = True, - show_merged: bool = True, - show_images: bool = False, - show_knots: bool = True, - **kwargs, -): +def _grid_sample_mse( + grid_row: torch.Tensor, + grid_col: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, +) -> torch.Tensor: + """Shared tail: stack the (col, row) grid, sample, and return the MSE. + + Inlined into both compiled kernels (`@torch.compile` follows the call). + Pulling the grid_sample + MSE out of the two K-paths means the only + difference between K=1 and K>=2 kernels is grid construction. + + The MSE is averaged over both the batch (N images) and the spatial dims, + so each image's gradient is scaled by 1/N relative to a per-image solve. + Adam's adaptive step size absorbs the constant rescale; LBFGS line search + rescales itself. """ - Non-rigid drift correction using PyTorch (default) or SciPy backend. + grid = torch.stack([grid_col, grid_row], dim=-1) + warped = F.grid_sample( + ref_t, grid, mode='bilinear', align_corners=True, padding_mode='border')[:, 0] + return ((warped - target_batch) ** 2).mean() - Parameters - ---------- - backend : str, default "pytorch" - Optimization backend. - - "pytorch": GPU-accelerated batched optimization. Single-knot only. - - "scipy": CPU L-BFGS row-by-row. Use when you need multi-knot - mode (``number_knots > 1``), which the pytorch path does not - yet support. - optimizer_name : str, default "adam" - PyTorch optimizer (ignored if backend="scipy"). - - **"adam"** - first-order momentum optimizer. Default. Best when: - - You want the fastest possible runtime, especially at image - sizes ≤512 px where the per-step grid_sample is small and - Adam's tight inner loop wins on launch overhead. - - You're confident ``max_image_shift`` reflects the true drift - bound (Adam's auto-lr derives from it; if it's too small, - Adam silently under-converges). - - You want bit-reproducible results across runs (LBFGS line - search has subtle non-determinism from Wolfe condition checks). - - **Provisional override guidance** (validated on one real-data - pair - Bob's gold-nanoparticle HAADF on a spectra background - - and the synthetic chevron test; needs broader testing on - diverse datasets before being treated as authoritative). If - you override ``lr`` manually, the rough formula is - ``expected_drift_px / (num_iterations * adam_steps)``. - Indicative starting values for the default ``num_iterations=8, - adam_steps=30`` (240 total steps): - * ~5 px drift (synthetic chevron, small drift): ``lr≈0.02`` - * ~50-100 px drift (gold-nanoparticle HAADF, real STEM): ``lr≈0.5`` - * larger / unknown drift: prefer ``optimizer_name="lbfgs"`` - which auto-scales via line search and doesn't need this - per-dataset tuning. - - **"lbfgs"** - quasi-Newton optimizer with strong-Wolfe line search. - Best when: - - The image is ≥512 px and you don't mind paying Python closure - overhead for fewer total steps (typically 2-3× faster than - Adam at 2048+ px because it converges in ~30 steps not 240). - - You're unsure about the drift magnitude or don't want to think - about ``lr`` tuning - LBFGS line search auto-scales the step - without any hand-tuning. - - You want quality over speed. - - **Failure modes to avoid:** - - **Don't normalize inputs to [0, 1] when using LBFGS** - - strong-Wolfe's curvature condition needs absolute gradient - magnitude above a threshold; with normalized intensities the - gradient is ~1e-4 and the line search returns step=0, - producing zero correction silently. Adam is unaffected. - - **Don't set ``max_image_shift`` smaller than your actual drift - if using Adam with default ``lr=None``** - Adam's auto-derived - lr scales with max_image_shift, so a too-small bound silently - clamps how much drift Adam can recover. LBFGS is unaffected. - - If unsure, start with Adam (the default) for ≤1024 px images and - switch to LBFGS for ≥2048 px or for unknown-drift exploratory work. - - Shared Parameters - ----------------- - num_iterations : int, default 8 - Number of outer iterations for alternating optimization. - regularization_sigma_px : float, default 16.0 - Gaussian smoothing sigma for knot regularization. - regularization_update_step_size : float, default 0.8 - Step size for knot updates (0-1, lower = more conservative). - regularization_poly_order : int, default 1 - Polynomial order for trend removal in knot regularization - (used by both pytorch and scipy backends). - max_image_shift : float, default 32.0 - Maximum shift for translation alignment between iterations. - PyTorch Parameters (ignored if backend="scipy") - ----------------------------------------------- - adam_steps : int, default 30 - Number of Adam optimization steps per outer iteration. - lr : float or None, default None - Learning rate for Adam. When None (default), auto-derived as - ``max_image_shift / (num_iterations * adam_steps * 4)``. - - **Why auto-derive?** Adam's ``m/sqrt(v)`` update self-normalizes - the gradient, so each step moves a knot by ~``lr`` pixels - regardless of image intensity scale. The total movement budget - is ``lr × num_iterations × adam_steps`` and is hard-bounded: - Adam cannot find drift larger than that budget no matter how - many iterations you give it. This means ``lr`` must be matched - to the EXPECTED DRIFT MAGNITUDE IN PIXELS, not to gradient - magnitude - the same default value that works on small synthetic - drift will silently under-converge on real data with larger drift. - - The auto-derived formula reserves half the total step budget for - search (covering up to ``max_image_shift / 2`` of nonlinear drift) - and the other half for refinement near the minimum. - - Override with an explicit float when you know the actual drift - magnitude - e.g. ``lr=2.0`` for very-large-drift in-situ data, - or ``lr=0.005`` for atomic-resolution stable samples. - lbfgs_max_iter : int, default 20 - Maximum LBFGS iterations per outer iteration (line search probes - within each iter happen automatically). Only used when - optimizer="lbfgs". - - SciPy Parameters (ignored if backend="pytorch") - ----------------------------------------------- - max_optimize_iterations : int, default 10 - Maximum L-BFGS iterations per row. - regularization_max_image_shift_px : float, optional - Maximum allowed shift per iteration. - solve_individual_rows : bool, default True - If True, optimize each row independently. - - Display Parameters - ------------------ - show_merged : bool, default True - Show merged image after alignment. - show_images : bool, default False - Show individual aligned images. - show_knots : bool, default True - Overlay knot positions on visualizations. +def _grid_sample_ncc( + grid_row: torch.Tensor, + grid_col: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, + ref_coverage_t: torch.Tensor, +) -> torch.Tensor: + """Warp ``ref_t`` and return coverage-weighted ``1 - mean NCC``. - Notes - ----- - With backend="pytorch", ``self.images_warped`` is left STALE after - the loop and refreshed lazily on first access via plot methods or - ``calculate_error()``. Code that reads ``self.images_warped.array`` - directly should call ``self._ensure_warped_images()`` first, or - use ``generate_corrected_image()`` which builds its own warps from - ``self.knots``. + Reference coverage is sampled through the same grid and detached before + reduction. This excludes canvas fill without rewarding knot motion merely + for changing the valid area. Per-image weighted zero-mean correlation is + then averaged over the batch. """ - if not hasattr(self, "knots"): - raise RuntimeError( - "No knots found. Call .preprocess() before running alignment." - ) - if backend == "pytorch": - device = self._device - dtype = self._dtype - num_images = self.shape[0] - canvas_shape = (self.shape[1], self.shape[2]) - if any(self.knots[i].shape[2] != 1 for i in range(num_images)): - raise NotImplementedError( - "PyTorch backend only supports single knot. " - "Use backend='scipy' for multiple knots.") - knots_batch = torch.tensor( - np.stack([self.knots[i][:, :, 0] for i in range(num_images)]), - dtype=dtype, device=device, requires_grad=True) - num_rows_knot = knots_batch.shape[2] - target_batch = torch.stack(self.images_t) - # Build u tensors once and reuse - same scan-position vector projects - # onto row and col offsets via the per-image scan_fast components. - u_t = [ - torch.as_tensor(self.interpolator[i].u, dtype=dtype, device=device) - for i in range(num_images) - ] - row_scan_offsets = torch.stack([ - u_t[i] * (self.interpolator[i].scan_fast[0] * (self.images[i].shape[0] - 1)) - for i in range(num_images) - ]) - col_scan_offsets = torch.stack([ - u_t[i] * (self.interpolator[i].scan_fast[1] * (self.images[i].shape[1] - 1)) - for i in range(num_images) - ]) - row_scale = 2.0 / (canvas_shape[0] - 1) - col_scale = 2.0 / (canvas_shape[1] - 1) - if optimizer_name == "adam": - # Auto-derive lr so the total movement budget covers a quarter - # of max_image_shift. The safety factor of 4 (not 2) prevents - # over-shooting at small image sizes where actual drift is well - # below max_image_shift; at large sizes the same factor still - # converges because the loss surface is smoother. See the `lr` - # parameter docstring for the full rationale. - adam_lr = lr if lr is not None else max_image_shift / (num_iterations * adam_steps * 4) - optimizer = torch.optim.Adam([knots_batch], lr=adam_lr, fused=True) - elif optimizer_name == "lbfgs": - optimizer = torch.optim.LBFGS( - [knots_batch], lr=1.0, max_iter=lbfgs_max_iter, - line_search_fn="strong_wolfe") - else: - raise ValueError(f"optimizer_name must be 'adam' or 'lbfgs', got {optimizer_name!r}") - if regularization_sigma_px is not None and regularization_sigma_px > 0: - x_knot = torch.arange(num_rows_knot, dtype=dtype, device=device) - x_norm = (x_knot - x_knot.mean()) / x_knot.std() - vander = torch.stack([x_norm ** p for p in range(regularization_poly_order + 1)], dim=1) - warped_t = self._warp_and_translate_torch( - max_image_shift, upsample_factor=8, knots_batch=knots_batch) - error_buffer = [] - for _ in tqdm(range(num_iterations), desc=f"Solving nonrigid drift ({optimizer_name})"): - # Build the reference under no_grad: arithmetic on warped_t (an - # inference tensor) would otherwise return an autograd-tracked - # leaf, and the optimizer would build a graph through it. - with torch.no_grad(): - warped_sum = warped_t.sum(0) - ref_batch = (warped_sum[None] - warped_t) / (num_images - 1) - knots_prev = knots_batch.detach().clone() - # Regularization alters the loss surface between outer iters, so - # stale momentum / curvature history would push knots the wrong way. - optimizer.state.clear() - if optimizer_name == "adam": - self._optimize_knots_adam( - ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, adam_steps) - else: - self._optimize_knots_lbfgs( - ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer) - self._regularize_knots( - knots_batch, knots_prev, vander, - regularization_max_image_shift_px, - regularization_sigma_px, - regularization_update_step_size) - warped_t = self._warp_and_translate_torch( - max_image_shift, upsample_factor=8, knots_batch=knots_batch) - # Per-iter error stays on GPU; sync once after the loop - images_mean = warped_t.mean(dim=0) - error_buffer.append(torch.mean(torch.abs(warped_t - images_mean[None]), dim=(1, 2))) - # Sync knots back to numpy; leave images_warped lazy so callers - # that never plot avoid the GPU→CPU transfer of the warped stack. - knots_final = knots_batch.detach().cpu().numpy() - for img_idx in range(num_images): - self.knots[img_idx][:, :, 0] = knots_final[img_idx] - self._images_warped_stale = True - self._max_image_shift_cached = max_image_shift - if error_buffer: - # Build all error rows in one DtoH transfer + one vstack, instead of - # the quadratic vstack-per-iteration pattern used by calculate_error. - errors_np = torch.stack(error_buffer).cpu().numpy() # (num_iterations, num_images) - mode_col = np.full((len(errors_np), 1), 2.0) - mean_col = errors_np.mean(axis=1, keepdims=True) - new_rows = np.hstack((mode_col, mean_col, errors_np)) - if not hasattr(self, "error_track"): - self.error_track = new_rows - else: - self.error_track = np.vstack((self.error_track, new_rows)) - else: - for _ in tqdm(range(num_iterations), desc="Solving nonrigid drift (scipy)"): - for ind in range(self.shape[0]): - image_ref = np.delete(self.images_warped.array, ind, axis=0).mean(axis=0) - knots_updated = self._optimize_knots_scipy( - ind, image_ref, self.knots[ind], - max_optimize_iterations=max_optimize_iterations, - solve_individual_rows=solve_individual_rows) - if regularization_max_image_shift_px is not None: - knots_shift = knots_updated - self.knots[ind] - knots_dist = np.sqrt(np.sum(knots_shift**2, axis=0)) - sub = knots_dist > regularization_max_image_shift_px - knots_updated[0][sub] = (self.knots[ind][0][sub] - + knots_shift[0][sub] * regularization_max_image_shift_px / knots_dist[sub]) - knots_updated[1][sub] = (self.knots[ind][1][sub] - + knots_shift[1][sub] * regularization_max_image_shift_px / knots_dist[sub]) - if regularization_sigma_px is not None and regularization_sigma_px > 0: - knots_smoothed = knots_updated.copy() - for dim in range(2): - x = np.arange(knots_updated.shape[1]) - for knot_ind in range(knots_updated.shape[2]): - y = knots_updated[dim, :, knot_ind] - coefs = np.polyfit(x, y, deg=regularization_poly_order) - trend = np.polyval(coefs, x) - residual = y - trend - residual_smooth = gaussian_filter(residual, sigma=regularization_sigma_px) - knots_smoothed[dim, :, knot_ind] = residual_smooth + trend - knots_updated = knots_smoothed - if regularization_update_step_size is not None: - knots_updated = (self.knots[ind] - + (knots_updated - self.knots[ind]) * regularization_update_step_size) - self.knots[ind] = knots_updated - warped_t = self._warp_and_translate_torch(max_image_shift, upsample_factor=8) - self.calculate_error(2, _warped_t=warped_t) - - diagnostics._record_stage(self, "nonrigid") - - if show_merged: - self.plot_merged_images( - show_knots=show_knots, - title="Merged: non-rigid", - **kwargs, - ) + grid = torch.stack([grid_col, grid_row], dim=-1) + warped = F.grid_sample( + ref_t, grid, mode='bilinear', align_corners=True, padding_mode='border')[:, 0] + coverage = F.grid_sample( + ref_coverage_t[:, None], + grid, + mode='bilinear', + align_corners=True, + padding_mode='zeros', + )[:, 0].detach().clamp(0.0, 1.0) + # (N, H, W) → (N, H*W), with weighted means/norms on measured pixels. + w = warped.reshape(warped.shape[0], -1) + t = target_batch.reshape(target_batch.shape[0], -1) + mask = coverage.reshape(coverage.shape[0], -1) + count = mask.sum(dim=-1, keepdim=True).clamp_min(1.0) + w0 = w - (w * mask).sum(dim=-1, keepdim=True) / count + t0 = t - (t * mask).sum(dim=-1, keepdim=True) / count + num = (mask * w0 * t0).sum(dim=-1) + den = ( + (mask * w0.square()).sum(dim=-1).sqrt() + * (mask * t0.square()).sum(dim=-1).sqrt() + ).clamp_min(1e-12) + ncc = (num / den).mean() + return 1.0 - ncc - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: non-rigid" for i in range(self.shape[0])], - **kwargs, - ) - return self +@torch.compile(mode="reduce-overhead", dynamic=False) +def _compiled_loss_fn_single( + knots_batch: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, + row_scan_offsets: torch.Tensor, + col_scan_offsets: torch.Tensor, + row_scale: float, + col_scale: float, +) -> torch.Tensor: + """Fused K=1 forward pass: knot anchor + scan_fast walk → MSE. + + ``knots_batch`` shape ``(N, 2, num_rows, 1)``. Each scanline has one + anchor knot and the per-pixel canvas position is filled in by adding + the precomputed ``scan_fast`` walk. ``correct_nonrigid`` selects this + kernel when ``K == 1`` and :func:`_compiled_loss_fn_multi` for ``K > 1``. + """ + grid_row = (knots_batch[:, 0, :, :] + row_scan_offsets[:, None, :]) * row_scale - 1.0 + grid_col = (knots_batch[:, 1, :, :] + col_scan_offsets[:, None, :]) * col_scale - 1.0 + return _grid_sample_mse(grid_row, grid_col, ref_t, target_batch) + + +@torch.compile(mode="reduce-overhead", dynamic=False) +def _compiled_loss_fn_multi( + knots_batch: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, + seg_idx: torch.Tensor, + seg_frac: torch.Tensor, + row_scale: float, + col_scale: float, +) -> torch.Tensor: + """Fused K-knot forward pass with linear knot interpolation along scanline. + + ``knots_batch`` shape ``(N, 2, num_rows, K)`` with ``K >= 2``. + ``seg_idx`` (long, shape ``(num_cols,)``) and ``seg_frac`` (shape + ``(num_cols,)``) precompute, per output column, which adjacent knot + pair to interpolate and the local fraction. Both are constant for + the lifetime of the optimization, so we lift them out of the loop. + """ + knot_lo = knots_batch[:, :, :, seg_idx] + knot_hi = knots_batch[:, :, :, seg_idx + 1] + interp = knot_lo + (knot_hi - knot_lo) * seg_frac[None, None, None, :] + grid_row = interp[:, 0] * row_scale - 1.0 + grid_col = interp[:, 1] * col_scale - 1.0 + return _grid_sample_mse(grid_row, grid_col, ref_t, target_batch) + + +@torch.compile(mode="reduce-overhead", dynamic=False) +def _compiled_loss_fn_single_ncc( + knots_batch: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, + ref_coverage_t: torch.Tensor, + row_scan_offsets: torch.Tensor, + col_scan_offsets: torch.Tensor, + row_scale: float, + col_scale: float, +) -> torch.Tensor: + """K=1 path with ``1 - NCC`` loss (brightness-invariant).""" + grid_row = (knots_batch[:, 0, :, :] + row_scan_offsets[:, None, :]) * row_scale - 1.0 + grid_col = (knots_batch[:, 1, :, :] + col_scan_offsets[:, None, :]) * col_scale - 1.0 + return _grid_sample_ncc( + grid_row, grid_col, ref_t, target_batch, ref_coverage_t + ) + + +@torch.compile(mode="reduce-overhead", dynamic=False) +def _compiled_loss_fn_multi_ncc( + knots_batch: torch.Tensor, + ref_t: torch.Tensor, + target_batch: torch.Tensor, + ref_coverage_t: torch.Tensor, + seg_idx: torch.Tensor, + seg_frac: torch.Tensor, + row_scale: float, + col_scale: float, +) -> torch.Tensor: + """K>=2 path with ``1 - NCC`` loss.""" + knot_lo = knots_batch[:, :, :, seg_idx] + knot_hi = knots_batch[:, :, :, seg_idx + 1] + interp = knot_lo + (knot_hi - knot_lo) * seg_frac[None, None, None, :] + grid_row = interp[:, 0] * row_scale - 1.0 + grid_col = interp[:, 1] * col_scale - 1.0 + return _grid_sample_ncc( + grid_row, grid_col, ref_t, target_batch, ref_coverage_t + ) def _optimize_knots_adam( - self, ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, adam_steps, + ref_batch, target_batch, knots_batch, + loss_fn, loss_args, + optimizer, steps, grad_mask=None, ): - """Run ``adam_steps`` of Adam on a batched knot tensor against ``_compiled_loss_fn``.""" + """Run ``steps`` of Adam on a batched knot tensor against ``loss_fn``.""" ref_t = ref_batch[:, None] - for _ in range(adam_steps): + for _ in range(steps): optimizer.zero_grad() - loss = self._compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale) + loss = loss_fn(knots_batch, ref_t, target_batch, *loss_args) loss.backward() + if grad_mask is not None: + knots_batch.grad.mul_(grad_mask) optimizer.step() -@torch.compile(mode="reduce-overhead", dynamic=False) -def _compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, -): - """Fused forward pass: knot offsets → grid → grid_sample → MSE loss. - - The MSE is averaged over both the batch (N images) and the spatial - dims, so each image's gradient is scaled by 1/N relative to a - per-image solve. Adam's adaptive step size absorbs the constant - rescale; LBFGS line search rescales itself. - """ - grid_row = (knots_batch[:, 0, :, None] + row_scan_offsets[:, None, :]) * row_scale - 1.0 - grid_col = (knots_batch[:, 1, :, None] + col_scan_offsets[:, None, :]) * col_scale - 1.0 - grid = torch.stack([grid_col, grid_row], dim=-1) - warped = torch.nn.functional.grid_sample( - ref_t, grid, mode='bilinear', align_corners=True, padding_mode='border')[:, 0] - return ((warped - target_batch) ** 2).mean() - - def _optimize_knots_lbfgs( - self, ref_batch, target_batch, knots_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale, - optimizer, + ref_batch, target_batch, knots_batch, + loss_fn, loss_args, + optimizer, grad_mask=None, ): """Run one LBFGS outer step (line search re-evaluates the closure several times).""" ref_t = ref_batch[:, None] def closure(): optimizer.zero_grad() - loss = self._compiled_loss_fn( - knots_batch, ref_t, target_batch, - row_scan_offsets, col_scan_offsets, row_scale, col_scale) + loss = loss_fn(knots_batch, ref_t, target_batch, *loss_args) loss.backward() + if grad_mask is not None: + knots_batch.grad.mul_(grad_mask) return loss optimizer.step(closure) +def sobel_gradient_magnitude( + images: torch.Tensor, + pre_smooth: float, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Compute per-image Sobel gradient magnitude with optional Gaussian pre-smooth. + + Returns ``(N, H, W)`` z-score-normalized per image so each output has + zero mean and unit variance — removes gain/offset sensitivity for + cross-detector loss comparisons. + """ + img = images[:, None] # (N, 1, H, W) for conv2d + if pre_smooth > 0: + ks = max(3, int(6 * pre_smooth) | 1) # odd kernel size + x = torch.arange(ks, dtype=dtype, device=device) - ks // 2 + g = torch.exp(-0.5 * (x / max(pre_smooth, 1e-6)) ** 2) + g = g / g.sum() + pad_h = ks // 2 + img = F.pad(img, (pad_h, pad_h, 0, 0), mode='reflect') + img = F.conv2d(img, g.reshape(1, 1, 1, -1)) + img = F.pad(img, (0, 0, pad_h, pad_h), mode='reflect') + img = F.conv2d(img, g.reshape(1, 1, -1, 1)) + sobel_column = torch.tensor( + [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], + dtype=dtype, device=device).reshape(1, 1, 3, 3) + sobel_row = torch.tensor( + [[-1, -2, -1], [0, 0, 0], [1, 2, 1]], + dtype=dtype, device=device).reshape(1, 1, 3, 3) + img_pad = F.pad(img, (1, 1, 1, 1), mode='reflect') + gradient_column = F.conv2d(img_pad, sobel_column) + gradient_row = F.conv2d(img_pad, sobel_row) + grad_mag = (gradient_row**2 + gradient_column**2).sqrt()[:, 0] + mean = grad_mag.mean(dim=(-2, -1), keepdim=True) + std = grad_mag.std(dim=(-2, -1), keepdim=True).clamp(min=1e-8) + return (grad_mag - mean) / std + + def _regularize_knots( - self, knots_batch, knots_prev, vander, - max_shift_px, sigma_px, step_size, -): - """Apply per-iteration knot regularization (in-place on ``knots_batch``). + knots_batch: torch.Tensor, + knots_prev: torch.Tensor, + vander: torch.Tensor | None, + max_shift_px: float | None, + sigma_px: float | None, + step_size: float | None, +) -> None: + """Apply per-iteration knot regularization (in place on ``knots_batch``). Three independent stages, each gated by its parameter being non-None: 1. Per-knot shift cap: clamp ``|new - prev|`` to ``max_shift_px`` @@ -389,62 +251,527 @@ def _regularize_knots( 3. Step-size blend: ``new = prev + step_size · (new - prev)``, under-relaxes the update for stability across outer iterations. """ - num_images, _, num_rows_knot = knots_batch.shape + # Knots are 4D ``(N, 2, R, K)`` everywhere — K=1 just has trailing 1. + num_images, _, num_rows_knot, K = knots_batch.shape with torch.no_grad(): if max_shift_px is not None: shift = knots_batch - knots_prev dist = torch.norm(shift, dim=1, keepdim=True) scale_factor = torch.clamp(max_shift_px / dist.clamp(min=1e-8), max=1.0) knots_batch.copy_(knots_prev + shift * scale_factor) - if sigma_px is not None and sigma_px > 0: - # Detrend + smooth all (N*2, num_rows) knots in one batched lstsq + smooth - knots_flat = knots_batch.reshape(-1, num_rows_knot).T # (num_rows, N*2) - coefs, _, _, _ = torch.linalg.lstsq(vander, knots_flat) - trend = (vander @ coefs).T # (N*2, num_rows) - residual = knots_batch.reshape(-1, num_rows_knot) - trend - smoothed = gaussian_smooth_1d(residual, sigma_px) - knots_batch.copy_((smoothed + trend).reshape(num_images, 2, num_rows_knot)) + if sigma_px is not None and sigma_px > 0 and vander is not None: + # Smooth/detrend along the row axis. Treat each (axis, intra-row + # knot) slot as an independent series along rows by moving the row + # dim last and flattening the leading channels. + knots_flat = knots_batch.permute(0, 1, 3, 2).reshape(-1, num_rows_knot).T + if vander.device.type == "mps": + # MPS does not implement lstsq. The normalized polynomial basis + # has at most four columns, so its full-rank system is small. + if vander.shape[0] < vander.shape[1]: + coefficients = torch.linalg.lstsq( + vander.cpu(), knots_flat.cpu() + ).solution.to(vander.device) + else: + normal_matrix = vander.T @ vander + coefficients = torch.linalg.solve( + normal_matrix, vander.T @ knots_flat + ) + else: + coefficients = torch.linalg.lstsq(vander, knots_flat).solution + trend = (vander @ coefficients).T + residual = knots_flat.T - trend + smoothed = drift_knots.gaussian_smooth_1d(residual, sigma_px) + knots_batch.copy_( + (smoothed + trend) + .reshape(num_images, 2, K, num_rows_knot) + .permute(0, 1, 3, 2)) if step_size is not None: knots_batch.copy_(knots_prev + (knots_batch - knots_prev) * step_size) -def _optimize_knots_scipy( - self, idx: int, image_ref: np.ndarray, knots_init: np.ndarray, - max_optimize_iterations: int = 10, solve_individual_rows: bool = True, -) -> np.ndarray: - """SciPy L-BFGS optimization for one image.""" - shape_knots = knots_init.shape - options = {"maxiter": max_optimize_iterations} if max_optimize_iterations else {} - if solve_individual_rows: - knots_updated = np.zeros_like(knots_init) - for row_ind in range(knots_init.shape[1]): - x0 = knots_init[:, row_ind, :].ravel() - def cost_function(x): - knots_row = x.reshape(shape_knots[0], shape_knots[2]) - xa, ya = self.interpolator[idx].transform_rows(knots_row) - xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) - yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) - dx, dy = xa - xf, ya - yf - warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) - + image_ref[xf + 1, yf] * dx * (1 - dy) - + image_ref[xf, yf + 1] * (1 - dx) * dy - + image_ref[xf + 1, yf + 1] * dx * dy) - return np.sum((warped - self.images[idx].array[row_ind, :]) ** 2) - result = minimize(cost_function, x0, method="L-BFGS-B", options=options) - knots_updated[:, row_ind, :] = result.x.reshape((2, -1)) +def setup_loss_kernel( + self, + K: int, + canvas_shape: tuple[int, int], + loss: str = "mse", +): + """Pick the K-aware compiled loss kernel and precompute its constants. + + Returns ``(loss_fn, loss_args)`` ready to be passed to + ``_optimize_knots_adam`` / ``_optimize_knots_lbfgs``. + + K=1 path: scan_fast walk offsets per image (``row_scan_offsets`` / + ``col_scan_offsets``). K>=2 path: per-output-column segment indices + + local fractions for the linear knot interpolation. Constants are + lifted out of the inner Adam / LBFGS loop so the compiled kernel + sees them as static. + + ``loss`` selects the scalar: ``"mse"`` / ``"gradient_mse"`` use MSE + kernels; ``"ncc"`` uses ``1 - mean NCC`` (brightness-invariant). + """ + device, dtype = self._device, self._dtype + num_images = self.shape[0] + row_scale = 2.0 / (canvas_shape[0] - 1) + col_scale = 2.0 / (canvas_shape[1] - 1) + use_ncc = loss == "ncc" + if K > 1: + # Identical to _transform_coordinates_multi_knot's geometry so + # apply_correction downstream inverts it cleanly. + num_cols = self.imgs[0].shape[1] + t = torch.linspace(0, 1, num_cols, dtype=dtype, device=device) * (K - 1) + seg_idx = torch.clamp(t.long(), max=K - 2) + seg_frac = t - seg_idx.to(dtype) + fn = ( + _compiled_loss_fn_multi_ncc + if use_ncc + else _compiled_loss_fn_multi + ) + return (fn, (seg_idx, seg_frac, row_scale, col_scale)) + # K=1: same scan-position vector projects onto row/col via scan_fast. + u_t = [ + torch.as_tensor(self.u_per_image[i], dtype=dtype, device=device) + for i in range(num_images) + ] + row_scan_offsets = torch.stack( + [ + u_t[i] * (self.scan_fast[i][0] * (self.imgs[i].shape[0] - 1)) + for i in range(num_images) + ] + ) + col_scan_offsets = torch.stack( + [ + u_t[i] * (self.scan_fast[i][1] * (self.imgs[i].shape[1] - 1)) + for i in range(num_images) + ] + ) + fn = ( + _compiled_loss_fn_single_ncc + if use_ncc + else _compiled_loss_fn_single + ) + return (fn, (row_scan_offsets, col_scan_offsets, row_scale, col_scale)) + + +def correct_nonrigid( + self, + *, + num_knots: int | None = None, + optimizer: str = "adam", + num_refine_cycles: int = 16, + knot_smoothing_sigma: float = 8.0, + update_fraction: float | None = 0.8, + trend_order: int = 1, + max_image_shift: float = 32.0, + optimizer_steps: int | None = 30, + learning_rate: float | None = None, + max_knot_step: float | None = None, + fixed_scans: list[int] | None = None, + loss: str = "auto", + edge_smoothing_sigma: float = 1.0, + early_stop_patience: int = 3, + early_stop_rtol: float = 1e-4, + min_iterations: int = 4, + show_combined: bool = True, + show_scans: bool = False, + show_knots: bool = True, + show_knot_plot: bool = False, + show_report: bool = False, + verbose: bool = True, +): + """Remove residual scanline drift after affine or strip correction. + + The optimizer updates scanline knots against the other corrected scans or + a fixed reference. Smooth knot regularization suppresses scanline jitter + while retaining drift that changes gradually across the acquisition. + + Prefer :meth:`correct_strip` first when residual after affine is a + smooth function of slow-scan position (deceleration / bend), especially + for reference-mode EDS on pure lattices - free nonrigid with a large + ``max_image_shift`` can lattice-alias. On atomic lattices keep + ``max_image_shift`` small (e.g. 2). + + Calls are cumulative. On an atomic lattice, begin with + ``max_image_shift=1`` or ``2`` and strong regularization (8-16 rows), then + accept the stage only when common-mask NCC and the RGB overlay improve. + In reference mode, the affine/strip solution already establishes the + external reference frame, so nonrigid iterations preserve that frame + instead of re-solving a potentially lattice-aliased global translation. + + Parameters + ---------- + num_knots : int or None, default None + Knots along every fast-scan line. ``None`` keeps the current layout; + use more than one when residual drift changes within a scanline. The + existing affine or strip field is preserved before non-rigid + optimization begins. + optimizer : str, default "adam" + ``"adam"`` (first-order momentum) or ``"lbfgs"`` (quasi-Newton + with strong-Wolfe line search). Adam is the fastest default for + ≤1024 px images. LBFGS auto-scales the step size and is preferred + for ≥2048 px or when the drift magnitude is unknown. Don't normalize + inputs to [0, 1] when using LBFGS - strong-Wolfe needs absolute + gradient magnitude and silently returns step=0 on unit-variance images. + num_refine_cycles : int, default 16 + Outer iterations for alternating reference build + knot update. + knot_smoothing_sigma : float, default 8.0 + Gaussian smoothing sigma for knot regularization. 4-12 typical + for STEM data; smaller = finer per-row correction. + update_fraction : float, default 0.8 + Step size for knot updates (0-1, lower = more conservative). + trend_order : int, default 1 + Polynomial order for trend removal in knot regularization. + max_image_shift : float, default 32.0 + Maximum shift for translation alignment between iterations. + Adam's auto-``learning_rate`` derives from this - set close to your expected + drift bound, otherwise Adam silently under-converges. + optimizer_steps : int or None, default 30 + Inner optimizer steps per refinement cycle. ``None`` selects 30 + for Adam and 20 for LBFGS. + learning_rate : float or None, default None + Adam learning rate. When ``None``, it is derived as + ``max_image_shift / (num_refine_cycles * optimizer_steps * 4)``. + Adam's ``m/sqrt(v)`` update self-normalizes the gradient, so each + step moves a knot by ~``learning_rate`` pixels regardless of intensity scale. + Override only when you know the actual drift magnitude. + fixed_scans : list[int] or None, default None + Indices of images whose knots are frozen (their mean becomes the + target for every moving image). Use ``[0]`` for single-sided + alignment with a HAADF reference. Auto-set to ``[0]`` in + reference mode (when the constructor was called with a 2-D ref + + ≥3-D drifted dataset). + loss : str, default "auto" + Most drift workflows compare **similar detectors** (0°/90° HAADF, + survey HAADF vs EDS-session HAADF, etc.). ``"auto"`` follows that + typical case and resolves to ``"ncc"`` (coverage-masked normalized + cross-correlation), so the train loss matches regional NCC reports. + Explicit choices: + + - ``"ncc"`` - coverage-masked ``1 - mean NCC``; brightness / offset + invariant; excludes unmeasured canvas fill. Default via ``"auto"`` + for similar-detector pairs. + - ``"mse"`` - raw-intensity MSE when both images share the same + detector and nearly identical gain (no intensity bias). + - ``"gradient_mse"`` - Sobel-gradient MSE after pre-smooth + + per-image z-score; only needed for **dissimilar / cross-detector** + pairs (e.g. HAADF + virtual dark field). + + edge_smoothing_sigma : float, default 1.0 + Gaussian sigma applied before Sobel when ``loss="gradient_mse"``. + Set to 0 to disable. Ignored for ``"mse"`` / ``"ncc"``. + early_stop_patience : int, default 3 + Stop when ``patience`` consecutive iterations show no improvement. + Set to ``num_refine_cycles`` to disable. + early_stop_rtol : float, default 1e-4 + Minimum relative improvement to count as progress. + min_iterations : int, default 4 + Floor before early stopping can trigger. + show_combined, show_scans : bool + Display knobs forwarded to the plot helpers. + show_report : bool, default False + Print all completed regional NCC checkpoints, including the + nonrigid result. + verbose : bool, default True + Show refinement-cycle progress with the current and best mean + absolute alignment error. Set to ``False`` for compact notebooks. + + Returns + ------- + Self + For method chaining. + + Examples + -------- + Similar-detector HAADF pair (``loss="auto"`` → NCC): + + >>> dc = DriftCorrection(im0, im1, scan_direction_degrees=[0, 90]) + >>> dc.correct_affine(show_combined=False) + >>> dc.correct_nonrigid() # auto → ncc + + Atomic lattice: keep the shift tiny so nonrigid cannot lattice-alias: + + >>> dc.correct_nonrigid(max_image_shift=2) + + Let residual motion vary along the fast-scan direction: + + >>> dc.correct_nonrigid(num_knots=6, max_image_shift=2) + + Dissimilar / cross-detector only (HAADF + VDF): + + >>> dc.correct_nonrigid(loss="gradient_mse", knot_smoothing_sigma=8.0) + + Notes + ----- + Plotting refreshes the cached warped images lazily. ``corrected()`` builds + its result directly from the fitted knots. + """ + if not hasattr(self, "knots"): + raise RuntimeError("No knots found. Call .preprocess() before running alignment.") + # Reloaded tensors start on the host; move solve state to the selected + # device before optimization. + self.imgs_t = [t.to(self._device) for t in self.imgs_t] + self.knots = [k.to(self._device) for k in self.knots] + for attr in ("_knots_after_affine", "_knots_after_strip", "_initial_knots"): + snapshot = getattr(self, attr, None) + if snapshot is not None: + setattr(self, attr, [k.to(self._device) for k in snapshot]) + if num_knots is not None: + drift_knots.resize_scanline_knots(self, num_knots) + if loss == "auto": + loss = "ncc" + valid_losses = ("mse", "gradient_mse", "ncc") + if loss not in valid_losses: + raise ValueError(f"loss must be one of {valid_losses!r} or 'auto', got {loss!r}") + if optimizer == "lbfgs" and self._normalized: + import warnings + + warnings.warn( + "normalize=True + LBFGS can cause silent convergence failure. " + "Wolfe line search may return step=0 on unit-variance images. " + "Consider using optimizer='adam' or normalize=False.", + UserWarning, + stacklevel=2, + ) + # Reference-mode auto-anchors the reference image (index 0). + if fixed_scans is None and self._reference_mode: + fixed_scans = [0] + fixed_set = frozenset(fixed_scans) if fixed_scans is not None else frozenset() + moving_indices = [i for i in range(self.shape[0]) if i not in fixed_set] + if fixed_set and not moving_indices: + raise ValueError( + f"All {self.shape[0]} images are in fixed_scans - nothing left to " + "optimize. fixed_scans must leave at least one moving image." + ) + device = self._device + dtype = self._dtype + num_images = self.shape[0] + canvas_shape = (self.shape[1], self.shape[2]) + K_per_image = {self.knots[i].shape[2] for i in range(num_images)} + if len(K_per_image) > 1: + raise ValueError(f"All images must use the same number of knots, got {K_per_image}") + K = K_per_image.pop() + # Knots are 4D throughout the optimizer - K=1 just keeps the trailing 1 + # so downstream code (regularizer, warp, sync) doesn't branch on shape. + knots_batch = ( + torch.stack([self.knots[i] for i in range(num_images)]).detach().requires_grad_(True) + ) + num_rows_knot = knots_batch.shape[2] + # a reloaded (AutoSerialize) object leaves imgs_t on CPU; alignment + # tensors must live on the solve device or Sobel/warp kernels mismatch + target_batch = torch.stack(self.imgs_t).to(self._device) + loss_fn, loss_args = setup_loss_kernel(self, K, canvas_shape, loss=loss) + optimizer_steps = ( + optimizer_steps if optimizer_steps is not None else (30 if optimizer == "adam" else 20) + ) + if optimizer == "adam": + # Auto-derive lr so the total movement budget covers a quarter + # of max_image_shift. The safety factor of 4 (not 2) prevents + # over-shooting at small image sizes where actual drift is well + # below max_image_shift; at large sizes the same factor still + # converges because the loss surface is smoother. See the `lr` + # parameter docstring for the full rationale. + adam_lr = ( + learning_rate + if learning_rate is not None + else max_image_shift / (num_refine_cycles * optimizer_steps * 4) + ) + torch_optimizer = torch.optim.Adam([knots_batch], lr=adam_lr, fused=True) + elif optimizer == "lbfgs": + torch_optimizer = torch.optim.LBFGS( + [knots_batch], lr=1.0, max_iter=optimizer_steps, line_search_fn="strong_wolfe" + ) + else: + raise ValueError(f"optimizer must be 'adam' or 'lbfgs', got {optimizer!r}") + if knot_smoothing_sigma is not None and knot_smoothing_sigma > 0: + x_knot = torch.arange(num_rows_knot, dtype=dtype, device=device) + x_norm = (x_knot - x_knot.mean()) / x_knot.std() + vander = torch.stack([x_norm**p for p in range(trend_order + 1)], dim=1) + else: + vander = None + # For gradient_mse, warp the edge-filtered images instead of the + # raw ones. Knots are spatial transforms independent of image + # content, so optimizing in gradient space yields the same drift + # field while being robust to intensity/contrast differences. + # imgs_t_override threads Sobel images through warp_and_translate + # without mutating self.imgs_t (which would silently corrupt + # apply_correction, visualization, and error metrics afterwards). + if loss == "gradient_mse": + sobel_batch = sobel_gradient_magnitude( + target_batch, edge_smoothing_sigma, device, dtype + ) + imgs_t_override = [sobel_batch[i] for i in range(num_images)] + target_batch = sobel_batch else: - x0 = knots_init.ravel() - def cost_function(x): - knots = x.reshape(shape_knots) - xa, ya = self.interpolator[idx].transform_coordinates(knots) - xf = np.clip(np.floor(xa).astype(int), 0, self.shape[1] - 2) - yf = np.clip(np.floor(ya).astype(int), 0, self.shape[2] - 2) - dx, dy = xa - xf, ya - yf - warped = (image_ref[xf, yf] * (1 - dx) * (1 - dy) - + image_ref[xf + 1, yf] * dx * (1 - dy) - + image_ref[xf, yf + 1] * (1 - dx) * dy - + image_ref[xf + 1, yf + 1] * dx * dy) - return np.sum((warped - self.images[idx].array) ** 2) - result = minimize(cost_function, x0, method="L-BFGS-B", options=options) - knots_updated = result.x.reshape(shape_knots) - return knots_updated + imgs_t_override = None + # Affine (and an optional preceding strip stage) already establishes + # the absolute frame in reference mode. Re-solving a global shift here + # can hop by one lattice period after an otherwise good correction. + # Mutual multi-scan mode still solves translation every outer step. + solve_global_translation = not self._reference_mode + warp_result = warp_and_translate( + self, + max_image_shift, + upsample_factor=8, + knots_batch=knots_batch, + solve_translation=solve_global_translation, + fixed_indices=fixed_set, + imgs_t_override=imgs_t_override, + return_weights=loss == "ncc", + ) + if loss == "ncc": + warped_t, coverage_weights_t = warp_result + else: + warped_t = warp_result + coverage_weights_t = None + # Build a boolean mask on device to zero fixed gradients efficiently. + # knots_batch is always 4D ``(N, 2, R, K)``; broadcast over (2, R, K). + if fixed_set: + grad_mask = torch.ones(num_images, 1, 1, 1, dtype=dtype, device=device) + for idx in fixed_set: + grad_mask[idx] = 0.0 + error_buffer = [] + best_error = float("inf") + patience_counter = 0 + pbar = tqdm( + range(num_refine_cycles), + desc=f"Solving nonrigid drift ({optimizer})", + disable=not verbose, + ) + for iter_idx in pbar: + # Build the reference under no_grad: arithmetic on warped_t (an + # inference tensor) would otherwise return an autograd-tracked + # leaf, and the optimizer would build a graph through it. + with torch.no_grad(): + if fixed_set: + # Fixed images define the reference for all moving images. + fixed_mean = warped_t[sorted(fixed_set)].mean(0) + ref_batch = fixed_mean[None].expand(num_images, -1, -1) + else: + warped_sum = warped_t.sum(0) + ref_batch = (warped_sum[None] - warped_t) / (num_images - 1) + knots_prev = knots_batch.detach().clone() + if loss == "ncc": + # ``bilinear_kde_batch`` considers weights >= 1e-3 covered. + # Reuse that exact ramp so the optimizer and renderer agree + # about which reference pixels are measured rather than fill. + coverage_batch = (coverage_weights_t / 1e-3).clamp(0.0, 1.0) + if fixed_set: + fixed_coverage = coverage_batch[sorted(fixed_set)].amin(0) + ref_coverage_batch = fixed_coverage[None].expand( + num_images, -1, -1 + ) + else: + ref_coverage_batch = torch.stack( + [ + coverage_batch[ + [j for j in range(num_images) if j != i] + ].amin(0) + for i in range(num_images) + ] + ) + cycle_loss_args = (ref_coverage_batch, *loss_args) + else: + cycle_loss_args = loss_args + # Regularization alters the loss surface between outer iters, so + # stale momentum / curvature history would push knots the wrong way. + torch_optimizer.state.clear() + if optimizer == "adam": + _optimize_knots_adam( + ref_batch, + target_batch, + knots_batch, + loss_fn, + cycle_loss_args, + torch_optimizer, + optimizer_steps, + grad_mask=grad_mask if fixed_set else None, + ) + else: + _optimize_knots_lbfgs( + ref_batch, + target_batch, + knots_batch, + loss_fn, + cycle_loss_args, + torch_optimizer, + grad_mask=grad_mask if fixed_set else None, + ) + _regularize_knots( + knots_batch, + knots_prev, + vander, + max_knot_step, + knot_smoothing_sigma, + update_fraction, + ) + # Restore fixed knots - regularization is a global smooth that + # would subtly shift them via polynomial detrend + Gaussian blur. + if fixed_set: + with torch.no_grad(): + for idx in fixed_set: + knots_batch[idx] = knots_prev[idx] + warp_result = warp_and_translate( + self, + max_image_shift, + upsample_factor=8, + knots_batch=knots_batch, + solve_translation=solve_global_translation, + fixed_indices=fixed_set, + imgs_t_override=imgs_t_override, + return_weights=loss == "ncc", + ) + if loss == "ncc": + warped_t, coverage_weights_t = warp_result + else: + warped_t = warp_result + # Per-iter error stays on GPU; sync once after the loop + images_mean = warped_t.mean(dim=0) + iter_error = torch.mean(torch.abs(warped_t - images_mean[None]), dim=(1, 2)) + error_buffer.append(iter_error) + # Early stopping: monitor post-iteration alignment quality + current_error = float(iter_error.mean()) + if current_error < best_error * (1 - early_stop_rtol): + best_error = current_error + patience_counter = 0 + else: + patience_counter += 1 + pbar.set_postfix_str( + f"error={current_error:.4g}, best={best_error:.4g}", + refresh=verbose, + ) + if iter_idx >= min_iterations - 1 and patience_counter >= early_stop_patience: + pbar.set_postfix_str( + f"converged at cycle {iter_idx + 1}, " + f"error={current_error:.4g}, best={best_error:.4g}", + refresh=verbose, + ) + break + # Sync knots back; leave imgs_warped lazy so callers + # that never plot avoid the GPU→CPU transfer of the warped stack. + knots_final = knots_batch.detach() + for img_idx in range(num_images): + self.knots[img_idx][...] = knots_final[img_idx] + self._images_warped_stale = True + self._max_image_shift_cached = max_image_shift + if error_buffer: + # Transfer and append the full convergence history once. + errors_np = torch.stack(error_buffer).cpu().numpy() # (num_iterations, num_images) + mode_col = np.full((len(errors_np), 1), 2.0) + mean_col = errors_np.mean(axis=1, keepdims=True) + new_rows = np.hstack((mode_col, mean_col, errors_np)) + if not hasattr(self, "error_track"): + self.error_track = new_rows + else: + self.error_track = np.vstack((self.error_track, new_rows)) + + drift_plot.show_after_step( + self, + "non-rigid", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + if show_report: + print(self.report().to_string()) + return self diff --git a/src/quantem/imaging/drift/core/strip.py b/src/quantem/imaging/drift/core/strip.py index 0c644ead..14a5dae6 100644 --- a/src/quantem/imaging/drift/core/strip.py +++ b/src/quantem/imaging/drift/core/strip.py @@ -1 +1,839 @@ -"""Package boundary for the strip-correction stage added in a later PR.""" +"""Piecewise-rigid correction for slow-scan residual drift. + +After affine correction, scan deceleration can leave displacement that changes +smoothly with slow-scan position. Masked NCC measures one rigid shift per strip, +then interpolation maps those measurements back to the scanline knots. +""" + +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np +import torch +from tqdm import tqdm + +import quantem.imaging.drift.core.warping as warping +import quantem.imaging.drift.plot as drift_plot + + +@dataclass(frozen=True, slots=True) +class StripPass: + """Describe one coarse-to-fine strip correction pass. + + Strip passes remove slow-scan-dependent residual displacement left after + affine correction. Wider searches recover large residuals; later passes + use more strips and narrower bounds to refine local alignment without + introducing scanline jitter. + + Parameters + ---------- + num_strips : int, optional + Number of slow-scan regions measured independently. At least two are + required to measure variation along the slow-scan direction. Default + is 24. + max_row_shift, max_column_shift : int, optional + Search bounds in scan pixels. Defaults are 8 and 80. + correction_start_fraction : float or None, optional + Fraction of the scan held fixed before residual correction begins. + ``None`` corrects the full field. Default is ``None``. + ramp_fraction : float, optional + Fraction of the scan used to blend into the corrected region. + Default is 0.08. + smoothing_sigma : float, optional + Gaussian smoothing along the slow-scan direction in scanlines. + Default is 12. + update_fraction : float, optional + Fraction of the measured residual applied in this pass. Default is 1. + """ + + num_strips: int = 24 + max_row_shift: int = 8 + max_column_shift: int = 80 + correction_start_fraction: float | None = None + ramp_fraction: float = 0.08 + smoothing_sigma: float = 12.0 + update_fraction: float = 1.0 + + def __post_init__(self): + if self.num_strips < 2: + raise ValueError( + f"num_strips must be at least 2, got {self.num_strips}" + ) + + +# --------------------------------------------------------------------------- +# pure numerics (no DriftCorrection dependency) +# --------------------------------------------------------------------------- + + +def free_weight( + n_rows: int, + *, + free_from_frac: float | None = 0.55, + ramp_frac: float = 0.08, +) -> np.ndarray: + """Per-scanline weight in ``[0, 1]``: 0 = freeze, 1 = apply residual. + + ``free_from_frac=None`` → all ones (full FOV residual). + """ + if free_from_frac is None: + return np.ones(n_rows, dtype=np.float32) + free0 = int(np.clip(free_from_frac, 0.0, 1.0) * n_rows) + ramp = max(8, int(ramp_frac * n_rows)) + w = np.zeros(n_rows, dtype=np.float32) + w[free0:] = 1.0 + r0 = max(0, free0 - ramp) + if free0 > r0: + t = np.linspace(0.0, 1.0, free0 - r0, dtype=np.float32) + t = t * t * (3.0 - 2.0 * t) # smoothstep + w[r0:free0] = t + return w + + +def _as_torch2d(x, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if isinstance(x, torch.Tensor): + t = x.detach().to(device=device, dtype=dtype) + else: + t = torch.as_tensor(np.asarray(x), device=device, dtype=dtype) + if t.ndim != 2: + raise ValueError(f"expected 2-D image, got shape {tuple(t.shape)}") + return t + + +def _masked_ncc_batch( + ref_s: torch.Tensor, + mov_s: torch.Tensor, + mask_s: torch.Tensor, +) -> torch.Tensor: + """Masked NCC for a batch of strips. Shapes ``(S, Hs, W)`` → ``(S,)``.""" + m = mask_s.to(dtype=ref_s.dtype) + n = m.sum(dim=(-2, -1)).clamp_min(1.0) + ref_mu = (ref_s * m).sum(dim=(-2, -1), keepdim=True) / n.view(-1, 1, 1) + mov_mu = (mov_s * m).sum(dim=(-2, -1), keepdim=True) / n.view(-1, 1, 1) + r = (ref_s - ref_mu) * m + v = (mov_s - mov_mu) * m + num = (r * v).sum(dim=(-2, -1)) + den = r.norm(dim=(-2, -1)) * v.norm(dim=(-2, -1)) + return num / den.clamp_min(1e-12) + + +def _ncc_at_shift( + ref_s: torch.Tensor, + mov_s: torch.Tensor, + msk_s: torch.Tensor, + dr: int, + dc: int, +) -> torch.Tensor: + """Masked NCC after rolling all strips by the same ``(dr, dc)``. + + Mask semantics match the historical brute search: + ``roll(mask, dr, row) * roll(mask, dc, col)`` (not a single 2-D roll). + """ + mov_rc = torch.roll(torch.roll(mov_s, shifts=dr, dims=-2), shifts=dc, dims=-1) + msk_rc = torch.roll(msk_s, shifts=dr, dims=-2) * torch.roll(msk_s, shifts=dc, dims=-1) + n_pix = msk_rc.sum(dim=(-2, -1)) + ncc = _masked_ncc_batch(ref_s, mov_rc, msk_rc) + return torch.where(n_pix >= 64.0, ncc, torch.full_like(ncc, -1.0e9)) + + +def _roll2d_per_strip(x: torch.Tensor, dr: torch.Tensor, dc: torch.Tensor) -> torch.Tensor: + """Circular 2-D roll with a different ``(dr, dc)`` per strip. ``x`` is ``(S, H, W)``.""" + S, H, W = x.shape + # torch.roll(x, +k) → x[(i - k) % n]; same via advanced indexing, one gather kernel. + rows = (torch.arange(H, device=x.device)[None, :] - dr.to(dtype=torch.long)[:, None]) % H + cols = (torch.arange(W, device=x.device)[None, :] - dc.to(dtype=torch.long)[:, None]) % W + s_idx = torch.arange(S, device=x.device)[:, None, None] + r_idx = rows[:, :, None].expand(S, H, W) + c_idx = cols[:, None, :].expand(S, H, W) + return x[s_idx, r_idx, c_idx] + + +def _ncc_at_per_strip_shifts( + ref_s: torch.Tensor, + mov_s: torch.Tensor, + msk_s: torch.Tensor, + dr: torch.Tensor, + dc: torch.Tensor, +) -> torch.Tensor: + """Masked NCC with a different integer shift per strip (``dr``, ``dc`` shape ``(S,)``).""" + dr = dr.to(dtype=torch.long, device=mov_s.device) + dc = dc.to(dtype=torch.long, device=mov_s.device) + mov_rc = _roll2d_per_strip(mov_s, dr, dc) + # Historical mask: roll(m, dr, row) * roll(m, dc, col) - product of 1-D rolls. + msk_r = _roll2d_per_strip(msk_s, dr, torch.zeros_like(dr)) + msk_c = _roll2d_per_strip(msk_s, torch.zeros_like(dc), dc) + msk_rc = msk_r * msk_c + n_pix = msk_rc.sum(dim=(-2, -1)) + ncc = _masked_ncc_batch(ref_s, mov_rc, msk_rc) + return torch.where(n_pix >= 64.0, ncc, torch.full_like(ncc, -1.0e9)) + + +def _search_shifts_brute( + ref_s: torch.Tensor, + mov_s: torch.Tensor, + msk_s: torch.Tensor, + max_shift_row: int, + max_shift_col: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Exhaustive integer search; strips batched, shifts in Python (small caps only).""" + S = ref_s.shape[0] + device, dtype = ref_s.device, ref_s.dtype + best_ncc = torch.full((S,), -1.0e9, device=device, dtype=dtype) + best_dr = torch.zeros(S, device=device, dtype=torch.long) + best_dc = torch.zeros(S, device=device, dtype=torch.long) + for dr in range(-int(max_shift_row), int(max_shift_row) + 1): + for dc in range(-int(max_shift_col), int(max_shift_col) + 1): + ncc = _ncc_at_shift(ref_s, mov_s, msk_s, dr, dc) + improved = ncc > best_ncc + best_ncc = torch.where(improved, ncc, best_ncc) + best_dr = torch.where(improved, torch.full_like(best_dr, dr), best_dr) + best_dc = torch.where(improved, torch.full_like(best_dc, dc), best_dc) + return best_dr, best_dc, best_ncc + + +def _search_shifts_fft( + ref_s: torch.Tensor, + mov_s: torch.Tensor, + msk_s: torch.Tensor, + max_shift_row: int, + max_shift_col: int, + *, + refine_radius: int = 5, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """GPU FFT phase-corr candidate + local true masked-NCC refine. + + Cost is ~one ``rfft2`` pair over all strips plus ``(2R+1)²`` gather-NCC + evals - not ``(2·max_row+1)·(2·max_col+1)`` full-image rolls. + ``refine_radius=5`` matches exhaustive brute on EDS residual tests. + """ + S, hs, W = ref_s.shape + device, dtype = ref_s.device, ref_s.dtype + max_r = int(max_shift_row) + max_c = int(max_shift_col) + R = int(refine_radius) + + # Zero-mean within mask → circular cross-correlation peak ≈ best shift. + n = msk_s.sum(dim=(-2, -1), keepdim=True).clamp_min(1.0) + ref_mu = (ref_s * msk_s).sum(dim=(-2, -1), keepdim=True) / n + mov_mu = (mov_s * msk_s).sum(dim=(-2, -1), keepdim=True) / n + r = (ref_s - ref_mu) * msk_s + v = (mov_s - mov_mu) * msk_s + + # sum_i ref[i] * mov[i - lag] peaks at lag = roll amount of mov toward ref. + # ifft(fft(r) * conj(fft(v)))[lag] ≈ that sum → lag is our (dr, dc). + Fr = torch.fft.rfft2(r) + Fm = torch.fft.rfft2(v) + corr = torch.fft.irfft2(Fr * Fm.conj(), s=(hs, W)) # (S, hs, W) + + # Restrict peak search to the allowed window (positive + wrapped negative). + # Build a boolean window once, gather scores. + rr = torch.arange(hs, device=device) + cc = torch.arange(W, device=device) + # lag dr maps to index dr % hs; allow dr in [-max_r, max_r] + if max_r == 0: + row_ok = rr == 0 + else: + row_ok = (rr <= max_r) | (rr >= hs - max_r) + if max_c == 0: + col_ok = cc == 0 + else: + col_ok = (cc <= max_c) | (cc >= W - max_c) + window = row_ok[:, None] & col_ok[None, :] # (hs, W) + neg_inf = torch.finfo(dtype).min + corr_win = torch.where(window[None], corr, torch.full_like(corr, neg_inf)) + flat = corr_win.reshape(S, -1) + peak = flat.argmax(dim=-1) + peak_dr = peak // W # 0..hs-1 + peak_dc = peak % W + + # Convert FFT indices → signed roll amounts in [-max, max] + cand_dr = torch.where(peak_dr <= max_r, peak_dr, peak_dr - hs) + cand_dc = torch.where(peak_dc <= max_c, peak_dc, peak_dc - W) + cand_dr = cand_dr.clamp(-max_r, max_r) + cand_dc = cand_dc.clamp(-max_c, max_c) + + # Local true masked-NCC refine around the FFT peak (handles mask / sign edge cases). + best_ncc = torch.full((S,), -1.0e9, device=device, dtype=dtype) + best_dr = cand_dr.clone() + best_dc = cand_dc.clone() + for ddr in range(-R, R + 1): + for ddc in range(-R, R + 1): + dr = (cand_dr + ddr).clamp(-max_r, max_r) + dc = (cand_dc + ddc).clamp(-max_c, max_c) + ncc = _ncc_at_per_strip_shifts(ref_s, mov_s, msk_s, dr, dc) + improved = ncc > best_ncc + best_ncc = torch.where(improved, ncc, best_ncc) + best_dr = torch.where(improved, dr, best_dr) + best_dc = torch.where(improved, dc, best_dc) + return best_dr, best_dc, best_ncc + + +@torch.no_grad() +def measure_strip_residual_torch( + reference: np.ndarray | torch.Tensor, + moving: np.ndarray | torch.Tensor, + mask: np.ndarray | torch.Tensor, + *, + n_strips: int = 12, + max_shift_col: int = 80, + max_shift_row: int = 8, + device: str | torch.device | None = None, + min_mask_frac: float = 0.25, + method: str = "auto", +) -> dict[str, object]: + """Batched strip residual of ``moving`` vs fixed ``reference`` (GPU torch). + + For each horizontal strip, find integer shifts + ``(dr, dc) ∈ [-max_shift_row, max_shift_row] × [-max_shift_col, max_shift_col]`` + maximizing masked NCC. All strips share the same frozen images. + + Parameters + ---------- + reference, moving + 2-D images in the same frame (typically reference HAADF and + affine-corrected EDS HAADF). + mask + Boolean coverage / common FOV. False pixels are excluded from NCC. + n_strips + Number of horizontal bands (slow-scan blocks). + max_shift_col, max_shift_row + Integer search half-width in pixels. Defaults are large enough that + post-affine EDS residual is not clipped (was ±40/±4 historically). + device + Torch device; default CUDA if available else CPU. + min_mask_frac + Skip strips with valid fraction below this (no residual written). + method + ``"auto"`` (default) - FFT phase-corr + local masked-NCC refine for + large windows (fast on GPU); exact brute for tiny windows. + ``"fft"`` - always FFT+refine. + ``"brute"`` - exhaustive (slow for large caps; exact). + + Returns + ------- + dict + ``centers`` (S,), ``drow`` (S,), ``dcol`` (S,), ``ncc`` (S,), + ``valid`` (S,) bool, ``n_strips``, ``strip_height``. + """ + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(device) + dtype = torch.float32 + + ref = _as_torch2d(reference, device, dtype) + mov = _as_torch2d(moving, device, dtype) + if isinstance(mask, torch.Tensor): + msk = mask.to(device=device, dtype=dtype) + if msk.ndim != 2: + raise ValueError(f"expected 2-D mask, got shape {tuple(msk.shape)}") + else: + msk = _as_torch2d(np.asarray(mask, dtype=np.float32), device, dtype) + if ref.shape != mov.shape or ref.shape != msk.shape: + raise ValueError( + f"shape mismatch ref{tuple(ref.shape)} mov{tuple(mov.shape)} mask{tuple(msk.shape)}" + ) + + H, W = ref.shape + if n_strips < 2: + raise ValueError(f"n_strips must be >= 2, got {n_strips}") + hs = H // n_strips + if hs < 1: + raise ValueError(f"strip height {hs} too small for H={H}, n_strips={n_strips}") + + S = n_strips + H_use = S * hs + ref_s = ref[:H_use].reshape(S, hs, W) + mov_s = mov[:H_use].reshape(S, hs, W) + msk_s = msk[:H_use].reshape(S, hs, W) + mask_frac = msk_s.mean(dim=(-2, -1)) # (S,) + + n_shifts = (2 * int(max_shift_row) + 1) * (2 * int(max_shift_col) + 1) + method = (method or "auto").lower() + if method == "auto": + # Prefer FFT on CUDA for anything beyond a tiny window - brute launches + # one full-strip roll kernel per (dr, dc) and wastes the GPU on large caps. + if device.type == "cuda" and n_shifts > 81: + method = "fft" + elif n_shifts > 200: + method = "fft" + else: + method = "brute" + + if method == "brute": + best_dr, best_dc, best_ncc = _search_shifts_brute( + ref_s, mov_s, msk_s, max_shift_row, max_shift_col + ) + elif method == "fft": + best_dr, best_dc, best_ncc = _search_shifts_fft( + ref_s, mov_s, msk_s, max_shift_row, max_shift_col + ) + else: + raise ValueError( + f"unknown method {method!r}; expected 'auto', 'fft', or 'brute'" + ) + + valid = (mask_frac >= min_mask_frac) & (best_ncc > -1.0e8) + centers = (torch.arange(S, device=device, dtype=dtype) + 0.5) * hs + + return { + "centers": centers.detach().cpu().numpy(), + "drow": best_dr.detach().cpu().numpy().astype(np.float64), + "dcol": best_dc.detach().cpu().numpy().astype(np.float64), + "ncc": best_ncc.detach().cpu().numpy().astype(np.float64), + "valid": valid.detach().cpu().numpy().astype(bool), + "mask_frac": mask_frac.detach().cpu().numpy().astype(np.float64), + "n_strips": S, + "strip_height": hs, + "image_shape": (H, W), + "method": method, + } + + +def interpolate_residual_to_rows( + centers: np.ndarray, + drow: np.ndarray, + dcol: np.ndarray, + valid: np.ndarray, + n_rows: int, + *, + smooth_sigma_rows: float = 48.0, +) -> tuple[np.ndarray, np.ndarray]: + """Interpolate strip residuals to every scanline; Gaussian-smooth along slow axis. + + Always returns length ``n_rows``. (``np.convolve(..., mode="same")`` can grow + to the kernel length when the kernel is longer than the signal - e.g. σ=48 + on a 64-px test FOV - so we pad + ``mode="valid"`` instead.) + """ + if not np.any(valid): + return np.zeros(n_rows, dtype=np.float32), np.zeros(n_rows, dtype=np.float32) + c = centers[valid] + dr = drow[valid] + dc = dcol[valid] + rows = np.arange(n_rows, dtype=np.float64) + dr_i = np.interp(rows, c, dr, left=dr[0], right=dr[-1]) + dc_i = np.interp(rows, c, dc, left=dc[0], right=dc[-1]) + if smooth_sigma_rows and smooth_sigma_rows > 0: + rad = int(max(1, round(3 * smooth_sigma_rows))) + x = np.arange(-rad, rad + 1, dtype=np.float64) + ker = np.exp(-(x**2) / (2 * smooth_sigma_rows**2)) + ker /= ker.sum() + # edge-pad so valid convolution is exactly n_rows (not max(n, ker)) + dr_i = np.convolve(np.pad(dr_i, rad, mode="edge"), ker, mode="valid") + dc_i = np.convolve(np.pad(dc_i, rad, mode="edge"), ker, mode="valid") + if dr_i.shape[0] != n_rows or dc_i.shape[0] != n_rows: + raise RuntimeError( + f"interpolate_residual_to_rows length bug: n_rows={n_rows} " + f"drow={dr_i.shape[0]} dcol={dc_i.shape[0]}" + ) + return dr_i.astype(np.float32), dc_i.astype(np.float32) + + +def region_ncc( + reference: np.ndarray | torch.Tensor, + moving: np.ndarray | torch.Tensor, + mask: np.ndarray | torch.Tensor, + *, + device: str | torch.device | None = None, +) -> dict[str, float]: + """Common + top/middle/bottom masked NCC. + + Uses torch on CUDA when available (or when ``device`` is set); falls back + to numpy on CPU-only hosts. Same numeric definition either way. + """ + if device is None: + use_torch = torch.cuda.is_available() + device = torch.device("cuda") if use_torch else torch.device("cpu") + else: + device = torch.device(device) + use_torch = True + + if use_torch: + ref = _as_torch2d(reference, device, torch.float32) + mov = _as_torch2d(moving, device, torch.float32) + if isinstance(mask, torch.Tensor): + m = mask.to(device=device) + else: + m = torch.as_tensor(np.asarray(mask), device=device) + m = (m > 0.5).to(dtype=torch.float32) # 0/1 multiply-mask (no bool gather) + h = ref.shape[0] + + def _ncc_tensor(a, b, mm): + n = mm.sum().clamp_min(1.0) + xa = (a - (a * mm).sum() / n) * mm + yb = (b - (b * mm).sum() / n) * mm + num = (xa * yb).sum() + den = (xa.norm() * yb.norm()).clamp_min(1e-12) + return torch.where(n >= 64.0, num / den, torch.full((), float("nan"), device=a.device)) + + scores = [ + _ncc_tensor(ref, mov, m), + _ncc_tensor(ref, mov, m * torch.cat([ + torch.ones(h // 3, device=device), + torch.zeros(h - h // 3, device=device), + ])[:, None]), + _ncc_tensor(ref, mov, m * torch.cat([ + torch.zeros(h // 3, device=device), + torch.ones(h // 3, device=device), + torch.zeros(h - 2 * (h // 3), device=device), + ])[:, None]), + _ncc_tensor(ref, mov, m * torch.cat([ + torch.zeros(2 * (h // 3), device=device), + torch.ones(h - 2 * (h // 3), device=device), + ])[:, None]), + m.mean(), + ] + # single host sync for all diagnostics + vals = torch.stack([s.reshape(()) for s in scores]).detach().cpu().tolist() + return { + "common": float(vals[0]), + "top": float(vals[1]), + "middle": float(vals[2]), + "bottom": float(vals[3]), + "mask_frac": float(vals[4]), + } + + ref = np.asarray(reference, dtype=np.float64) + mov = np.asarray(moving, dtype=np.float64) + m = np.asarray(mask, dtype=bool) + h = ref.shape[0] + + def _ncc(a, b, mm): + if mm.sum() < 64: + return float("nan") + x, y = a[mm], b[mm] + x = x - x.mean() + y = y - y.mean() + return float(x @ y / max(np.linalg.norm(x) * np.linalg.norm(y), 1e-12)) + + out = {"common": _ncc(ref, mov, m)} + for name, (r0, r1) in zip( + ("top", "middle", "bottom"), + ((0, h // 3), (h // 3, 2 * h // 3), (2 * h // 3, h)), + ): + band = np.zeros_like(m) + band[r0:r1] = True + out[name] = _ncc(ref, mov, m & band) + out["mask_frac"] = float(m.mean()) + return out + + +def apply_row_residual_to_knots( + dc, + drow: np.ndarray, + dcol: np.ndarray, + weight: np.ndarray, + *, + moving_index: int = 1, +) -> None: + """Add per-scanline residual to knot coordinates in-place (topology unchanged). + + ``weight[r]`` scales the residual on scanline ``r`` (0 = leave knot as-is). + """ + k = dc.knots[moving_index] + if k.shape[1] != len(drow) or k.shape[1] != len(weight): + raise ValueError( + f"row residual length mismatch: knots H={k.shape[1]} " + f"drow={len(drow)} weight={len(weight)}" + ) + with torch.no_grad(): + dev, dtype = k.device, k.dtype + w = torch.as_tensor(weight, device=dev, dtype=dtype)[:, None] # (H, 1) + # K trailing dim: broadcast over knots along the scanline + k[0].add_(torch.as_tensor(drow, device=dev, dtype=dtype)[:, None] * w) + k[1].add_(torch.as_tensor(dcol, device=dev, dtype=dtype)[:, None] * w) + + +def correct_strip( + self, + *, + num_strips: int = 24, + max_row_shift: int = 8, + max_column_shift: int = 80, + smoothing_sigma: float = 12.0, + update_fraction: float = 1.0, + num_refine_cycles: int = 1, + passes: Sequence[StripPass] | None = None, + fixed_scans: list[int] | None = None, + correction_start_fraction: float | None = None, + ramp_fraction: float = 0.08, + min_overlap_fraction: float = 0.25, + show_combined: bool = True, + show_scans: bool = False, + show_knots: bool = True, + show_knot_plot: bool = False, + show_report: bool = False, + verbose: bool = True, +): + """Remove slow-scan-dependent residual drift after affine correction. + + Horizontal strips provide stable local displacement measurements when a + single affine rate cannot describe scan acceleration or bending. Each strip + contributes one masked-NCC ``(row, col)`` shift; interpolation and Gaussian + smoothing turn those measurements into a continuous scanline correction. + + Reference mode fixes image 0 automatically. Mutual multi-scan mode measures + each free scan against the leave-one-out mean of the other corrected scans. + Use :meth:`correct_nonrigid` when independent per-scanline motion is needed. + + Parameters + ---------- + num_strips : int, default 24 + Number of horizontal bands for independent rigid shifts. + max_column_shift, max_row_shift : int + Integer NCC search range (px) per band. Defaults are wide so + post-affine residual is not clipped; search is FFT-accelerated. + correction_start_fraction : float or None, default None + If set, only apply residual from this fraction of the FOV + downward (slow-scan) with a smooth ramp - useful when only the + bottom of a long scan still bends. ``None`` = full FOV. + smoothing_sigma : float, default 12.0 + Gaussian sigma (rows) when expanding strip shifts to scanlines. + num_refine_cycles : int, default 1 + Outer rebuilds of the corrected stack (re-measure residual after + applying knot deltas). Multipass helps large decelerating residuals. + update_fraction : float, default 1.0 + Multiplier applied to each measured row/column shift before it is + written into the knots. Values below one damp a pass, trading + convergence speed for stability. Used only by the scalar API. + passes : sequence of StripPass, optional + Heterogeneous coarse-to-fine recipe. Scalar pass settings must remain + at their defaults when this is supplied. + fixed_scans : list[int] or None + Images held fixed. Default ``None`` → auto ``[0]`` in reference + mode; otherwise leave-one-out mutual mode. + min_overlap_fraction : float, default 0.25 + Minimum valid-mask fraction required to measure a strip. + show_combined, show_scans, show_knots, show_knot_plot + Display knobs (same spirit as :meth:`correct_nonrigid`). + show_report : bool, default False + Print the before/affine/strip regional NCC table after alignment. + verbose : bool, default True + Show strip-pass progress and print per-pass region NCC deltas when + available. + + Returns + ------- + Self + For method chaining. + + Examples + -------- + >>> drift.correct_affine(show_combined=False) + >>> drift.correct_strip(num_strips=24, show_combined=True) + + """ + if not hasattr(self, "knots") or not hasattr(self, "_knots_after_affine"): + raise RuntimeError("correct_strip requires correct_affine() first.") + scalar_values = { + "num_strips": num_strips, + "max_row_shift": max_row_shift, + "max_column_shift": max_column_shift, + "smoothing_sigma": smoothing_sigma, + "update_fraction": update_fraction, + "num_refine_cycles": num_refine_cycles, + "correction_start_fraction": correction_start_fraction, + "ramp_fraction": ramp_fraction, + } + scalar_defaults = { + "num_strips": 24, + "max_row_shift": 8, + "max_column_shift": 80, + "smoothing_sigma": 12.0, + "update_fraction": 1.0, + "num_refine_cycles": 1, + "correction_start_fraction": None, + "ramp_fraction": 0.08, + } + if passes is None: + if num_refine_cycles < 1: + raise ValueError(f"num_refine_cycles must be >= 1, got {num_refine_cycles!r}.") + pass_configs = [ + StripPass( + num_strips=num_strips, + max_row_shift=max_row_shift, + max_column_shift=max_column_shift, + correction_start_fraction=correction_start_fraction, + ramp_fraction=ramp_fraction, + smoothing_sigma=smoothing_sigma, + update_fraction=update_fraction, + ) + for _ in range(num_refine_cycles) + ] + else: + mixed = [ + name for name, value in scalar_values.items() if value != scalar_defaults[name] + ] + if mixed: + raise ValueError( + "passes= cannot be combined with scalar pass settings: " + + ", ".join(mixed) + + ". Put those values on each StripPass instead." + ) + pass_configs = list(passes) + if not pass_configs: + raise ValueError("passes must contain at least one StripPass.") + if not all(isinstance(config, StripPass) for config in pass_configs): + raise TypeError( + "passes must contain only StripPass objects. Import with " + "`from quantem.imaging.drift import StripPass`." + ) + + # Same anchoring rule as correct_affine / correct_nonrigid. + if fixed_scans is None and self._reference_mode: + fixed_scans = [0] + fixed_set = frozenset(fixed_scans) if fixed_scans is not None else frozenset() + num_images = self.shape[0] + moving_indices = [i for i in range(num_images) if i not in fixed_set] + if not moving_indices: + raise ValueError( + "All images are fixed - nothing to strip-align. " + "fixed_scans must leave at least one moving image." + ) + if fixed_set and any(i < 0 or i >= num_images for i in fixed_set): + raise ValueError( + f"fixed_scans out of range for {num_images} images: {sorted(fixed_set)}" + ) + device = self._device + method = "auto" + + # Keep solve tensors on the active device (reload can leave CPU knots). + self.imgs_t = [t.to(self._device) for t in self.imgs_t] + self.knots = [k.to(self._device) for k in self.knots] + for attr in ("_knots_after_affine", "_initial_knots"): + snapshot = getattr(self, attr, None) + if snapshot is not None: + setattr(self, attr, [k.to(self._device) for k in snapshot]) + + direct_reference_mode = self._reference_mode and fixed_set == frozenset({0}) + reference_mask = ( + np.asarray(self.coverage_mask(), dtype=bool) if direct_reference_mode else None + ) + + strip_progress = tqdm( + enumerate(pass_configs), + total=len(pass_configs), + desc="Solving strip drift", + unit="pass", + disable=not verbose, + ) + for pass_idx, config in strip_progress: + # In reference mode, keep the external reference frame fixed and + # warp the moving alignment image directly, exactly as downstream + # spectrum/map channels are warped. Mutual mode still needs the + # co-registered leave-one-out canvas. + corrected = ( + warping.reference_scan_stack(self) + if direct_reference_mode + else warping.co_registered_scan_stack(self, fixed_set=fixed_set) + ) + shapes = {im.shape for im in corrected} + if len(shapes) != 1: + raise ValueError(f"correct_strip: corrected scans have mixed shapes {shapes}") + mask = ( + reference_mask + if reference_mask is not None + else np.asarray(self.coverage_mask(), dtype=bool) + ) + if mask.shape != corrected[0].shape: + raise ValueError( + f"coverage_mask shape {mask.shape} != scan shape {corrected[0].shape}" + ) + + # Measure all free scans against frozen refs, then apply together. + pending = [] + for i in moving_indices: + strip_progress.set_postfix_str( + f"strips={config.num_strips}, scan={i}", + refresh=verbose, + ) + if fixed_set: + ref = np.mean([corrected[j] for j in sorted(fixed_set)], axis=0).astype( + np.float32 + ) + else: + others = [corrected[j] for j in range(num_images) if j != i] + ref = np.mean(others, axis=0).astype(np.float32) + mov = corrected[i] + ncc_before = region_ncc(ref, mov, mask, device=device) + measured = measure_strip_residual_torch( + ref, + mov, + mask, + n_strips=config.num_strips, + max_shift_col=config.max_column_shift, + max_shift_row=config.max_row_shift, + device=device, + min_mask_frac=min_overlap_fraction, + method=method, + ) + H = ref.shape[0] + drow, dcol = interpolate_residual_to_rows( + measured["centers"], + measured["drow"], + measured["dcol"], + measured["valid"], + H, + smooth_sigma_rows=config.smoothing_sigma, + ) + weight = free_weight( + H, + free_from_frac=config.correction_start_fraction, + ramp_frac=config.ramp_fraction, + ) + drow = drow * config.update_fraction + dcol = dcol * config.update_fraction + pending.append( + ( + i, + drow, + dcol, + weight, + ncc_before, + ) + ) + + for i, drow, dcol, weight, _ in pending: + apply_row_residual_to_knots(self, drow, dcol, weight, moving_index=i) + + corrected_after = ( + warping.reference_scan_stack(self) + if direct_reference_mode + else warping.co_registered_scan_stack(self, fixed_set=fixed_set) + ) + for i, _, _, _, ncc_before in pending: + if fixed_set: + ref = np.mean( + [corrected[j] for j in sorted(fixed_set)], axis=0 + ).astype(np.float32) + else: + others = [corrected[j] for j in range(num_images) if j != i] + ref = np.mean(others, axis=0).astype(np.float32) + ncc_after = region_ncc( + ref, corrected_after[i], mask, device=device + ) + if verbose: + common_b = ncc_before.get("common", float("nan")) + common_a = ncc_after.get("common", float("nan")) + strip_progress.write( + f"correct_strip pass {pass_idx + 1}/{len(pass_configs)} " + f"image {i}: NCC common {common_b:.4f} → {common_a:.4f} " + f"(Δ={common_a - common_b:+.4f})" + ) + + self._images_warped_stale = True + # Snapshot knots after strip for optional stage plots (affine / strip / NR). + self._knots_after_strip = [k.detach().clone() for k in self.knots] + + drift_plot.show_after_step( + self, + "strip", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + if show_report: + print(self.report().to_string()) + return self diff --git a/src/quantem/imaging/drift/core/warping.py b/src/quantem/imaging/drift/core/warping.py index 8b741044..d8926426 100644 --- a/src/quantem/imaging/drift/core/warping.py +++ b/src/quantem/imaging/drift/core/warping.py @@ -1,9 +1,29 @@ -"""Cross-correlation and translation refinement for drift correction.""" +"""Warping primitives: cross-correlation translation + backward resampling. + +These are the fused torch operations used after the knot's forward warp has +put scan images on the same canvas. Forward warping itself lives in +``core.knots``; these functions run once the warped images need alignment and the learned drift +needs to be applied to raw data. + +Contains: + +* :func:`cross_corr_batch` + :func:`translate_align` — sub-pixel translation + alignment from FFT cross-correlation peaks (and the candidate scoring + signal for ``correct_affine``'s grid search). +* :func:`backward_warp` + :func:`backward_warp_grid_search` — bicubic + backward resampling that applies the learned drift to raw data, and + the affine candidate scoring loop that bypasses canvas KDE. +* Private DFT-upsample + parabolic-peak helpers used by the above. +""" import math +import numpy as np import torch from torch.fft import fft2, fftfreq, ifft2, ifftshift +from tqdm.auto import tqdm + +from quantem.imaging.drift.core import knots as drift_knots def cross_corr_batch( @@ -49,33 +69,17 @@ def cross_corr_batch( -------- >>> ref = torch.randn(5, 64, 64, dtype=torch.float64) >>> mov = torch.randn(5, 64, 64, dtype=torch.float64) - >>> cost = cross_corr_batch(ref, mov, 8, 32.0) + >>> cost = cross_corr_batch(ref, mov, 8) >>> cost.shape torch.Size([5]) """ - num_test_drifts, num_rows, num_cols = ref_images.shape + _, num_rows, num_cols = ref_images.shape dtype = ref_images.dtype mov_fft = fft2(mov_images) cross_corr_fft = fft2(ref_images) * mov_fft.conj() - cross_corr = ifft2(cross_corr_fft).real - # Reject correlation peaks beyond max shift to avoid locking onto - # periodic lattice repeats or noise peaks far from the true shift - if max_shift_mask is not None: - cross_corr.masked_fill_(max_shift_mask[None], 0.0) - # Find best-matching shift: integer peak → parabola (~0.1 px) → DFT (~0.01 px) - peak_flat_idx = cross_corr.flatten(1).argmax(dim=1) - peak_row = peak_flat_idx // num_cols - peak_col = peak_flat_idx % num_cols - batch_idx = torch.arange(num_test_drifts, device=ref_images.device) - refined_row, refined_col = _parabolic_peak_2d( - cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx + image_shifts = _translation_from_cross_correlation( + cross_corr_fft, upsample_factor, max_shift_mask ) - image_shifts = _dft_refine_shifts( - cross_corr_fft, refined_row, refined_col, upsample_factor - ) - # Wrap from [0, N) to [-N/2, N/2) so shifts represent actual displacement - image_shifts[:, 0] = ((image_shifts[:, 0] + num_rows / 2) % num_rows) - num_rows / 2 - image_shifts[:, 1] = ((image_shifts[:, 1] + num_cols / 2) % num_cols) - num_cols / 2 if freq_grids is not None: freq_row, freq_col = freq_grids else: @@ -89,15 +93,104 @@ def cross_corr_batch( return torch.mean(torch.abs(ref_images - aligned_images), dim=(1, 2)) +def fixed_overlap_ncc( + ref_images: torch.Tensor, + mov_images: torch.Tensor, + scan_shape: tuple[int, int], + max_image_shift: float | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Rank translations without letting padded pixels choose the lattice branch. + + A centered reference crop is compared with equally sized windows from the + moving scan. Every shift therefore uses the same number of measured pixels, + and normalized correlation removes intensity-scale differences. + """ + scan_rows, scan_cols = scan_shape + canvas_rows, canvas_cols = ref_images.shape[-2:] + row_start = (canvas_rows - scan_rows) // 2 + col_start = (canvas_cols - scan_cols) // 2 + ref = ref_images[ + :, row_start : row_start + scan_rows, col_start : col_start + scan_cols + ] + mov = mov_images[ + :, row_start : row_start + scan_rows, col_start : col_start + scan_cols + ] + + shift_limit = ( + min(scan_shape) / 4 if max_image_shift is None else float(max_image_shift) + ) + margin = min( + max(1, int(math.ceil(shift_limit))), + (min(scan_shape) - 2) // 2, + ) + template = ref[:, margin:-margin, margin:-margin] + template = template - template.mean(dim=(-2, -1), keepdim=True) + template_rows, template_cols = template.shape[-2:] + + fft_rows = 1 << (scan_rows + template_rows - 2).bit_length() + fft_cols = 1 << (scan_cols + template_cols - 2).bit_length() + numerator = torch.fft.irfft2( + torch.fft.rfft2(mov, s=(fft_rows, fft_cols)) + * torch.fft.rfft2( + template.flip((-2, -1)), s=(fft_rows, fft_cols) + ), + s=(fft_rows, fft_cols), + ) + numerator = numerator[ + :, + template_rows - 1 : scan_rows, + template_cols - 1 : scan_cols, + ] + + integral = torch.nn.functional.pad(mov, (1, 0, 1, 0)) + integral = integral.cumsum(-2).cumsum(-1) + integral_sq = torch.nn.functional.pad(mov.square(), (1, 0, 1, 0)) + integral_sq = integral_sq.cumsum(-2).cumsum(-1) + + def window_sum(table): + return ( + table[:, template_rows:, template_cols:] + - table[:, :-template_rows, template_cols:] + - table[:, template_rows:, :-template_cols] + + table[:, :-template_rows, :-template_cols] + ) + + moving_sum = window_sum(integral) + moving_sum_sq = window_sum(integral_sq) + pixels = float(template_rows * template_cols) + moving_norm = torch.sqrt( + (moving_sum_sq - moving_sum.square() / pixels).clamp_min(0.0) + ) + template_norm = template.norm(dim=(-2, -1), keepdim=True) + ncc = numerator / (template_norm * moving_norm).clamp_min(1e-12) + + shifts = torch.arange( + -margin, + margin + 1, + device=ncc.device, + dtype=ncc.dtype, + ) + allowed = shifts[:, None].square() + shifts[None, :].square() <= shift_limit**2 + ncc.masked_fill_(~allowed[None], -torch.inf) + flat_index = ncc.flatten(1).argmax(dim=1) + peak_row = flat_index // ncc.shape[-1] + peak_col = flat_index % ncc.shape[-1] + batch = torch.arange(ncc.shape[0], device=ncc.device) + best_ncc = ncc[batch, peak_row, peak_col] + image_shifts = -torch.stack( + (peak_row - margin, peak_col - margin), dim=1 + ).to(ncc.dtype) + return 1.0 - best_ncc, image_shifts, best_ncc - ncc[:, margin, margin] + + def translate_align( warped_images: torch.Tensor, upsample_factor: int, max_image_shift: float | None, - anchor_first: bool = False, ) -> torch.Tensor: """Pairwise translation alignment of warped images via cross-correlation. - Called by ``_warp_and_translate_torch`` after each affine warp to + Called by :func:`warp_and_translate` after each canvas warp to remove residual translational misalignment between the image pair. Without this step, the merged image would be blurred by the remaining translation offset even after the affine drift is corrected. @@ -114,9 +207,6 @@ def translate_align( Sub-pixel precision (1/N pixel) for DFT refinement. max_image_shift : float or None Maximum allowed shift in pixels. Peaks beyond this radius are masked. - anchor_first : bool, default False - Keep image 0 fixed and apply the full relative translation to later - images. Used when a reference image defines the output coordinates. Returns ------- @@ -140,24 +230,9 @@ def translate_align( for img_idx in range(1, num_images): mov_fft = fft2(warped_images[img_idx]) cross_corr_fft = ref_fft * mov_fft.conj() - cross_corr = ifft2(cross_corr_fft).real - if shift_mask is not None: - cross_corr.masked_fill_(shift_mask, 0.0) - # Find the integer peak, refine with parabola to ~0.1 px - peak_flat_idx = cross_corr.flatten().argmax() - peak_row = peak_flat_idx[None] // num_cols - peak_col = peak_flat_idx[None] % num_cols - batch_idx = torch.zeros(1, dtype=torch.long, device=device) - refined_row, refined_col = _parabolic_peak_2d( - cross_corr[None], peak_row, peak_col, num_rows, num_cols, batch_idx - ) - # Zoom into peak neighborhood with DFT to get ~0.01 px precision, - # then wrap from [0, N) to centered [-N/2, N/2) convention - refined_shift = _dft_refine_shifts( - cross_corr_fft[None], refined_row, refined_col, upsample_factor - ) - image_shifts[img_idx, 0] = ((refined_shift[0, 0] + num_rows / 2) % num_rows) - num_rows / 2 - image_shifts[img_idx, 1] = ((refined_shift[0, 1] + num_cols / 2) % num_cols) - num_cols / 2 + image_shifts[img_idx] = _translation_from_cross_correlation( + cross_corr_fft[None], upsample_factor, shift_mask + )[0] # Apply the recovered shift to current image via Fourier shift theorem, # then blend into running average so later images align to the cumulative mean phase = torch.exp( @@ -166,18 +241,552 @@ def translate_align( ) ) ref_fft = ref_fft * img_idx / (img_idx + 1) + mov_fft * phase / (img_idx + 1) - if not anchor_first: - # Mutual correction has no absolute reference frame. - image_shifts -= image_shifts.mean(dim=0) + # Remove mean so shifts are relative (no absolute reference frame) + image_shifts -= image_shifts.mean(dim=0) return image_shifts +def translate_align_pair_batch( + warped_pairs: torch.Tensor, + upsample_factor: int, + max_image_shift: float | None, +) -> torch.Tensor: + """Solve translations for a batch of independent two-image pairs. + + This is the candidate-batched equivalent of :func:`translate_align` for + the common two-scan case. It preserves the same FFT, parabolic peak, DFT + refinement, shift bounds, and zero-mean convention while avoiding a + Python-level full-canvas solve for every affine validation candidate. + + Parameters + ---------- + warped_pairs : torch.Tensor + Candidate image pairs, shape ``(N, 2, H, W)``. + upsample_factor : int + Sub-pixel precision (1/N pixel). + max_image_shift : float or None + Maximum allowed translational shift in pixels. + + Returns + ------- + torch.Tensor + Zero-mean shifts with shape ``(N, 2, 2)`` in ``(row, col)`` order. + """ + if warped_pairs.ndim != 4 or warped_pairs.shape[1] != 2: + raise ValueError( + "warped_pairs must have shape (N, 2, H, W); got " + f"{tuple(warped_pairs.shape)}." + ) + num_pairs, _, num_rows, num_cols = warped_pairs.shape + dtype = warped_pairs.dtype + device = warped_pairs.device + ref_fft = fft2(warped_pairs[:, 0]) + mov_fft = fft2(warped_pairs[:, 1]) + cross_corr_fft = ref_fft * mov_fft.conj() + shift_mask = None + if max_image_shift is not None: + dist_row = fftfreq( + num_rows, 1.0 / num_rows, device=device, dtype=dtype, + ) + dist_col = fftfreq( + num_cols, 1.0 / num_cols, device=device, dtype=dtype, + ) + shift_mask = ( + dist_row[:, None] ** 2 + dist_col[None, :] ** 2 + >= max_image_shift**2 + ) + pair_shift = _translation_from_cross_correlation( + cross_corr_fft, upsample_factor, shift_mask + ) + shifts = torch.zeros( + num_pairs, 2, 2, dtype=dtype, device=device, + ) + shifts[:, 1] = pair_shift + shifts -= shifts.mean(dim=1, keepdim=True) + return shifts + + +@torch.inference_mode() +def warp_and_translate( + correction, + max_image_shift: float | None, + upsample_factor: int = 8, + knots_batch: torch.Tensor | None = None, + solve_translation: bool = True, + fixed_indices: frozenset[int] | None = None, + imgs_t_override: list[torch.Tensor] | None = None, + return_weights: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Warp scans to their shared canvas and remove residual translation. + + Affine, strip, and non-rigid stages all need the same scientific operation: + render the current scan geometry, estimate the remaining rigid offset, move + the knots, and render once more. Keeping that operation here prevents the + affine search from owning general correction mechanics. + """ + device = correction._device + dtype = correction._dtype + num_images = correction.shape[0] + canvas_shape = (correction.shape[1], correction.shape[2]) + fixed_set = fixed_indices if fixed_indices else frozenset() + imgs_t = imgs_t_override if imgs_t_override is not None else correction.imgs_t + # AutoSerialize restores tensors on the host so one archive can be opened + # on CUDA, MPS, or CPU. Move the small alignment images at the shared warp + # boundary; callers should not need to repair a reloaded correction before + # requesting a report, coverage mask, or figure. + imgs_t = [image.to(device=device, dtype=dtype) for image in imgs_t] + if imgs_t_override is None: + correction.imgs_t = imgs_t + if knots_batch is None: + correction.knots = [ + knots.to(device=device, dtype=dtype) for knots in correction.knots + ] + else: + knots_batch = knots_batch.to(device=device, dtype=dtype) + + def render(warped_t, weights_t): + for img_idx in range(num_images): + knots_img = ( + knots_batch[img_idx].detach() if knots_batch is not None else None + ) + warped, image_weights = drift_knots.interpolator( + correction, img_idx, knots_img + ).warp_to_canvas( + imgs_t[img_idx], + canvas_shape, + correction.kde_sigma, + correction.pad_value[img_idx], + ) + warped_t[img_idx] = warped + weights_t[img_idx] = image_weights + + warped_t = torch.zeros(num_images, *canvas_shape, dtype=dtype, device=device) + weights_t = torch.zeros_like(warped_t) + render(warped_t, weights_t) + if not solve_translation: + if knots_batch is None: + correction.imgs_warped.array[:] = warped_t.cpu().numpy() + return (warped_t, weights_t) if return_weights else warped_t + + shifts_t = translate_align(warped_t, upsample_factor, max_image_shift) + if fixed_set: + fixed_idx_list = sorted(fixed_set) + anchor = shifts_t[fixed_idx_list].mean(0) + shifts_t -= anchor + for idx in fixed_set: + shifts_t[idx] = 0.0 + + if knots_batch is not None: + knots_batch[:, 0] += shifts_t[:, 0, None, None] + knots_batch[:, 1] += shifts_t[:, 1, None, None] + else: + for img_idx in range(num_images): + correction.knots[img_idx][0] += shifts_t[img_idx, 0] + correction.knots[img_idx][1] += shifts_t[img_idx, 1] + + render(warped_t, weights_t) + if knots_batch is None: + correction.imgs_warped.array[:] = warped_t.cpu().numpy() + return (warped_t, weights_t) if return_weights else warped_t + + +def backward_warp( + images: torch.Tensor, + drift: tuple[float, float] | torch.Tensor, + rigid_shift: tuple[float, float] = (0.0, 0.0), + mode: str = "bilinear", +) -> torch.Tensor: + """Apply drift correction via backward interpolation (``grid_sample``). + + Builds a per-scanline sampling grid that undoes the estimated drift + and optional rigid translation, then resamples with the chosen + interpolation kernel. + + Parameters + ---------- + images : torch.Tensor + Images to correct, shape ``(N, H, W)`` or ``(H, W)``. + For 4D-STEM, pass detector-pixel slices in chunks. + drift : tuple[float, float] | torch.Tensor + **Affine mode** — ``(row_slope, col_slope)`` scalar drift rate + in pixels per scan line, as returned by ``correct_affine``. + + **Tensor mode** — per-row ``(2, H)`` for K=1 (one shift per + scanline, broadcast across columns) or per-pixel ``(2, H, W)`` + for K>=2 (drift varies along the fast axis). + :class:`DriftKnot` returns the right shape per K, so + callers don't materialize the per-pixel tensor when the per-row + form suffices. *rigid_shift* is ignored in tensor mode. + rigid_shift : tuple[float, float], default (0.0, 0.0) + Global ``(row, col)`` translation to apply (affine mode only). + mode : str, default "bilinear" + Interpolation kernel passed to ``grid_sample``. + + Returns + ------- + torch.Tensor + Corrected images, same shape as *images*. + """ + squeeze = images.dim() == 2 + if squeeze: + images = images[None] + n, h, w = images.shape + device, dtype = images.device, images.dtype + + if isinstance(drift, torch.Tensor): + drift = drift.to(device=device, dtype=dtype) + if drift.shape == (2, h): + # Per-row (K=1): broadcast a single shift across columns. + row_shift = drift[0][:, None] + col_shift = drift[1][:, None] + elif drift.shape == (2, h, w): + # Per-pixel (K>=2): drift varies along fast axis. + row_shift = drift[0] + col_shift = drift[1] + else: + raise ValueError( + f"Drift must be (2, {h}) for K=1 or (2, {h}, {w}) for K>=2; " + f"got {tuple(drift.shape)}") + sample_r = torch.arange(h, device=device, dtype=dtype)[:, None].expand(-1, w) - row_shift + sample_c = torch.arange(w, device=device, dtype=dtype)[None, :].expand(h, -1) - col_shift + else: + # Affine mode: drift is (row_slope, col_slope) + offset = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2 + drift_row, drift_col = drift + shift_row, shift_col = rigid_shift + sample_r = ( + torch.arange(h, device=device, dtype=dtype)[:, None].expand(-1, w) + - drift_row * offset[:, None] + - shift_row + ) + sample_c = ( + torch.arange(w, device=device, dtype=dtype)[None, :].expand(h, -1) + - drift_col * offset[:, None] + - shift_col + ) + + grid_row = 2.0 * sample_r / (h - 1) - 1.0 + grid_col = 2.0 * sample_c / (w - 1) - 1.0 + # Pass as (1, N, H, W) with a single (1, H, W, 2) grid so grid_sample applies + # one grid to all N channels in one kernel call. The alternative (N, 1, H, W) + # with (N, H, W, 2) materialises N identical grids — 4096× more memory for EDS. + grid = torch.stack([grid_col, grid_row], dim=-1)[None] # (1, H, W, 2) — col first per grid_sample convention + + out = torch.nn.functional.grid_sample( + images[None], grid, mode=mode, + align_corners=True, padding_mode="border", + )[0] # (N, H, W) + return out[0] if squeeze else out + + +def canvas_center_to_scan(correction, canvas_image) -> np.ndarray: + """Crop a solver canvas to the acquired scan field of view.""" + array = ( + canvas_image.detach().cpu().numpy() + if isinstance(canvas_image, torch.Tensor) + else np.asarray(canvas_image) + ) + scan_rows, scan_columns = correction.imgs[0].shape[:2] + row = (array.shape[0] - scan_rows) // 2 + column = (array.shape[1] - scan_columns) // 2 + return np.ascontiguousarray( + array[row : row + scan_rows, column : column + scan_columns], + dtype=np.float32, + ) + + +def co_registered_scan_stack( + correction, + *, + fixed_set: frozenset[int], + max_image_shift: float | None = 32.0, + solve_translation: bool = True, + knots: list[torch.Tensor] | None = None, +) -> list[np.ndarray]: + """Warp scans into one scan-sized frame for residual measurements.""" + knots_batch = torch.stack( + [knot.detach() for knot in knots or correction.knots] + ).to(device=correction._device, dtype=correction._dtype) + warped = warp_and_translate( + correction, + max_image_shift=max_image_shift if solve_translation else None, + upsample_factor=8, + knots_batch=knots_batch, + solve_translation=solve_translation, + fixed_indices=fixed_set if fixed_set else None, + ) + if solve_translation: + with torch.no_grad(): + for index in range(len(correction.knots)): + correction.knots[index][...] = knots_batch[index] + return [canvas_center_to_scan(correction, image) for image in warped] + + +def reference_scan_stack( + correction, + knots: list[torch.Tensor] | None = None, +) -> list[np.ndarray]: + """Warp a moving scan while preserving its external reference frame.""" + reference = np.asarray(correction.imgs[0].array, dtype=np.float32) + moving = backward_warp( + correction.imgs_t[1], + drift=drift_knots.interpolator( + correction, 1, (knots or correction.knots)[1] + ).drift_raw(correction._initial_knots[1]), + mode="bilinear", + ).detach().cpu().numpy() + return [ + np.ascontiguousarray(reference), + np.ascontiguousarray(moving, dtype=np.float32), + ] + + +def knot_fingerprint(correction): + """Summarize current knots for validating the cached warped images.""" + parts = [] + for knot in correction.knots: + array = ( + knot.detach().cpu().numpy() + if hasattr(knot, "detach") + else np.asarray(knot) + ) + parts.append( + (array.shape, float(array.sum()), float(np.abs(array).max())) + ) + return tuple(parts) + + +def ensure_warped_images(correction): + """Keep the displayed warped stack synchronized with the solved knots.""" + fingerprint = knot_fingerprint(correction) + if ( + not getattr(correction, "_images_warped_stale", True) + and getattr(correction, "_warped_fingerprint", None) == fingerprint + ): + return + + if getattr(correction, "_reference_mode", False): + scans = np.stack(reference_scan_stack(correction)).astype( + np.float32, copy=False + ) + canvas_rows, canvas_columns = correction.shape[-2:] + scan_rows, scan_columns = scans.shape[-2:] + row = (canvas_rows - scan_rows) // 2 + column = (canvas_columns - scan_columns) // 2 + for index in range(scans.shape[0]): + correction.imgs_warped.array[index].fill( + float(correction.pad_value[index]) + ) + correction.imgs_warped.array[ + index, + row : row + scan_rows, + column : column + scan_columns, + ] = scans[index] + else: + warp_and_translate( + correction, + getattr(correction, "_max_image_shift_cached", None), + upsample_factor=8, + solve_translation=False, + ) + correction._images_warped_stale = False + correction._warped_fingerprint = knot_fingerprint(correction) + + + + +@torch.inference_mode() +def backward_warp_grid_search( + ref_image: torch.Tensor, + mov_image: torch.Tensor, + drift_vectors: torch.Tensor, + upsample_factor: int, + max_image_shift: float | None, + chunk_size: int | None = None, + progress_desc: str | None = None, +) -> tuple[int, torch.Tensor]: + """Score drift candidates by backward-warping the moving image. + + For each candidate ``(dr, dc)``, builds a sampling grid that undoes the + drift:: + + sample_row[i, j] = i - dr * (i - center) + sample_col[i, j] = j - dc * (i - center) + + then backward-warps the moving image with ``grid_sample`` (bicubic) and + scores alignment with the reference via ``cross_corr_batch``. + + This avoids forward-scatter KDE artifacts that bias the cost when only + one image's geometry changes (``fixed_indices`` mode). The periodic + wrapping in ``bilinear_kde_batch`` creates geometry-dependent seam + artifacts that differ between the fixed reference and the drift-shifted + moving image, making the MAE minimum diverge from the true drift. + Backward-warp scoring eliminates this by working at original resolution + without any canvas or KDE. + + Parameters + ---------- + ref_image : torch.Tensor + Reference image, shape ``(H, W)``. + mov_image : torch.Tensor + Moving image to correct, shape ``(H, W)``. + drift_vectors : torch.Tensor + Candidate drift rates, shape ``(N, 2)`` — columns ``(row_rate, col_rate)``. + upsample_factor : int + Sub-pixel precision for cross-correlation refinement. + max_image_shift : float or None + Maximum allowed translational shift in pixels. + chunk_size : int or None + Candidates per GPU pass. ``None`` auto-selects based on free memory. + progress_desc : str or None + Description for a progress bar shown when the search needs multiple + chunks. ``None`` disables progress reporting. + + Returns + ------- + tuple[int, torch.Tensor] + Index of the best candidate and full cost tensor of shape ``(N,)``. + """ + device = ref_image.device + dtype = ref_image.dtype + h, w = ref_image.shape + num_candidates = drift_vectors.shape[0] + center = (h - 1) / 2.0 + + rows = torch.arange(h, device=device, dtype=dtype) + cols = torch.arange(w, device=device, dtype=dtype) + offset = rows - center # (H,) + + shift_mask = None + if max_image_shift is not None: + dist_r = fftfreq(h, 1.0 / h, device=device, dtype=dtype) + dist_c = fftfreq(w, 1.0 / w, device=device, dtype=dtype) + shift_mask = dist_r[:, None] ** 2 + dist_c[None, :] ** 2 >= max_image_shift ** 2 + freq_grids = ( + fftfreq(h, device=device, dtype=dtype)[:, None], + fftfreq(w, device=device, dtype=dtype)[None, :], + ) + + if chunk_size is None: + if device.type == "cuda": + bytes_per_element = torch.finfo(dtype).bits // 8 + # grid (H*W*2) + warped (H*W) + FFT buffers (H*W*8*2) + per_cand_bytes = h * w * bytes_per_element * (2 + 1 + 16) + free_bytes, _ = torch.cuda.mem_get_info(device) + chunk_size = max(1, int(free_bytes * 0.4 / per_cand_bytes)) + chunk_size = min(chunk_size, num_candidates) + else: + chunk_size = num_candidates + + all_costs = [] + chunk_starts = range(0, num_candidates, chunk_size) + show_progress = progress_desc is not None and len(chunk_starts) > 1 + pbar = tqdm( + total=num_candidates, + desc=progress_desc, + unit="candidate", + disable=not show_progress, + ) + try: + for chunk_start in chunk_starts: + chunk_end = min(chunk_start + chunk_size, num_candidates) + n_chunk = chunk_end - chunk_start + drift_chunk = drift_vectors[chunk_start:chunk_end] + + drift_row_rate = drift_chunk[:, 0] # (n_chunk,) + drift_col_rate = drift_chunk[:, 1] # (n_chunk,) + + row_shift = drift_row_rate[:, None] * offset[None, :] # (n_chunk, H) + col_shift = drift_col_rate[:, None] * offset[None, :] # (n_chunk, H) + + # A candidate that moves most source pixels outside the detector + # footprint can look deceptively good because ``padding_mode`` + # repeats a nearly constant border. Reject rates with less than + # half-field geometric overlap before ranking their image cost. + sample_row_1d = rows[None] - row_shift + valid_rows = ( + (sample_row_1d >= 0.0) & (sample_row_1d <= h - 1) + ).to(dtype) + lower_col = torch.maximum( + torch.zeros_like(col_shift), col_shift, + ) + upper_col = torch.minimum( + torch.full_like(col_shift, w - 1), + (w - 1) + col_shift, + ) + valid_columns = torch.clamp( + upper_col - lower_col + 1.0, min=0.0, max=float(w), + ) + overlap_fraction = torch.sum( + valid_rows * valid_columns, dim=1, + ) / float(h * w) + + sample_rows = rows[None, :, None].expand(n_chunk, h, w) - row_shift[:, :, None] + sample_cols = cols[None, None, :].expand(n_chunk, h, w) - col_shift[:, :, None] + + grid_row = 2.0 * sample_rows / (h - 1) - 1.0 + grid_col = 2.0 * sample_cols / (w - 1) - 1.0 + grid = torch.stack( + [grid_col, grid_row], dim=-1 + ) # col first per grid_sample convention + + warped = torch.nn.functional.grid_sample( + mov_image[None, None].expand(n_chunk, 1, h, w), grid, + mode="bicubic", align_corners=True, padding_mode="border", + )[:, 0] + + ref_batch = ref_image[None].expand(n_chunk, -1, -1) + costs = cross_corr_batch( + ref_batch, warped, upsample_factor, + max_shift_mask=shift_mask, freq_grids=freq_grids, + ) + costs.masked_fill_(overlap_fraction < 0.5, torch.inf) + all_costs.append(costs) + pbar.update(n_chunk) + finally: + pbar.close() + + all_costs = torch.cat(all_costs) + return torch.argmin(all_costs).item(), all_costs + + +# --------------------------------------------------------------------------- +# Building blocks - used internally by the public API functions above +# --------------------------------------------------------------------------- + + +def _translation_from_cross_correlation( + cross_corr_fft: torch.Tensor, + upsample_factor: int, + max_shift_mask: torch.Tensor | None, +) -> torch.Tensor: + # Keep every solver on one integer → parabolic → DFT refinement and the + # same centered (row, col) shift convention. + num_images, num_rows, num_cols = cross_corr_fft.shape + cross_corr = ifft2(cross_corr_fft).real + if max_shift_mask is not None: + cross_corr.masked_fill_(max_shift_mask[None], 0.0) + peak_flat_idx = cross_corr.flatten(1).argmax(dim=1) + peak_row = peak_flat_idx // num_cols + peak_col = peak_flat_idx % num_cols + batch_idx = torch.arange(num_images, device=cross_corr_fft.device) + refined_row, refined_col = _parabolic_peak_2d( + cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx + ) + shifts = _dft_refine_shifts( + cross_corr_fft, refined_row, refined_col, upsample_factor + ) + shifts[:, 0] = ((shifts[:, 0] + num_rows / 2) % num_rows) - num_rows / 2 + shifts[:, 1] = ((shifts[:, 1] + num_cols / 2) % num_cols) - num_cols / 2 + return shifts + + def _dft_refine_shifts( - cross_corr_fft, - peak_row, - peak_col, - upsample_factor, -): + cross_corr_fft: torch.Tensor, + peak_row: torch.Tensor, + peak_col: torch.Tensor, + upsample_factor: int, +) -> torch.Tensor: """Refine coarse sub-pixel shifts using DFT upsampling + parabolic fit. After ``_parabolic_peak_2d`` gives a coarse sub-pixel position, this @@ -306,8 +915,19 @@ def _dft_upsample_batch( # (N,P,M) @ (N,M,K) @ (N,K,P) -> (N,P,P) return (kern_row @ cross_corr_fft @ kern_col).real +# --------------------------------------------------------------------------- +# Primitives - lowest-level operations +# --------------------------------------------------------------------------- -def _parabolic_peak_2d(cross_corr, peak_row, peak_col, num_rows, num_cols, batch_idx): + +def _parabolic_peak_2d( + cross_corr: torch.Tensor, + peak_row: torch.Tensor, + peak_col: torch.Tensor, + num_rows: int, + num_cols: int, + batch_idx: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: """Refine an integer cross-correlation peak to sub-pixel precision. Extracts the 3-point stencil along each axis and fits a parabola. @@ -345,7 +965,12 @@ def _parabolic_peak_2d(cross_corr, peak_row, peak_col, num_rows, num_cols, batch return refined_row, refined_col -def _parabolic_sub_pixel(val_m1, val_0, val_p1, mask=None): +def _parabolic_sub_pixel( + val_m1: torch.Tensor, + val_0: torch.Tensor, + val_p1: torch.Tensor, + mask: torch.Tensor | None = None, +) -> torch.Tensor: """Sub-pixel offset from a 3-point stencil via parabolic interpolation. Cross-correlation peaks fall on integer pixel positions, but the true diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index c0beb7e1..57e73084 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -1,379 +1,468 @@ -from collections.abc import Sequence +"""Orchestration for :class:`DriftCorrection`. + +This module owns the public workflow class. Numerical stages live in +``quantem.imaging.drift.core``, plotting in ``plot``, reporting in ``report``, +and 4D-STEM helpers in ``fourdstem``. Notebooks should import +``DriftCorrection`` from ``quantem.imaging`` or ``quantem.imaging.drift`` and +use ``from_emd`` / ``correct_affine`` / ``plot_combined`` / ``show`` / ``save``. +""" + from typing import Self import numpy as np import torch from numpy.typing import NDArray +import quantem.imaging.drift.apply as drift_apply +import quantem.imaging.drift.core.affine as affine +import quantem.imaging.drift.core.nonrigid as nonrigid +import quantem.imaging.drift.core.strip as strip +import quantem.imaging.drift.diagnostics as diagnostics +import quantem.imaging.drift.fourdstem as fourdstem +import quantem.imaging.drift.plot as drift_plot +import quantem.imaging.drift.preparation as preparation +import quantem.imaging.drift.report as drift_report from quantem.core.config import validate_device from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.dataset3d import Dataset3d from quantem.core.io.serialize import AutoSerialize -from quantem.core.utils.compound_validators import ( - validate_list_of_dataset2d, -) -from quantem.core.utils.validators import ensure_valid_array - -from . import apply as drift_apply -from . import fourdstem, preparation -from . import plot as drift_plot -from .core import affine, nonrigid class DriftCorrection(AutoSerialize): - """ - DriftCorrection provides translation, affine, and non-rigid drift correction for - sequential 2D images using scan direction metadata and flexible spatial interpolation. + """GPU-accelerated multi-angle scan drift correction. - This class supports input data as numpy arrays, Dataset2d, or Dataset3d instances, - with various padding strategies and configurable spline interpolation of scanline - trajectories via Bézier knot control. + Aligns scans at different scan directions, recovers per-scanline drift, + and forms a corrected product (HAADF pairs, reference EDS/EELS, 4D-STEM). - Features - -------- - - Load data from arrays or files - - Apply initial scanline resampling using Bézier curves - - Align images using translation, affine, or non-rigid optimization - - Visualize intermediate and final results with optional knot overlays - - Serialize state with `.save()` and restore with `.load()` - - Parameters (via `from_data` or `from_file`) - ------------------------------------------- - images : list of 2D arrays, Dataset2d, Dataset3d, or file names, or a 3D numpy array - The image stack to correct for drift. - scan_direction_degrees : list of float - The scan direction angle (in degrees) for each image, measured relative to vertical. - pad_fraction : float, default 0.25 - Fraction of padding to add around each image during interpolation. - pad_value : str, float, or list of float, default 'median' - How to pad outside the image area during warping. Can be: - - One of: 'median', 'mean', 'min', 'max' - - A float quantile value (e.g., 0.25) - - A list of per-image float values - number_knots : int, default 1 - Number of knots to use for Bézier interpolation of scanline trajectories. - We strongly recommend using `number_knots = 1` unless the fast scan direction is - expected to vary within the image. - - Example - ------- - Instantiate the DriftCorrection class, run preprocessing and alignment, and save/load results: - - >>> drift = DriftCorrection.from_data( - ... images=[ - ... image0, # 2D numpy array or Dataset2d - ... image1, - ... ], - ... scan_direction_degrees=[0, 90], - ... ).preprocess( - ... pad_fraction=0.25, - ... pad_value='median', - ... number_knots=1, - ... ) - - >>> drift.align_affine() - >>> drift.align_nonrigid() - >>> drift.plot_merged_images() - >>> image_corr = drift.generate_corrected_image() - - >>> drift.save("drift_result.zip") - >>> drift_reloaded = quantem.io.load("drift_result.zip") - - >>> image_corr.save("image_corrected.zip") - >>> image_corr_reloaded = quantem.io.load("image_corrected.zip") - - Notes - ----- - - Use `align_translation()` for rigid shifts, `align_affine()` for scan-shear or uniform drift, - and `align_nonrigid()` for flexible per-row or per-image correction. - - The class stores resampled images in `self.images_warped` and the control knots in `self.knots`. - - Visualization is supported through `plot_merged_images()` and `plot_transformed_images()`. - - Performance - ----------- - ``align_affine`` uses PyTorch to run all heavy operations on GPU - (works on CUDA, MPS, and CPU). The key optimizations are: - - - **Batched grid search**: all ~97 candidate drift vectors are warped - and scored in parallel, instead of one-at-a-time in a Python loop. - - **Batched bilinear KDE** (``core.knots.bilinear_kde_batch``): - scatter-based image warping via ``scatter_add_`` with int32 indices. - - **Batched FFT cross-correlation** (``core.warping.cross_corr_batch``): - sub-pixel alignment using DFT upsampling across all candidates at once. - - **Zero CPU round-trips**: coordinate transforms, Gaussian smoothing, - translation alignment, and error computation all stay on GPU until - the final sync. - - This gives ~300× speedup over the original NumPy implementation - (e.g. 436 s → 1.5 s on 2048×2048 image pairs). - - Memory is automatically chunked when the full batch doesn't fit. - Approximate memory per candidate at common sizes: - - ========== =========== ================ - Input size Canvas size Mem / candidate - ========== =========== ================ - 1024×1024 1280×1280 85 MB - 2048×2048 2560×2560 341 MB - 4096×4096 5120×5120 1.36 GB - ========== =========== ================ - """ - - _token = object() + Typical chain: :meth:`from_emd` → :meth:`correct_affine` → :meth:`plot_combined` + / :meth:`show` / :meth:`save`. Residual polish: :meth:`correct_strip`, + :meth:`correct_nonrigid` (use small ``max_image_shift`` on lattices). + """ def __init__( self, - images: list[Dataset2d], - scan_direction_degrees: NDArray, + *datasets: Dataset2d | Dataset3d | NDArray, + scan_direction_degrees: list[float] | NDArray | float | None = None, + alignment_image: NDArray | None = None, device: str | None = None, - _token: object | None = None, ): - if _token is not self._token: - raise RuntimeError( - "Use DriftCorrection.from_data() or .from_file() to instantiate this class." - ) - - self.images = images - self.scan_direction_degrees = ensure_valid_array(scan_direction_degrees, ndim=1) - self._reference_mode = False - + """Parameters + ---------- + *datasets : ndarray or Dataset2d or Dataset3d or Dataset4d + Two or more inputs (first = reference / first scan). + scan_direction_degrees : float or sequence of float, optional + Angle per dataset. Required for bare arrays; omit when Datasets + already have ``metadata["scan_rotation_deg"]``. + alignment_image : 2-D ndarray, optional + ≥3-D reference mode: 2-D partner for alignment (else auto virtual image). + device : str, optional + ``None`` → cuda/mps/cpu; ``"cpu"`` / ``"gpu"`` / ``"cuda:N"`` to pin. + """ + # Core state (always set so all code paths can rely on these). + self._datasets: list[np.ndarray | None] | None = None + self._datasets_consumed: bool = False + # Records that this instance was constructed from a heavy dataset + # (4D-STEM collection or reference mode). save() drops _datasets, so + # this flag is the only surviving signal that a reloaded instance + # needs its datasets re-attached before apply_correction() can run. + self._built_from_datasets: bool = False + self._normalized: bool = False + self._reference_mode: bool = False + # device=None auto-picks (cuda -> mps -> cpu). Pass "cuda:1" etc. to + # pin a specific GPU when one process drives several (HPC schedulers + # that set CUDA_VISIBLE_DEVICES per job need nothing here). device, _ = validate_device(device) self._device = device self._dtype = torch.float32 - @classmethod - def from_file( - cls, - file_paths: Sequence[str], - scan_direction_degrees: Sequence[float] | NDArray, - file_type: str | None = None, - ) -> Self: - image_list = [Dataset2d.from_file(fp, file_type=file_type) for fp in file_paths] - return cls.from_data( - image_list, - scan_direction_degrees, + if not datasets and alignment_image is None: + return + preparation.prepare_inputs( + self, datasets, scan_direction_degrees, alignment_image ) + @property + def device(self) -> str: + """Normalized compute device used by this correction. + + Returns ``"cuda:N"``, ``"mps"``, or ``"cpu"``. CUDA indices follow + PyTorch's visible-device numbering, including any remapping performed + by ``CUDA_VISIBLE_DEVICES``. + """ + return str(self._device) + @classmethod - def from_data( + def from_images( cls, - images: list[Dataset2d] | list[NDArray] | Dataset3d | NDArray, - scan_direction_degrees: list[float] | NDArray, + *images, + scan_direction_degrees: list[float] | NDArray | tuple[float, ...] | None = None, device: str | None = None, ) -> Self: - validated_images = validate_list_of_dataset2d(images) + """Create drift correction from already-loaded 2-D scans. - return cls( - images=validated_images, - scan_direction_degrees=scan_direction_degrees, - device=device, - _token=cls._token, - ) + Prefer :meth:`from_emd` for Velox files. Use this when data is already + in memory. Scan angles resolve as: + + 1. Explicit ``scan_direction_degrees`` (always wins). + 2. Each Dataset's ``metadata["scan_rotation_deg"]`` (from ``read_emd``). + 3. Neither → ``TypeError`` (bare arrays never get a silent default). + + Parameters + ---------- + *images + Two or more 2-D ``ndarray`` or :class:`Dataset2d` objects. + scan_direction_degrees : sequence of float or None, default None + One angle per image when metadata is missing. Omit when every + Dataset already carries ``scan_rotation_deg``. + device : str or None, default None + ``None`` auto-selects ``cuda`` → ``mps`` → ``cpu``. Use ``"cpu"`` to + force CPU, ``"gpu"`` to require a GPU, or ``"cuda:N"`` for multi-GPU. + + Returns + ------- + DriftCorrection + + Examples + -------- + >>> from quantem.imaging.drift.io import read_emd + >>> from quantem.gpu.device import profile + >>> device = profile()["device"] # or device="cpu" + >>> d0, d1 = read_emd(f0), read_emd(f90) + >>> dc = DriftCorrection.from_images(d0, d1, device=device) + + Bare arrays require explicit angles: + + >>> dc = DriftCorrection.from_images( + ... arr0, arr90, scan_direction_degrees=(0, 90), device="cpu") + """ + return cls(*images, scan_direction_degrees=scan_direction_degrees, device=device) @classmethod - def from_reference( + def from_emd(cls, *paths, verbose: bool = True, device: str | None = None) -> Self: + """Load Velox EMD files and build a :class:`DriftCorrection`. + + Each path is read with ``read_emd`` (scan angle from Velox metadata), then + handed to :meth:`from_images`. Never type angles for EMD files. Path order + does not matter. + + Parameters + ---------- + *paths : str or Path + Two or more EMD files with distinct scan angles. + verbose : bool, default True + Print shape, pixel size, and scan angle for each file. + device : str or None, default None + Compute device. ``None`` auto-selects ``cuda`` → ``mps`` → ``cpu``. + ``"gpu"`` requires CUDA/MPS; ``"cpu"`` forces CPU; ``"cuda:1"`` pins a + device in multi-GPU processes. + + Device helpers: ``from quantem.gpu.device import profile`` then + ``device=profile()["device"]``. + + Returns + ------- + DriftCorrection + Ready for :meth:`correct_affine`. + + Examples + -------- + >>> from quantem.gpu.device import profile + >>> device = profile()["device"] # or device="cpu" + >>> dc = DriftCorrection.from_emd( + ... "scan_0.emd", "scan_90.emd", device=device, verbose=False) + >>> dc.correct_affine(show_combined=False, verbose=False) + >>> dc.plot_combined(stage=("initial", "affine"), interactive=True) + >>> dc.save("drift.zip", mode="o") + """ + if len(paths) < 2: + raise TypeError(f"from_emd requires at least 2 EMD files, got {len(paths)}") + # Local import: drift.io imports quantem at module load, so importing it + # at correction module scope would be circular. + from quantem.imaging.drift.io import read_emd + + data = [read_emd(p) for p in paths] + if len(data) == 2: + data[0], data[1] = preparation.match_scan_shapes( + data[0], data[1], verbose=verbose + ) + if verbose: + for ds in data: + print( + f"{tuple(ds.shape)} {float(ds.sampling[0]) * 1e3:.2f} pm " + f"scan {ds.metadata['scan_rotation_deg']:.1f} deg (read from EMD)" + ) + return cls.from_images(*data, device=device) + + @classmethod + def from_4dstem( cls, - reference_image: Dataset2d | NDArray | Self, - alignment_image: Dataset2d | NDArray, - *, - scan_direction_degrees: float | None = None, + *datasets, + scan_direction_degrees: list[float] | NDArray | tuple[float, ...] = (0.0, 90.0), + scan_sampling: float | tuple[float, float] | None = None, + scan_units: str | tuple[str, str] = "pixels", device: str | None = None, ) -> Self: - """Fit a fixed-reference drift field from two-dimensional images. + """Build a 0°/90° 4D-STEM *collection* correction (two drifted scans). - The two-dimensional ``reference_image`` defines the output frame. - ``alignment_image`` is the HAADF image acquired with a spectrum image. - Pass the spectrum image to :meth:`apply_correction` after fitting. + Alignment uses auto-extracted virtual images; apply fields with + :meth:`corrected_4dstem`. Distinct from :meth:`from_reference` (fixed 2-D + reference + one multi-D target). Currently requires exactly two orthogonal + 4D-STEM datasets. Parameters ---------- - reference_image : Dataset2d, numpy.ndarray, or DriftCorrection - Fixed two-dimensional reference image. A solved correction is - converted to its corrected native-size image automatically. - alignment_image : Dataset2d or numpy.ndarray - HAADF image acquired with the spectrum image. - scan_direction_degrees : float or None, default None - Recorded scan rotation. Omit when ``alignment_image`` carries - ``metadata["scan_rotation_deg"]``. + *datasets + Two 4D-STEM datasets with leading scan axes + ``(scan_row, scan_col, det_row, det_col)``. + scan_direction_degrees : sequence of float, default (0.0, 90.0) + Scan direction per dataset in degrees. + scan_sampling : float or 2-tuple of float, optional + Real-space scan sampling. Scalar applies to both axes. ``None`` leaves + uncalibrated sampling. + scan_units : str or 2-tuple of str, default "pixels" + Units for ``scan_sampling`` (for example ``"nm"``). + device : str or None, default None + ``None`` → cuda/mps/cpu. ``"cpu"`` / ``"gpu"`` / ``"cuda:N"`` as in + :meth:`from_emd`. Returns ------- DriftCorrection - Correction ready for ``preprocess`` and ``align_affine``. Examples -------- - >>> drift = DriftCorrection.from_reference( - ... reference, - ... haadf, - ... scan_direction_degrees=0.0, + >>> from quantem.gpu.device import profile + >>> device = profile()["device"] # or "cpu" + >>> dc = DriftCorrection.from_4dstem( + ... data_0, data_1, + ... scan_direction_degrees=(0.0, 90.0), + ... scan_sampling=0.05, + ... scan_units="nm", + ... device=device, ... ) - >>> drift.preprocess(show_merged=False, show_images=False) - >>> drift.align_affine(show_merged=False, show_images=False) - >>> corrected = drift.apply_correction(spectrum_image) + >>> dc.correct_affine(show_combined=False, verbose=False) + >>> result = dc.corrected_4dstem(merge=True, verbose=True) """ - if isinstance(reference_image, cls): - reference_image = reference_image.generate_corrected( - upsample_factor=1, - strip_padding=True, - mask_output=False, - show_merged=False, - ) - reference_array = ( - reference_image.array - if isinstance(reference_image, Dataset2d) - else ensure_valid_array(reference_image, ndim=2) - ) - alignment_array = ( - alignment_image.array - if isinstance(alignment_image, Dataset2d) - else ensure_valid_array(alignment_image, ndim=2) - ) - if reference_array.shape != alignment_array.shape: - raise ValueError( - "reference_image and alignment_image must have the same " - f"shape; got {reference_array.shape} and {alignment_array.shape}." - ) - if scan_direction_degrees is None: - metadata = getattr(alignment_image, "metadata", {}) - scan_direction_degrees = metadata.get("scan_rotation_deg") - if scan_direction_degrees is None: - raise TypeError( - "scan_direction_degrees is required when alignment_image " - "does not carry scan_rotation_deg metadata." - ) - result = cls.from_data( - images=[reference_image, alignment_image], - scan_direction_degrees=[scan_direction_degrees, scan_direction_degrees], + result = cls( + *datasets, + scan_direction_degrees=scan_direction_degrees, device=device, ) - result._reference_mode = True + if scan_sampling is not None: + units = (scan_units, scan_units) if isinstance(scan_units, str) else scan_units + for image in result.imgs: + image.sampling = scan_sampling + image.units = units return result @classmethod - def from_4dstem( + def from_reference( cls, - *datasets, - scan_direction_degrees: Sequence[float] | NDArray | None = None, - scan_sampling: float | tuple[float, float] | None = None, - scan_units: str | tuple[str, str] = "pixels", + reference_image, + drifted_dataset, + *, + alignment_image: NDArray | None = None, + scan_direction_degrees: list[float] | NDArray | float = 0.0, device: str | None = None, ) -> Self: - """Fit one correction from virtual images of two 4D-STEM scans. + """Build a *reference-anchored* correction (fixed frame + drifted data). + + ``reference_image`` is a 2-D HAADF (or a solved :class:`DriftCorrection` + whose merge becomes the frame). ``drifted_dataset`` may be 2-D, 3-D + (EDS/EELS), or 4-D STEM. Alignment uses ``alignment_image`` or an auto + virtual image. Image 0 is fixed automatically. + + Parameters + ---------- + reference_image + Fixed 2-D reference, or a solved :class:`DriftCorrection` used as a + provenance-chained reference. + drifted_dataset + Drifted 2-D image, 3-D spectrum image, or 4-D STEM dataset. + alignment_image : ndarray, optional + 2-D image used to estimate target drift (for example session HAADF). + If omitted, a virtual image is extracted automatically. + scan_direction_degrees : float or sequence of float, default 0.0 + Target scan direction(s) in degrees. + device : str or None, default None + ``None`` auto-selects device or inherits from a chained reference. + ``"cpu"`` forces CPU; ``"gpu"`` requires GPU. + + Returns + ------- + DriftCorrection - Only scan axes are transformed when the fitted fields are later - applied. Detector coordinates and native detector sampling are never - rotated, resampled, or binned by this constructor. + Examples + -------- + >>> from quantem.gpu.device import profile + >>> device = profile()["device"] + >>> dc = DriftCorrection.from_reference( + ... haadf_ref, + ... eds_cube, + ... alignment_image=eds_session_haadf, + ... device=device, + ... ) + >>> dc.correct_affine(show_combined=False) + >>> eds_corrected = dc.corrected() # Dataset3d + + Chained reference from a prior pair solve: + + >>> dc_pair = DriftCorrection.from_emd(f0, f90).correct_affine() + >>> dc_eds = DriftCorrection.from_reference( + ... dc_pair, eds_cube, alignment_image=eds_haadf) """ - if len(datasets) != 2: - raise ValueError( - f"from_4dstem requires exactly two datasets, got {len(datasets)}." + reference_dc = None + if isinstance( + reference_image, DriftCorrection + ) or AutoSerialize._is_autoserialize_instance(reference_image): + # Chained reference: derive the fixed frame from the solved pair's + # corrected merge, cropped to the drifted scan field of view. + reference_dc = reference_image + if device is None: + device = reference_dc.device + drifted_shape = preparation.input_array(drifted_dataset).shape[:2] + reference_downsample = preparation.reference_downsample( + tuple(int(value) for value in reference_dc.imgs[0].shape[:2]), + tuple(int(value) for value in drifted_shape), + reference_sampling=getattr(reference_dc.imgs[0], "sampling", None), + target_sampling=getattr(drifted_dataset, "sampling", None), ) - raw_datasets = [fourdstem.data_array(dataset) for dataset in datasets] - if any(data.ndim != 4 for data in raw_datasets): - raise ValueError( - "from_4dstem expects scan-axis-leading 4-D inputs; got " - f"{[tuple(data.shape) for data in raw_datasets]}." - ) - detector_shapes = [tuple(data.shape[2:]) for data in raw_datasets] - if detector_shapes[0] != detector_shapes[1]: - raise ValueError( - "4D-STEM inputs must share one detector shape; got " - f"{detector_shapes}." - ) - if scan_direction_degrees is None: - angles = [ - getattr(dataset, "metadata", {}).get("scan_rotation_deg") - for dataset in datasets - ] - if any(angle is None for angle in angles): - raise TypeError( - "scan_direction_degrees is required when 4D-STEM inputs " - "do not carry scan_rotation_deg metadata." + if reference_downsample > 1: + # Solve the reference pair on the target acquisition grid. On + # atomic images, aligning at the finer grid and averaging only + # the final merge can preserve a different lattice phase than + # the coarser EDS scan. Re-solving the already-small virtual + # images avoids that alias while keeping the operation fully + # automatic and costs well under a second on the microscope GPU. + scaled_reference_images = [ + preparation.average_downsample_2d( + np.asarray(image.array), + reference_downsample, + ) + for image in reference_dc.imgs + ] + reference_dc = cls.from_images( + *scaled_reference_images, + scan_direction_degrees=tuple(reference_dc.scan_direction_degrees), + device=device, ) - scan_direction_degrees = angles - virtual_images = [] - for index, (dataset, raw_dataset) in enumerate( - zip(datasets, raw_datasets, strict=True) - ): - image = Dataset2d.from_array( - fourdstem.integrate_virtual_detector(raw_dataset), - name=f"4D-STEM virtual image {index}", - ) - if scan_sampling is None and hasattr(dataset, "sampling"): - image.sampling = np.asarray(dataset.sampling[:2], dtype=float) - if isinstance(scan_units, str) and hasattr(dataset, "units"): - image.units = list(dataset.units[:2]) - image.metadata.update(dict(getattr(dataset, "metadata", {}))) - virtual_images.append(image) - if scan_sampling is not None: - sampling = ( - (float(scan_sampling), float(scan_sampling)) - if np.isscalar(scan_sampling) - else tuple(float(value) for value in scan_sampling) - ) - units = ( - (scan_units, scan_units) - if isinstance(scan_units, str) - else tuple(scan_units) + reference_dc.correct_affine( + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + reference_image = preparation.match_reference_image( + reference_dc.corrected().array, + tuple(int(value) for value in reference_dc.imgs[0].shape[:2]), + tuple(int(value) for value in drifted_shape), ) - for image in virtual_images: - image.sampling = sampling - image.units = units - result = cls.from_data( - virtual_images, + reference_image = reference_image.astype(np.float32, copy=False) + result = cls( + reference_image, + drifted_dataset, scan_direction_degrees=scan_direction_degrees, + alignment_image=alignment_image, device=device, ) - result._datasets = raw_datasets - result._datasets_consumed = False + # The corrected product must carry the drifted dataset's calibration + # (sampling, units, metadata). Without this, reference-mode corrected() + # falls back to {} and the EDS/4D output ships uncalibrated. + result._reference_dataset_info = drift_apply.dataset_info(drifted_dataset) + if not result._reference_mode: + raise ValueError( + "from_reference() requires a 2-D reference image and one " + "drifted target dataset. For a 2-D target, pass a scalar " + "scan_direction_degrees value or matching reference/target " + "scan directions so the call is unambiguously single-sided. " + "Use DriftCorrection.from_4dstem() for 0/90 4D-STEM " + "collection correction." + ) return result - @property - def imgs(self): - """Compatibility alias for the scan images used by the solver.""" - return self.images - - @property - def device(self) -> str: - """Compute device selected for drift fitting and data propagation.""" - return str(self._device) + drift_field = fourdstem.drift_field + probe_positions = fourdstem.probe_positions preprocess = preparation.preprocess - align_translation = preparation.align_translation - align_affine = affine.align_affine - _affine_grid_search_batch = affine._affine_grid_search_batch - _auto_chunk_size = staticmethod(affine._auto_chunk_size) - _warp_and_translate_torch = affine._warp_and_translate_torch - align_nonrigid = nonrigid.align_nonrigid - _optimize_knots_adam = nonrigid._optimize_knots_adam - _compiled_loss_fn = staticmethod(nonrigid._compiled_loss_fn) - _optimize_knots_lbfgs = nonrigid._optimize_knots_lbfgs - _regularize_knots = nonrigid._regularize_knots - _optimize_knots_scipy = nonrigid._optimize_knots_scipy - generate_corrected = drift_apply.generate_corrected - generate_corrected_image = drift_apply.generate_corrected_image + correct_affine = affine.correct_affine + + correct_strip = strip.correct_strip + + report = drift_report.report + + correct_nonrigid = nonrigid.correct_nonrigid + + diagnose_affine = diagnostics.diagnose_affine + diagnose_nonrigid = diagnostics.diagnose_nonrigid + + corrected = drift_apply.corrected apply_correction = drift_apply.apply_correction - drift_field = fourdstem.drift_field - probe_positions = fourdstem.probe_positions + + crop = drift_apply.crop + crop_slices = drift_apply.crop_slices + coverage_mask = drift_apply.coverage_mask + show = drift_plot.show + show_4dstem = drift_plot.show_4dstem + corrected_virtual_images = fourdstem.corrected_virtual_images regional_diffraction_patterns = fourdstem.regional_diffraction_patterns corrected_4dstem = fourdstem.corrected_4dstem + integrate_virtual_detector = staticmethod(fourdstem.integrate_virtual_detector) + # -- serialization ------------------------------------------------------- + def save(self, path, mode="w", store="auto", skip=(), compression_level=4): - """Serialize fitted state without embedding raw multidimensional data.""" + """Save a solved correction for later analysis and figure rendering. + + Persists knots and 2-D alignment state so figure notebooks can ``load`` + without re-solving. Large 4D-STEM cubes under ``_datasets`` are always + skipped. + + Parameters + ---------- + path : str or Path + Output path (typically ``drift.zip`` next to the raw data). + mode : str, default "w" + File mode for AutoSerialize. Tutorials often use ``"o"`` to overwrite. + store : str, default "auto" + Storage backend selection. + skip : sequence, default () + Extra attributes to omit. ``_datasets`` is always appended. + compression_level : int, default 4 + Archive compression level. + + Returns + ------- + None + + Examples + -------- + >>> dc = DriftCorrection.from_emd(f0, f90) + >>> dc.correct_affine(show_combined=False, verbose=False) + >>> dc.save("data/sample/drift.zip", mode="o") + >>> from quantem.core.io import load + >>> dc2 = load("data/sample/drift.zip") + + """ if isinstance(skip, (str, type)): skip = [skip] + skip = list(skip) + ["_datasets"] super().save( path, mode=mode, store=store, - skip=[*skip, "_datasets"], + skip=skip, compression_level=compression_level, ) - calculate_error = drift_apply.calculate_error - plot_transformed_images = drift_plot.plot_transformed_images + + drift_rate = property(affine.drift_rate) + + # -- visualization methods bound directly from drift_plot so hover + # shows the real signature + docstring (no `**kw` indirection). + plot_warped_images = drift_plot.plot_warped_images plot_convergence = drift_plot.plot_convergence - _ensure_warped_images = drift_plot._ensure_warped_images - plot_merged_images = drift_plot.plot_merged_images + # Primary registration QA plot (tutorials + paper combined wording). + plot_combined = drift_plot.plot_combined + plot_knots = drift_plot.plot_knots + plot_probe_positions = drift_plot.plot_probe_positions diff --git a/src/quantem/imaging/drift/diagnostics.py b/src/quantem/imaging/drift/diagnostics.py index b4d52343..0fbd7ea2 100644 --- a/src/quantem/imaging/drift/diagnostics.py +++ b/src/quantem/imaging/drift/diagnostics.py @@ -1,276 +1,474 @@ -"""Pure numerical diagnostics for drift-correction results. +"""Scientific diagnostics for affine and non-rigid drift correction.""" -This module measures existing correction checkpoints. It does not run an -optimizer, update knots, or replace the current warped-image cache. Public API -decisions are intentionally deferred to the final API integration change. -""" - -from collections.abc import Sequence +import copy +import time +import matplotlib.pyplot as plt import numpy as np +import pandas as pd import torch -from numpy.typing import NDArray - -_COVERAGE_THRESHOLD = 1e-3 -_MIN_CORRELATION_PIXELS = 64 - - -def _copy_knots(knots: Sequence[NDArray | torch.Tensor]) -> tuple[NDArray, ...]: - """Copy knot arrays to CPU without retaining mutable solver storage.""" - copied = [] - for knot in knots: - if isinstance(knot, torch.Tensor): - knot = knot.detach().cpu().numpy() - copied.append(np.asarray(knot).copy()) - return tuple(copied) - - -def _record_stage(correction, stage: str) -> None: - """Record one lightweight knot checkpoint after a completed stage.""" - if not hasattr(correction, "knots"): - raise RuntimeError( - "A diagnostic checkpoint requires prepared knots. Run preprocess() first." +from matplotlib.figure import Figure +from tqdm import tqdm + +import quantem.imaging.drift.apply as drift_apply +from quantem.imaging.drift.core import strip +from quantem.imaging.drift.core.warping import canvas_center_to_scan +from quantem.imaging.drift.plot import overlay_pair + + +def _diagnostic_grid(num_rows: int): + """Keep square diagnostic panels close enough for direct comparison.""" + figure, axes = plt.subplots( + num_rows, + 4, + figsize=(13.1, 3.8 * num_rows), + squeeze=False, + ) + figure.subplots_adjust( + left=0.075, + right=0.995, + bottom=0.02, + top=0.91, + wspace=0.015, + hspace=0.18, + ) + return figure, axes + + +def _copy_for_diagnosis(self): + """Isolate exploratory solves and rendering from the live correction.""" + memo = {} + if getattr(self, "_datasets", None) is not None: + # Diagnostics use the resident 2-D alignment images, not the potentially + # multi-gigabyte EDS or 4D-STEM source payload. + memo[id(self._datasets)] = None + return copy.deepcopy(self, memo) + + +def diagnose_affine( + self, + *, + stage: str | None = "affine", + grid_shape: tuple[int, int] = (2, 2), + smoothing_sigma: float | None = None, +) -> tuple[Figure, pd.DataFrame]: + """Show where one regional affine model leaves non-rigid disagreement + + Periodic atomic images can achieve a high global score even when defects + do not correspond everywhere. This view divides the corrected field into + regions so scientists can verify the trusted affine fit and see where + spatially varying drift remains for non-rigid correction. + + The correction is never modified. Each row represents one physical image + region. Four columns show corrected scan 0, corrected scan 1, their actual + RGB registration, and the absolute scan difference. No additional shift + is searched or applied. + + Parameters + ---------- + stage : {"initial", "affine", "strip", "nonrigid", None}, optional + Correction stage to diagnose. Default is ``"affine"``. ``None`` uses + the current knots. + grid_shape : tuple of int, optional + Number of regional rows and columns. Default is ``(2, 2)``. + smoothing_sigma : float or None, optional + Display and measurement smoothing in pixels. ``None`` uses the current + correction value. Use 0.5 to validate near-native atomic detail after + searching with stronger smoothing. + + Returns + ------- + matplotlib.figure.Figure + Four-column regional comparison of both scans, their actual RGB + registration, and their absolute difference. + pandas.DataFrame + Regional role and bounds, current NCC, yellow agreement, and mean + absolute difference. + + Examples + -------- + >>> drift.correct_affine(show_combined=False) + >>> figure, regions = drift.diagnose_affine() + >>> regions[["region", "current_ncc", "current_yellow"]] + """ + grid_rows, grid_columns = (int(value) for value in grid_shape) + if grid_rows < 1 or grid_columns < 1: + raise ValueError(f"grid_shape must contain positive counts, got {grid_shape}") + candidate = _copy_for_diagnosis(self) + if smoothing_sigma is not None: + candidate.kde_sigma = float(smoothing_sigma) + candidate._images_warped_stale = True + stack = drift_apply.warped_stack(candidate, stage) + if stack.shape[0] != 2: + raise ValueError( + "diagnose_affine() currently compares one image pair; " + f"this correction contains {stack.shape[0]} images" ) - if not hasattr(correction, "_diagnostic_knots"): - correction._diagnostic_knots = {} - stage = str(stage) - recorded = list(correction._diagnostic_knots) - if stage in recorded: - for stale_stage in recorded[recorded.index(stage) :]: - del correction._diagnostic_knots[stale_stage] - correction._diagnostic_knots[stage] = _copy_knots(correction.knots) - - -def _available_stages(correction) -> tuple[str, ...]: - """Return recorded correction stages in completion order.""" - snapshots = getattr(correction, "_diagnostic_knots", {}) - return tuple(snapshots) + reference, moving = [canvas_center_to_scan(candidate, image) for image in stack] + + # ----------------------------------------------------------------------- + # Divide the actual corrected field into physical regions + # ----------------------------------------------------------------------- + tile_rows = reference.shape[0] // grid_rows + tile_columns = reference.shape[1] // grid_columns + reference_tiles = [] + moving_tiles = [] + regions = [] + search = getattr(self, "affine_search_info", {}) + trusted_bounds = search.get("trusted_region_bounds_row_column") + vertical = ("top", "bottom") if grid_rows == 2 else () + horizontal = ("left", "right") if grid_columns == 2 else () + for row in range(grid_rows): + for column in range(grid_columns): + row_slice = slice(row * tile_rows, (row + 1) * tile_rows) + column_slice = slice( + column * tile_columns, + (column + 1) * tile_columns, + ) + reference_tiles.append(reference[row_slice, column_slice]) + moving_tiles.append(moving[row_slice, column_slice]) + name = ( + f"{vertical[row]} {horizontal[column]}" + if vertical and horizontal + else f"row {row}, column {column}" + ) + trusted = ( + search.get("strategy") == "trusted_region" + and trusted_bounds is not None + and row_slice.start >= trusted_bounds[0] + and row_slice.stop <= trusted_bounds[1] + and column_slice.start >= trusted_bounds[2] + and column_slice.stop <= trusted_bounds[3] + ) + regions.append( + ( + name, + row, + column, + row_slice, + column_slice, + "trusted affine fit" if trusted else "validation region", + ) + ) + # ----------------------------------------------------------------------- + # Plot only the delivered result and summarize the visible agreement + # ----------------------------------------------------------------------- + figure, axes = _diagnostic_grid(len(regions)) + records = [] + stage_label = "current final" if stage is None else str(stage) + directions = np.asarray(self.scan_direction_degrees, dtype=float) + for index, (name, row, column, row_slice, column_slice, role) in enumerate( + regions + ): + reference_tile = reference_tiles[index] + moving_tile = moving_tiles[index] + reference_display = reference_tile + moving_display = moving_tile + current_overlay = overlay_pair(reference_display, moving_display) + + red, green = current_overlay[..., 0], current_overlay[..., 1] + current_yellow = float( + np.minimum(red, green).sum() + / max(float(np.maximum(red, green).sum()), 1e-12) + ) + finite = np.isfinite(reference_display) & np.isfinite(moving_display) + first = reference_display[finite].astype(np.float64) + second = moving_display[finite].astype(np.float64) + first -= first.mean() + second -= second.mean() + current_ncc = float( + first @ second + / max(np.linalg.norm(first) * np.linalg.norm(second), 1e-12) + ) + records.append( + { + "region": name, + "region_role": role, + "row_start": row_slice.start, + "row_stop": row_slice.stop, + "column_start": column_slice.start, + "column_stop": column_slice.stop, + "grid_row": row, + "grid_column": column, + "current_ncc": current_ncc, + "current_yellow": current_yellow, + "mean_absolute_difference": float( + np.abs(reference_display[finite] - moving_display[finite]).mean() + ), + } + ) + contrast = np.concatenate( + (reference_display.ravel()[::16], moving_display.ravel()[::16]) + ) + low, high = np.percentile(contrast[np.isfinite(contrast)], (1, 99)) + scale = max(float(high - low), np.finfo(np.float32).eps) + difference = np.abs( + np.clip((reference_display - low) / scale, 0, 1) + - np.clip((moving_display - low) / scale, 0, 1) + ) + axes[index, 0].imshow( + reference_display, + cmap="gray", + vmin=low, + vmax=high, + ) + axes[index, 0].set_title( + f"Scan 0 ({directions[0]:g}°)\nactual {stage_label} image" + ) + axes[index, 1].imshow( + moving_display, + cmap="gray", + vmin=low, + vmax=high, + ) + axes[index, 1].set_title( + f"Scan 1 ({directions[1]:g}°)\nactual {stage_label} image" + ) + axes[index, 2].imshow(current_overlay) + axes[index, 2].set_title( + f"Actual {stage_label} RGB\n" + f"NCC {current_ncc:.3f}, yellow {current_yellow:.1%}" + ) + axes[index, 3].imshow(difference, cmap="magma", vmin=0, vmax=1) + axes[index, 3].set_title("Residual difference |scan 0 - scan 1|") + axes[index, 0].set_ylabel( + f"{name}\nrows {row_slice.start}:{row_slice.stop}\n" + f"columns {column_slice.start}:{column_slice.stop}\n{role}", + fontsize=11, + ) + for axis in axes[index]: + axis.set_xticks([]) + axis.set_yticks([]) -def _select_stages( - correction, - stages: Sequence[str] | None, -) -> tuple[str, ...]: - """Validate a requested stage sequence against recorded checkpoints.""" - available = _available_stages(correction) - if not available: + drift_rate = search.get( + "drift_rate_row_col", + getattr(self, "drift_rate", None), + ) + rate_text = ( + "unknown" + if drift_rate is None + else f"[{float(drift_rate[0]):+.5f}, {float(drift_rate[1]):+.5f}]" + ) + trusted_name = search.get("trusted_region", "whole image") + trusted_text = ( + trusted_name.replace("_", " ") + if trusted_bounds is None + else f"{trusted_name.replace('_', ' ')} {tuple(trusted_bounds)}" + ) + figure.suptitle( + f"Affine model: one drift vector {rate_text} px/scanline applied to " + f"the complete image\nfit region: {trusted_text}; " + f"displayed stage: {stage_label}", + fontsize=16, + ) + return figure, pd.DataFrame.from_records(records) + + +def diagnose_nonrigid( + self, + *, + num_knots: tuple[int, ...] = (1, 4, 6), + verbose: bool = True, + **nonrigid_options, +) -> tuple[Figure, pd.DataFrame]: + """Compare non-rigid knot counts without changing the correction + + Multiple knots let the displacement vary along each fast-scan line. This + diagnostic holds the affine or strip starting field and every optimizer + setting fixed, then compares knot counts using one common measured mask. + Use it when a single whole-field score cannot distinguish a physically + smooth correction from a periodic-lattice hop. + + The returned table pairs registration metrics with displacement-field + smoothness. A higher NCC alone is not sufficient evidence on an atomic + lattice; inspect the RGB rows and prefer the smallest knot count that + aligns distinctive features without introducing fast-direction roughness. + The supplied correction is unchanged. + + Parameters + ---------- + num_knots : tuple of int, optional + Knot counts to compare along every fast-scan line. Default is + ``(1, 4, 6)``. + verbose : bool, optional + Show one progress bar across knot-count candidates. Default is True. + **nonrigid_options + Options forwarded to :meth:`correct_nonrigid`, such as + ``num_refine_cycles``, ``knot_smoothing_sigma``, ``loss``, and + ``max_image_shift``. Diagnostic plotting is disabled inside each run. + + Returns + ------- + matplotlib.figure.Figure + One four-column row per knot count: both corrected scans, their RGB + agreement, and their absolute difference. + pandas.DataFrame + Common-mask NCC, yellow agreement, coverage, runtime, and residual + displacement smoothness for every knot count. ``fast_roughness_px`` + is the root-mean-square difference between neighboring knot + displacements along each fast-scan line. It is zero for one knot, + which has no neighbor, and does not measure image noise. + + Examples + -------- + >>> drift.correct_affine(show_combined=False) + >>> figure, metrics = drift.diagnose_nonrigid( + ... num_knots=(1, 4, 6), + ... num_refine_cycles=128, + ... knot_smoothing_sigma=8, + ... ) + >>> metrics[["num_knots", "common_ncc", "fast_roughness_px"]] + """ + if not hasattr(self, "_knots_after_affine"): raise RuntimeError( - "Drift diagnostics require a prepared correction. Run preprocess() first." + "diagnose_nonrigid() requires an affine correction. " + "Run correct_affine() first." ) - if stages is None: - return available - requested = tuple(dict.fromkeys(str(stage) for stage in stages)) - missing = [stage for stage in requested if stage not in available] - if missing: + counts = tuple(dict.fromkeys(int(value) for value in num_knots)) + if not counts or min(counts) < 1: raise ValueError( - f"Diagnostic stages are not available: {missing}. Available stages: {list(available)}." + f"num_knots must contain positive integers, got {num_knots!r}." ) - return requested - -def _stage_knots(correction, stage: str) -> tuple[NDArray, ...]: - """Return isolated knot arrays for one recorded correction stage.""" - selected = _select_stages(correction, (stage,)) - return _copy_knots(correction._diagnostic_knots[selected[0]]) - - -def _warp_stage( - correction, - stage: str, -) -> tuple[NDArray[np.float32], NDArray[np.float32]]: - """Render a checkpoint without changing the live correction object.""" - knots = _stage_knots(correction, stage) - images = [] - weights = [] - for image_index, stage_knots in enumerate(knots): - warped, coverage = correction.interpolator[image_index].warp_image( - np.asarray(correction.images[image_index].array), - stage_knots, + solve_options = { + **nonrigid_options, + "show_combined": False, + "show_scans": False, + "show_knots": False, + "show_knot_plot": False, + "show_report": False, + "verbose": False, + } + candidates = [] + progress = tqdm( + counts, + desc="Diagnosing non-rigid knot counts", + unit="candidate", + disable=not verbose or len(counts) == 1, + ) + for count in progress: + candidate = _copy_for_diagnosis(self) + start = ( + candidate._knots_after_strip + if hasattr(candidate, "_knots_after_strip") + else candidate._knots_after_affine + ) + candidate.knots = [value.clone() for value in start] + candidate._images_warped_stale = True + started = time.perf_counter() + candidate.correct_nonrigid(num_knots=count, **solve_options) + elapsed = time.perf_counter() - started + + stack = drift_apply.warped_stack(candidate, stage=None) + reference, moving = [ + canvas_center_to_scan(candidate, image) for image in stack + ] + baseline = ( + candidate._knots_after_strip + if hasattr(candidate, "_knots_after_strip") + else candidate._knots_after_affine + ) + residual = torch.stack( + [ + current - baseline + for current, baseline in zip( + candidate.knots, + baseline, + strict=True, + ) + ] + ).detach().cpu().numpy() + candidates.append( + { + "num_knots": count, + "reference": reference, + "moving": moving, + "mask": candidate.coverage_mask(), + "residual": residual, + "seconds": elapsed, + } ) - images.append(np.asarray(warped, dtype=np.float32)) - weights.append(np.asarray(coverage, dtype=np.float32)) - return np.stack(images), np.stack(weights) - - -def _masked_ncc( - reference: NDArray, - moving: NDArray, - mask: NDArray[np.bool_], -) -> float: - """Calculate normalized cross-correlation on explicitly measured pixels.""" - if int(mask.sum()) < _MIN_CORRELATION_PIXELS: - return float("nan") - first = np.asarray(reference, dtype=np.float64)[mask] - second = np.asarray(moving, dtype=np.float64)[mask] - first -= first.mean() - second -= second.mean() - denominator = np.linalg.norm(first) * np.linalg.norm(second) - if denominator <= np.finfo(np.float64).eps: - return float("nan") - return float(first @ second / denominator) - -def _masked_errors( - reference: NDArray, - moving: NDArray, - mask: NDArray[np.bool_], -) -> tuple[float, float]: - """Return mean absolute and root-mean-square disagreement.""" - if not np.any(mask): - return float("nan"), float("nan") - difference = ( - np.asarray(reference, dtype=np.float64)[mask] - np.asarray(moving, dtype=np.float64)[mask] + common_mask = np.logical_and.reduce( + [np.asarray(candidate["mask"], dtype=bool) for candidate in candidates] ) - return float(np.mean(np.abs(difference))), float(np.sqrt(np.mean(difference**2))) - - -def _scan_comparisons(stack: NDArray) -> list[tuple[NDArray, NDArray]]: - """Build pair or leave-one-out comparisons for a registered scan stack.""" - if stack.shape[0] < 2: - raise ValueError( - f"Registration diagnostics require at least two scans, got {stack.shape[0]}." + figure, axes = _diagnostic_grid(len(candidates)) + records = [] + for row, candidate in enumerate(candidates): + reference = candidate["reference"] + moving = candidate["moving"] + overlay = overlay_pair(reference, moving) + scores = strip.region_ncc( + reference, + moving, + common_mask, + device=self._device, ) - if stack.shape[0] == 2: - return [(stack[0], stack[1])] - return [ - ( - np.mean(np.delete(stack, image_index, axis=0), axis=0), - stack[image_index], + red = overlay[..., 0][common_mask] + green = overlay[..., 1][common_mask] + yellow = float( + np.minimum(red, green).sum() + / max(float(np.maximum(red, green).sum()), 1e-12) + ) + residual = candidate["residual"] + fast_first = ( + float(np.sqrt(np.mean(np.diff(residual, axis=3) ** 2))) + if residual.shape[3] > 1 + else 0.0 + ) + fast_second = ( + float(np.sqrt(np.mean(np.diff(residual, n=2, axis=3) ** 2))) + if residual.shape[3] > 2 + else 0.0 + ) + records.append( + { + "num_knots": candidate["num_knots"], + "common_ncc": scores["common"], + "top_ncc": scores["top"], + "middle_ncc": scores["middle"], + "bottom_ncc": scores["bottom"], + "yellow": yellow, + "coverage": float(common_mask.mean()), + "residual_rms_px": float(np.sqrt(np.mean(residual**2))), + "residual_max_px": float(np.max(np.abs(residual))), + "slow_roughness_px": float( + np.sqrt(np.mean(np.diff(residual, axis=2) ** 2)) + ), + "fast_roughness_px": fast_first, + "fast_curvature_px": fast_second, + "seconds": candidate["seconds"], + } ) - for image_index in range(stack.shape[0]) - ] - - -def _pair_metrics( - reference: NDArray, - moving: NDArray, - mask: NDArray[np.bool_], -) -> dict[str, float]: - """Measure one registered pair over common and regional coverage.""" - num_rows = mask.shape[0] - metrics = {"common_ncc": _masked_ncc(reference, moving, mask)} - for name, (row_start, row_stop) in zip( - ("top_ncc", "middle_ncc", "bottom_ncc"), - ( - (0, num_rows // 3), - (num_rows // 3, 2 * num_rows // 3), - (2 * num_rows // 3, num_rows), - ), - strict=True, - ): - regional_mask = np.zeros_like(mask) - regional_mask[row_start:row_stop] = mask[row_start:row_stop] - metrics[name] = _masked_ncc(reference, moving, regional_mask) - mean_absolute, root_mean_square = _masked_errors(reference, moving, mask) - metrics["mean_absolute_difference"] = mean_absolute - metrics["root_mean_square_difference"] = root_mean_square - return metrics - - -def _finite_mean(values: Sequence[float]) -> float: - """Average finite measurements without emitting all-NaN warnings.""" - finite = np.asarray(values, dtype=np.float64) - finite = finite[np.isfinite(finite)] - return float(finite.mean()) if finite.size else float("nan") - - -def _registration_data( - correction, - stages: Sequence[str] | None = None, -) -> tuple[ - tuple[str, ...], - dict[str, NDArray[np.float32]], - NDArray[np.bool_], - list[dict[str, float | str]], -]: - """Build common-mask registration evidence for completed stages.""" - selected = _select_stages(correction, stages) - stacks = {} - coverage_by_stage = {} - for stage in selected: - stack, weights = _warp_stage(correction, stage) - stacks[stage] = stack - coverage_by_stage[stage] = np.all(weights >= _COVERAGE_THRESHOLD, axis=0) - common_mask = np.logical_and.reduce([coverage_by_stage[stage] for stage in selected]) - - rows: list[dict[str, float | str]] = [] - for stage in selected: - pair_rows = [ - _pair_metrics(reference, moving, common_mask) - for reference, moving in _scan_comparisons(stacks[stage]) - ] - row: dict[str, float | str] = {"stage": stage} - for key in pair_rows[0]: - row[key] = _finite_mean([pair[key] for pair in pair_rows]) - row["coverage"] = float(common_mask.mean()) - rows.append(row) - return selected, stacks, common_mask, rows - -def _displacement_fields( - correction, - stage: str, -) -> NDArray[np.float64]: - """Return row/column knot displacement from the initial checkpoint.""" - initial = _stage_knots(correction, "initial") - selected = _stage_knots(correction, stage) - if any( - current.shape != baseline.shape - for current, baseline in zip(selected, initial, strict=True) - ): - raise ValueError("Displacement diagnostics require matching knot topology across stages.") - return np.stack( - [ - np.asarray(current, dtype=np.float64) - np.asarray(baseline, dtype=np.float64) - for current, baseline in zip(selected, initial, strict=True) - ] + contrast = np.concatenate((reference.ravel()[::16], moving.ravel()[::16])) + low, high = np.percentile(contrast[np.isfinite(contrast)], (1, 99)) + scale = max(float(high - low), np.finfo(np.float32).eps) + difference = np.abs( + np.clip((reference - low) / scale, 0, 1) + - np.clip((moving - low) / scale, 0, 1) + ) + axes[row, 0].imshow(reference, cmap="gray", vmin=low, vmax=high) + axes[row, 1].imshow(moving, cmap="gray", vmin=low, vmax=high) + axes[row, 2].imshow(overlay) + axes[row, 3].imshow(difference, cmap="magma", vmin=0, vmax=1) + count = candidate["num_knots"] + noun = "knot" if count == 1 else "knots" + axes[row, 0].set_ylabel(f"{count} {noun} per scanline", fontsize=11) + axes[row, 0].set_title("Corrected scan 0") + axes[row, 1].set_title("Corrected scan 1") + axes[row, 2].set_title( + f"RGB agreement\nNCC {scores['common']:.3f}, yellow {yellow:.1%}" + ) + axes[row, 3].set_title("Residual difference |scan 0 - scan 1|") + for axis in axes[row]: + axis.set_xticks([]) + axis.set_yticks([]) + + figure.suptitle( + "Non-rigid knot-count diagnosis\n" + "one affine/strip starting field and one common measured mask", + fontsize=16, ) - - -def _displacement_rows( - correction, - stages: Sequence[str] | None = None, -) -> list[dict[str, float | int | str]]: - """Summarize scan-line-origin displacement for each image and stage. - - Adjacent-change values are component-wise RMS measurements. Fast-knot - differences are not normalized by knot separation, so they are comparable - only when the fast-direction knot topology and spacing are unchanged. - """ - selected = _select_stages(correction, stages) - rows: list[dict[str, float | int | str]] = [] - for stage in selected: - fields = _displacement_fields(correction, stage) - for image_index, field in enumerate(fields): - magnitude = np.linalg.norm(field, axis=0) - endpoint_vector = np.mean(field[:, -1] - field[:, 0], axis=1) - adjacent_line_change = np.diff(field, axis=1) - adjacent_fast_knot_change = np.diff(field, axis=2) - fast_knot_second_difference = np.diff(field, n=2, axis=2) - rows.append( - { - "stage": stage, - "image": image_index, - "endpoint_row_displacement_px": float(endpoint_vector[0]), - "endpoint_column_displacement_px": float(endpoint_vector[1]), - "endpoint_displacement_px": float(np.linalg.norm(endpoint_vector)), - "rms_displacement_px": float(np.sqrt(np.mean(magnitude**2))), - "max_displacement_px": float(np.max(magnitude)), - "component_rms_adjacent_line_change_px": ( - float(np.sqrt(np.mean(adjacent_line_change**2))) - if adjacent_line_change.size - else 0.0 - ), - "component_rms_adjacent_fast_knot_change_px": ( - float(np.sqrt(np.mean(adjacent_fast_knot_change**2))) - if adjacent_fast_knot_change.size - else 0.0 - ), - "component_rms_fast_knot_second_difference_px": ( - float(np.sqrt(np.mean(fast_knot_second_difference**2))) - if fast_knot_second_difference.size - else 0.0 - ), - } - ) - return rows + return figure, pd.DataFrame.from_records(records) diff --git a/src/quantem/imaging/drift/fourdstem.py b/src/quantem/imaging/drift/fourdstem.py index 7dcfe2bb..92a5e5f3 100644 --- a/src/quantem/imaging/drift/fourdstem.py +++ b/src/quantem/imaging/drift/fourdstem.py @@ -7,36 +7,11 @@ from tqdm.auto import tqdm from quantem.imaging.drift.apply import ( - _apply_scan_field, - _raw_frame_displacement, + apply_correction_to_dataset, + crop_slices, + padding_offset, ) - - -def data_array(dataset): - """Return the resident array or tensor from a QuantEM dataset wrapper.""" - if isinstance(dataset, (np.ndarray, torch.Tensor)): - return dataset - array = getattr(dataset, "array", None) - if array is not None: - return array - tensor = getattr(dataset, "_tensor", None) - if tensor is not None: - return tensor - raise TypeError( - "4D-STEM input must be an array, tensor, or QuantEM dataset with " - "resident .array/.tensor data." - ) - - -def padding_offset( - canvas_shape: tuple[int, int], - scan_shape: tuple[int, int], -) -> tuple[float, float]: - """Return the centered scan offset in the padded solver canvas.""" - return ( - (canvas_shape[0] - scan_shape[0]) / 2.0, - (canvas_shape[1] - scan_shape[1]) / 2.0, - ) +from quantem.imaging.drift.core import knots as drift_knots @dataclass @@ -119,7 +94,13 @@ def integrate_virtual_detector( -------- >>> adf = integrate_virtual_detector(data, detector_mask=annulus) """ - dataset = data_array(dataset) + try: + from quantem.gpu.detector import masked_sum + except ModuleNotFoundError as exc: + if exc.name not in {"quantem.gpu", "quantem.gpu.detector"}: + raise + masked_sum = None + if reduce not in {"mean", "sum"}: raise ValueError(f"reduce must be 'mean' or 'sum', got {reduce!r}") scan_shape = tuple(int(value) for value in dataset.shape[:2]) @@ -140,17 +121,25 @@ def integrate_virtual_detector( num_selected = int(mask.sum()) if num_selected == 0: raise ValueError("detector_mask selects zero detector pixels") - if isinstance(dataset, torch.Tensor): + if isinstance(dataset, np.ndarray) and np.issubdtype( + dataset.dtype, + np.floating, + ): + flat = dataset.reshape(*scan_shape, -1) + image = flat[..., mask.reshape(-1)].sum(axis=-1, dtype=np.float32) + elif masked_sum is not None: + image = masked_sum(dataset, mask).reshape(scan_shape) + elif isinstance(dataset, torch.Tensor): mask_t = torch.as_tensor(mask, device=dataset.device) - flattened = dataset.reshape(*scan_shape, -1).to(torch.float32) - image = flattened[..., mask_t.reshape(-1)].sum(-1) - image = image.detach().cpu().numpy() + flat = dataset.reshape(*scan_shape, -1).to(torch.float32) + image = flat[..., mask_t.reshape(-1)].sum(-1).detach().cpu().numpy() else: array = dataset.get() if hasattr(dataset, "get") else np.asarray(dataset) - image = array.reshape(*scan_shape, -1)[..., mask.reshape(-1)].sum( - -1, - dtype=np.float64, - ).astype(np.float32) + flat = array.reshape(*scan_shape, -1) + image = flat[..., mask.reshape(-1)].sum( + axis=-1, + dtype=np.float32, + ) return image / num_selected if reduce == "mean" else image @@ -181,7 +170,9 @@ def drift_field(self, idx: int) -> torch.Tensor: "Run dc.preprocess().correct_affine() (and optionally " ".correct_nonrigid()) before drift_field()." ) - return _raw_frame_displacement(self, idx) + return drift_knots.interpolator(self, idx).drift_raw( + self._initial_knots[idx] + ) def probe_positions( @@ -229,47 +220,31 @@ def probe_positions( "dc.preprocess() before exporting nominal or corrected " "probe positions." ) - index = image_index % len(self.images) + index = image_index % len(self.imgs) knots = ( self.knots[index] if corrected else self._initial_knots[index] ) - row, column = self.interpolator[index].transform_coordinates(knots) - positions = np.stack([row, column], axis=-1).astype(np.float32) + row, column = drift_knots.interpolator(self, index, knots).to_canvas() + positions = torch.stack([row, column], dim=-1) if strip_padding: pad_row, pad_column = padding_offset( (self.shape[1], self.shape[2]), - self.images[0].shape[:2], + self.imgs[0].shape[:2], + ) + positions -= torch.tensor( + [pad_row, pad_column], + device=positions.device, + dtype=positions.dtype, ) - positions -= np.asarray([pad_row, pad_column], dtype=np.float32) - result = positions + result = to_numpy(positions, dtype=np.float32) if plot: - import matplotlib.pyplot as plt - - nominal = self.probe_positions( - index, - corrected=False, + self.plot_probe_positions( + image_index=index, strip_padding=strip_padding, - plot=False, - ) - plt.plot( - nominal[::stride, ::stride, 1], - nominal[::stride, ::stride, 0], - ".", - alpha=0.4, - label="nominal", - ) - plt.plot( - result[::stride, ::stride, 1], - result[::stride, ::stride, 0], - ".", - alpha=0.6, - label="corrected", + stride=stride, ) - plt.gca().invert_yaxis() - plt.gca().set_aspect("equal") - plt.legend() return result @@ -327,7 +302,7 @@ def corrected_virtual_images( "corrected_virtual_images() requires preprocess() and " "correct_affine() first." ) - if len(self.images) != 2: + if len(self.imgs) != 2: raise ValueError( "corrected_virtual_images() expects exactly two scan images" ) @@ -340,14 +315,11 @@ def corrected_virtual_images( np.asarray(image_0, dtype=np.float32), np.asarray(image_1, dtype=np.float32), ] - if ( - images[0].shape != self.images[0].shape - or images[1].shape != self.images[1].shape - ): + if images[0].shape != self.imgs[0].shape or images[1].shape != self.imgs[1].shape: raise ValueError( "virtual image shapes must match the raw scan images used for drift " f"alignment: got {images[0].shape}, {images[1].shape}; expected " - f"{self.images[0].shape}, {self.images[1].shape}" + f"{self.imgs[0].shape}, {self.imgs[1].shape}" ) if output_frame == "canvas": @@ -355,32 +327,42 @@ def corrected_virtual_images( components = [] coverages = [] for image_index, image in enumerate(images): - corrected, coverage = self.interpolator[image_index].warp_image( + image_t = torch.as_tensor( image, - self.knots[image_index], - output_shape=canvas_shape, - kde_sigma=self.kde_sigma, - pad_value=0.0, + device=self._device, + dtype=torch.float32, + ) + corrected, coverage = drift_knots.interpolator( + self, + image_index, + ).warp_to_canvas( + image_t, + canvas_shape, + self.kde_sigma, + 0.0, ) components.append(corrected) coverages.append(coverage) valid = [coverage >= 1e-3 for coverage in coverages] - contribution_count = valid[0].astype(np.float32) + valid[1].astype( - np.float32 + contribution_count = valid[0].to(torch.float32) + valid[1].to( + torch.float32 ) - merged = np.divide( - components[0] * valid[0] + components[1] * valid[1], - contribution_count, - out=np.zeros(canvas_shape, dtype=np.float32), - where=contribution_count > 0, + merged = torch.where( + contribution_count > 0, + ( + components[0] * valid[0] + + components[1] * valid[1] + ) + / contribution_count.clamp_min(1), + 0.0, ) return { "corrected_image": to_numpy(merged, dtype=np.float32), "corrected_image_0": to_numpy(components[0], dtype=np.float32), "corrected_image_1": to_numpy(components[1], dtype=np.float32), "coverage_image": to_numpy( - np.maximum(coverages[0], coverages[1]), + torch.maximum(coverages[0], coverages[1]), dtype=np.float32, ), "coverage_image_0": to_numpy(coverages[0], dtype=np.float32), @@ -389,17 +371,24 @@ def corrected_virtual_images( components = [] for image_index, image in enumerate(images): - corrected = _apply_scan_field( - self, + image_t = torch.as_tensor( image[..., None], + device=self._device, + dtype=torch.float32, + ) + corrected = apply_correction_to_dataset( + self, + image_t, image_index=image_index, + mode="bilinear", chunk_size=1, - output_dtype=np.float32, + output_dtype=torch.float32, + output_device=self._device, )[..., 0] if image_index == 1: rot_k = _rot90_to_image0_frame(self, image_index=1) if rot_k: - corrected = np.rot90(corrected, k=rot_k, axes=(0, 1)).copy() + corrected = torch.rot90(corrected, k=rot_k, dims=(0, 1)) components.append(corrected) merged = (components[0] + components[1]) * 0.5 @@ -491,16 +480,14 @@ def regional_diffraction_patterns( "regional_diffraction_patterns() needs the two raw 4D-STEM datasets. " "Pass datasets=(scan_0, scan_90) when using a saved correction." ) - source_datasets = tuple(data_array(dataset) for dataset in source_datasets) + source_datasets = tuple(source_datasets) if len(source_datasets) != 2 or any(dataset is None for dataset in source_datasets): raise ValueError( "regional_diffraction_patterns() expects exactly two raw 4D-STEM " f"datasets, got {len(source_datasets)}" ) scan_shapes = [tuple(int(value) for value in data.shape[:2]) for data in source_datasets] - expected_shapes = [ - tuple(int(value) for value in image.shape) for image in self.images - ] + expected_shapes = [tuple(int(value) for value in image.shape) for image in self.imgs] detector_shapes = [tuple(int(value) for value in data.shape[2:]) for data in source_datasets] if scan_shapes != expected_shapes: raise ValueError( @@ -592,6 +579,75 @@ def regional_diffraction_patterns( } +def corrected_4dstem_views(correction, *, det_bin: int = 1) -> list[np.ndarray]: + """Prepare a compact raw-to-corrected 4D-STEM comparison. + + Detector binning before correction preserves the correction field while + avoiding work on detector detail that the interactive viewer will discard. + The returned stages share image 0's scan frame and the solved crop. + """ + def detector_bin(cube): + if det_bin == 1: + return cube + detector_rows, detector_columns = cube.shape[-2:] + shape = cube.shape + reshaped = cube.reshape( + shape[0], + shape[1], + detector_rows // det_bin, + det_bin, + detector_columns // det_bin, + det_bin, + ) + if isinstance(reshaped, torch.Tensor): + dtype = reshaped.dtype if reshaped.is_floating_point() else torch.int32 + else: + dtype = ( + reshaped.dtype + if np.issubdtype(reshaped.dtype, np.floating) + else np.int32 + ) + return reshaped.sum((3, 5), dtype=dtype) + + raw_0, raw_1 = (detector_bin(dataset) for dataset in correction._datasets) + corrected_0 = correction.apply_correction( + raw_0, + image_index=0, + output_dtype="same", + output_device=correction.device, + verbose=False, + ) + corrected_1 = correction.apply_correction( + raw_1, + image_index=1, + output_dtype="same", + output_device=correction.device, + verbose=False, + ) + quarter_turns = _rot90_to_image0_frame(correction) + if quarter_turns: + corrected_1 = ( + torch.rot90(corrected_1, quarter_turns, dims=(0, 1)) + if isinstance(corrected_1, torch.Tensor) + else np.rot90(corrected_1, quarter_turns, axes=(0, 1)).copy() + ) + if isinstance(corrected_0, torch.Tensor): + if corrected_0.is_floating_point(): + merged = (corrected_0 + corrected_1) * 0.5 + else: + merged = ( + (corrected_0.to(torch.int64) + corrected_1.to(torch.int64)) >> 1 + ).to(corrected_0.dtype) + elif np.issubdtype(corrected_0.dtype, np.floating): + merged = (corrected_0 + corrected_1) * 0.5 + else: + merged = ( + (corrected_0.astype(np.int64) + corrected_1.astype(np.int64)) >> 1 + ).astype(corrected_0.dtype) + rows, columns = crop_slices(correction) + return [to_numpy(cube[rows, columns]) for cube in (raw_0, corrected_0, merged)] + + def corrected_4dstem( self, *, @@ -663,31 +719,21 @@ def corrected_4dstem( and isinstance(datasets[1], torch.Tensor) and datasets[1].is_cuda ) - corrected_4dstem_0 = _apply_scan_field( - self, - datasets[0], - image_index=0, - mode=mode, - chunk_size=chunk_size, - output_dtype=output_dtype, - output=output_0, + corrected_4dstem_0 = apply_correction_to_dataset( + self, None, image_index=0, mode=mode, chunk_size=chunk_size, + output_dtype=output_dtype, output_device=output_device, + output=output_0, verbose=verbose, + progress_desc="Correcting scan 1/2", ) - if output_device is not None and output_0 is None: - corrected_4dstem_0 = torch.as_tensor(corrected_4dstem_0).to(output_device) if inputs_on_device: self._datasets[0] = None torch.cuda.empty_cache() - corrected_4dstem_1 = _apply_scan_field( - self, - datasets[1], - image_index=1, - mode=mode, - chunk_size=chunk_size, - output_dtype=output_dtype, - output=output_1, + corrected_4dstem_1 = apply_correction_to_dataset( + self, None, image_index=1, mode=mode, chunk_size=chunk_size, + output_dtype=output_dtype, output_device=output_device, + output=output_1, verbose=verbose, + progress_desc="Correcting scan 2/2", ) - if output_device is not None and output_1 is None: - corrected_4dstem_1 = torch.as_tensor(corrected_4dstem_1).to(output_device) if inputs_on_device: self._datasets[1] = None self._datasets_consumed = True @@ -702,7 +748,7 @@ def corrected_4dstem( else: corrected_4dstem_1 = np.rot90( corrected_4dstem_1, k=rot_k, axes=(0, 1), - ) + ).copy() corrected_4dstem = None if merge: @@ -723,15 +769,21 @@ def corrected_4dstem( disable=not verbose or len(row_starts) <= 1, ) if isinstance(corrected_4dstem_0, torch.Tensor): - corrected_4dstem = torch.empty_like( - corrected_4dstem_0, - dtype=torch.float32, + merge_dtype = ( + corrected_4dstem_0.dtype + if corrected_4dstem_0.is_floating_point() + else torch.float32 + ) + corrected_4dstem = torch.empty( + corrected_4dstem_0.shape, + dtype=merge_dtype, + device=corrected_4dstem_0.device, ) for r0 in row_starts: r1 = min(r0 + row_block, Hm) corrected_4dstem[r0:r1] = ( - corrected_4dstem_0[r0:r1].to(torch.float32) - + corrected_4dstem_1[r0:r1].to(torch.float32) + corrected_4dstem_0[r0:r1].to(merge_dtype) + + corrected_4dstem_1[r0:r1].to(merge_dtype) ) * 0.5 merge_progress.update(r1 - r0) else: @@ -751,8 +803,8 @@ def corrected_4dstem( # Extract raw VDFs from the stored alignment images. The scan collection # reference is the scalar channel correction implied by the learned scan # drift fields, not an external ground truth. - alignment_vdf_0 = np.asarray(self.images[0].array) - alignment_vdf_1 = np.asarray(self.images[1].array) + alignment_vdf_0 = np.asarray(self.imgs[0].array) + alignment_vdf_1 = np.asarray(self.imgs[1].array) scalar_corrected_vdf = corrected_virtual_images( self, alignment_vdf_0, diff --git a/src/quantem/imaging/drift/io.py b/src/quantem/imaging/drift/io.py index a7f19e49..cef25850 100644 --- a/src/quantem/imaging/drift/io.py +++ b/src/quantem/imaging/drift/io.py @@ -1,4 +1,4 @@ -"""Velox EMD input and metadata pairing for drift correction.""" +"""Read Velox EMD images, spectrum images, and scan metadata for drift correction.""" import json import math @@ -7,12 +7,28 @@ import h5py import numpy as np -from quantem.core.datastructures.dataset2d import Dataset2d -from quantem.core.datastructures.dataset3d import Dataset3d + +def _axis_calibration( + axes: list[dict], + ndim: int, +) -> tuple[list[float], list[float], list[str]]: + """Return origin, sampling, and units in array-axis order.""" + indexed = sorted( + ( + (int(axis.get("index_in_array", fallback)), axis) + for fallback, axis in enumerate(axes[:ndim]) + ), + key=lambda item: item[0], + ) + ordered = [axis for _, axis in indexed] + origin = [float(axis.get("offset", 0.0)) for axis in ordered] + sampling = [float(axis.get("scale", 1.0)) for axis in ordered] + units = [(axis.get("units") or "pixels") for axis in ordered] + return origin, sampling, units def _deinterleave_velox_metadata(raw: bytes) -> bytes: - """Undo repeated-byte encodings found in some Velox metadata records.""" + """Undo Velox's repeated-byte metadata encoding for image stacks.""" for stride in range(2, 17): candidate = raw[::stride] if candidate.startswith(b"{"): @@ -21,16 +37,13 @@ def _deinterleave_velox_metadata(raw: bytes) -> bytes: def _read_velox_metadata_record(file_path: str | Path) -> dict: - """Decode the first image metadata record without loading image pixels.""" + """Decode the first Velox image metadata record without loading pixels.""" with h5py.File(file_path, "r") as handle: image_group = handle.get("Data/Image") if image_group is None: return {} for image_name in image_group: - metadata = image_group[image_name].get("Metadata") - if metadata is None: - continue - raw = bytes(metadata[:]).replace(b"\x00", b"") + raw = bytes(image_group[image_name]["Metadata"][:]).replace(b"\x00", b"") for candidate in (raw, _deinterleave_velox_metadata(raw)): try: return json.loads(candidate.decode("utf-8", errors="ignore")) @@ -39,8 +52,8 @@ def _read_velox_metadata_record(file_path: str | Path) -> dict: return {} -def read_emd_metadata(file_path: str | Path) -> dict[str, object]: - """Read normalized Velox acquisition metadata without loading pixels. +def read_emd_metadata(file_path: str | Path) -> dict: + """Read normalized Velox acquisition metadata without loading image data. Parameters ---------- @@ -49,17 +62,25 @@ def read_emd_metadata(file_path: str | Path) -> dict[str, object]: Returns ------- - dict[str, object] - Scan rotation, stage position, calibration, field of view, acquisition - time, acquisition context, and the decoded original metadata. + dict + Normalized fields: ``scan_rotation_deg``, ``magnification``, + ``stage_xy_m``, ``pixel_size_nm``, ``scan_shape``, ``fov_m``, + ``acquisition_timestamp``, and ``acquisition_context``. Missing Velox + fields are returned as ``None``. ``original_metadata`` retains the + decoded metadata tree for specialized consumers. + + Notes + ----- + Multi-stream EMD files can contain slightly different stage readouts on + individual image streams. Uses the first Velox image metadata record, + matching QuantEM Live's pairing policy. Examples -------- - >>> metadata = read_emd_metadata("scan.emd") + >>> metadata = read_emd_metadata("scan_0.emd") >>> metadata["scan_rotation_deg"] 0.0 """ - file_path = Path(file_path) metadata = _read_velox_metadata_record(file_path) with h5py.File(file_path, "r") as handle: acquisition_context = ( @@ -71,33 +92,41 @@ def read_emd_metadata(file_path: str | Path) -> dict[str, object]: rotation = metadata.get("Scan", {}).get("ScanRotation") magnification = metadata.get("Optics", {}).get("NominalMagnification") position = metadata.get("Stage", {}).get("Position", {}) or {} - stage_position = ( + stage_xy_m = ( (float(position["x"]), float(position["y"])) if "x" in position and "y" in position else None ) pixel_size_m = metadata.get("BinaryResult", {}).get("PixelSize", {}).get("width") scan_size = metadata.get("Scan", {}).get("ScanSize", {}) or {} - scan_columns = scan_size.get("width") - scan_rows = scan_size.get("height", scan_columns) + scan_width = scan_size.get("width") + scan_height = scan_size.get("height", scan_width) scan_shape = ( - (int(scan_rows), int(scan_columns)) - if scan_rows and scan_columns + (int(scan_height), int(scan_width)) + if scan_height and scan_width else None ) - timestamp = metadata.get("Acquisition", {}).get( - "AcquisitionStartDatetime", {} - ).get("DateTime") + timestamp = ( + metadata.get("Acquisition", {}) + .get("AcquisitionStartDatetime", {}) + .get("DateTime") + ) return { "path": str(file_path), - "scan_rotation_deg": None if rotation is None else math.degrees(float(rotation)), - "magnification": None if magnification is None else float(magnification), - "stage_position_m": stage_position, - "pixel_size_nm": None if pixel_size_m is None else float(pixel_size_m) * 1e9, + "scan_rotation_deg": ( + None if rotation is None else math.degrees(float(rotation)) + ), + "magnification": ( + None if magnification is None else float(magnification) + ), + "stage_xy_m": stage_xy_m, + "pixel_size_nm": ( + None if pixel_size_m is None else float(pixel_size_m) * 1e9 + ), "scan_shape": scan_shape, - "field_of_view_m": ( - float(pixel_size_m) * float(scan_columns) - if pixel_size_m and scan_columns + "fov_m": ( + float(pixel_size_m) * float(scan_width) + if pixel_size_m and scan_width else None ), "acquisition_timestamp": ( @@ -108,24 +137,9 @@ def read_emd_metadata(file_path: str | Path) -> dict[str, object]: } -def _axis_calibration(axes: list[dict], ndim: int) -> tuple[list[float], list[float], list[str]]: - """Return origin, sampling, and units in array-axis order.""" - indexed = sorted( - ( - (int(axis.get("index_in_array", fallback)), axis) - for fallback, axis in enumerate(axes[:ndim]) - ), - key=lambda item: item[0], - ) - ordered = [axis for _, axis in indexed] - origin = [float(axis.get("offset", 0.0)) for axis in ordered] - sampling = [float(axis.get("scale", 1.0)) for axis in ordered] - units = [(axis.get("units") or "pixels") for axis in ordered] - return origin, sampling, units - +def _image_dataset(stream, path, metadata): + from quantem.core.datastructures.dataset2d import Dataset2d -def _image_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset2d: - """Build a calibrated image dataset from one RosettaSciIO stream.""" origin, sampling, units = _axis_calibration(stream.get("axes", []), 2) title = stream.get("metadata", {}).get("General", {}).get("title") image = Dataset2d.from_array( @@ -135,57 +149,22 @@ def _image_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset2d: sampling=sampling, units=units, ) - image.file_path = str(path) + image.file_path = path image.metadata.update(metadata) return image -def read_emd(path: str | Path) -> Dataset2d: - """Read the calibrated HAADF image and recorded scan geometry. - - Parameters - ---------- - path : str or Path - Velox EMD acquisition. - - Returns - ------- - Dataset2d - HAADF image carrying normalized acquisition metadata. - - Examples - -------- - >>> image = read_emd("scan.emd") - >>> image.metadata["scan_rotation_deg"] - 0.0 - """ - from rsciio.emd import file_reader +def _spectrum_dataset(stream, path, metadata): + """Build a calibrated ``(scan_row, scan_col, energy)`` Dataset3d.""" + from quantem.core.datastructures.dataset3d import Dataset3d - streams = file_reader(str(path), select_type="images") - stream = next( - ( - item - for item in streams - if item["data"].ndim == 2 - and item.get("metadata", {}).get("General", {}).get("title") == "HAADF" - ), - next((item for item in streams if item["data"].ndim == 2), None), - ) - if stream is None: - raise ValueError(f"{path} contains no two-dimensional image stream.") - return _image_dataset(stream, path, read_emd_metadata(path)) - - -def _spectrum_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset3d: - """Build a scan-row, scan-column, energy Dataset3d.""" array = np.asarray(stream["data"]) - axes = stream.get("axes", []) - indexed_axes = [ + axes = [ axis for _, axis in sorted( ( (int(axis.get("index_in_array", fallback)), axis) - for fallback, axis in enumerate(axes[:3]) + for fallback, axis in enumerate(stream.get("axes", [])[:3]) ), key=lambda item: item[0], ) @@ -193,7 +172,7 @@ def _spectrum_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset energy_axis = next( ( index - for index, axis in enumerate(indexed_axes) + for index, axis in enumerate(axes) if (axis.get("units") or "").lower() in {"ev", "kev"} or "energy" in (axis.get("name") or "").lower() ), @@ -201,11 +180,10 @@ def _spectrum_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset ) if energy_axis != 2: array = np.moveaxis(array, energy_axis, 2) - axis = indexed_axes.pop(energy_axis) - indexed_axes.append(axis) - origin = [float(axis.get("offset", 0.0)) for axis in indexed_axes] - sampling = [float(axis.get("scale", 1.0)) for axis in indexed_axes] - units = [(axis.get("units") or "pixels") for axis in indexed_axes] + axes.append(axes.pop(energy_axis)) + origin = [float(axis.get("offset", 0.0)) for axis in axes] + sampling = [float(axis.get("scale", 1.0)) for axis in axes] + units = [(axis.get("units") or "pixels") for axis in axes] title = stream.get("metadata", {}).get("General", {}).get("title") spectrum = Dataset3d.from_array( np.ascontiguousarray(array, dtype=np.float32), @@ -220,41 +198,86 @@ def _spectrum_dataset(stream: dict, path: str | Path, metadata: dict) -> Dataset return spectrum +def read_emd(path: str | Path): + """Read the HAADF image and acquisition geometry from a Velox EMD file. + + Drift correction needs one calibrated scan image plus its recorded scan + direction. RosettaSciIO selects the image stream. ``read_emd_metadata`` + supplies rotation, stage position, magnification, and acquisition time. + + Parameters + ---------- + path : str or Path + Velox EMD acquisition. + + Returns + ------- + Dataset2d + Calibrated HAADF image carrying the normalized acquisition metadata. + + Examples + -------- + >>> image = read_emd("scan_0.emd") + >>> image.metadata["scan_rotation_deg"] + 0.0 + """ + from rsciio.emd import file_reader + + streams = file_reader(str(path), select_type="images") + stream = next( + ( + item + for item in streams + if item["data"].ndim == 2 + and item.get("metadata", {}).get("General", {}).get("title") == "HAADF" + ), + next((item for item in streams if item["data"].ndim == 2), None), + ) + if stream is None: + raise ValueError(f"{path} contains no two-dimensional image stream") + + return _image_dataset(stream, path, read_emd_metadata(path)) + + def read_emd_eds( path: str | Path, *, load_spectrum: bool = False, verbose: bool = True, ) -> dict[str, object]: - """Read one Velox EDS or EELS spectrum-image acquisition. + """Read images from a Velox EDS/EELS spectrum-image EMD. - The full spectrum is opt-in because it can be large. When loaded, it is a - calibrated :class:`Dataset3d` with ``(scan_row, scan_col, energy)`` axes. - Energy calibration and units are retained exactly as recorded. + Encapsulates the raw rsciio stream traversal so callers never write a + ``for ds in file_reader(...)`` loop: the simultaneously-acquired HAADF + survey and any Velox pre-quantified 2-D element maps are separated here + and returned by name. The full spectrum is opt-in because expanding it can + require tens of gigabytes while drift correction needs only the HAADF. + The scan angle comes from the EMD metadata (never typed). Parameters ---------- path : str or Path Velox spectrum-image EMD file. - load_spectrum : bool, default False - Load the full spectrum in addition to the HAADF and element maps. - verbose : bool, default True - Print a compact inventory. + load_spectrum : bool, optional + Load the full ``(row, col, energy)`` spectrum. The default ``False`` + loads only the HAADF and stored 2-D elemental maps. + verbose : bool, optional + Print a compact summary of the loaded images. Returns ------- dict[str, object] - HAADF image, optional spectrum, element maps, and acquisition metadata. + Calibrated HAADF image, optional spectrum and energy axis, stored + element maps, and the acquisition geometry needed for correction. Examples -------- - >>> acquisition = read_emd_eds("spectrum_image.emd", load_spectrum=True) - >>> acquisition["spectrum"].units[-1] - 'keV' + >>> acquisition = read_emd_eds("spectrum_image.emd") + >>> acquisition["haadf"].shape + (2048, 2048) """ from rsciio.emd import file_reader - metadata = read_emd_metadata(path) streams = list( file_reader( str(path), @@ -266,13 +289,14 @@ def read_emd_eds( None, ) if load_spectrum and spectrum_stream is None: - raise ValueError(f"{path} contains no three-dimensional spectrum.") + raise ValueError(f"{path} contains no 3-D spectrum") haadf_stream = next( ( stream for stream in streams if np.asarray(stream["data"]).ndim == 2 - and stream.get("metadata", {}).get("General", {}).get("title") == "HAADF" + and stream.get("metadata", {}).get("General", {}).get("title") + == "HAADF" ), None, ) @@ -282,34 +306,281 @@ def read_emd_eds( for stream in streams if np.asarray(stream["data"]).ndim == 2 and stream is not haadf_stream } - haadf = ( - None - if haadf_stream is None - else _image_dataset(haadf_stream, path, metadata) - ) + metadata = read_emd_metadata(path) + scan_rot = metadata["scan_rotation_deg"] + if scan_rot is None: + scan_rot = 0.0 + px_nm = metadata["pixel_size_nm"] + px_nm = float("nan") if px_nm is None else float(px_nm) + if haadf_stream is None: + haadf = None + else: + haadf = _image_dataset( + haadf_stream, + path, + metadata | { + "scan_rotation_deg": float(scan_rot), + "pixel_size_nm": px_nm, + }, + ) spectrum = ( None if spectrum_stream is None else _spectrum_dataset(spectrum_stream, path, metadata) ) - if haadf is None and not element_maps: - raise ValueError(f"{path} contains no HAADF image or element maps.") + if spectrum is None: + energy_axis = None + else: + scale = float(spectrum.sampling[-1]) + offset = float(spectrum.origin[-1]) + to_kev = 1e-3 if str(spectrum.units[-1]).lower() == "ev" else 1.0 + energy_axis = ( + offset + np.arange(spectrum.shape[-1], dtype=np.float64) * scale + ) * to_kev + image = haadf if haadf is not None else next(iter(element_maps.values()), None) + if image is None: + raise ValueError(f"{path} contains no 2-D HAADF or elemental maps") + shape = tuple(image.shape) if verbose: + spectrum_shape = "not loaded" if spectrum is None else str(spectrum.shape) print( f"HAADF {None if haadf is None else haadf.shape} " - f"elements {list(element_maps)} " - f"spectrum {None if spectrum is None else spectrum.shape} " - f"scan {metadata['scan_rotation_deg']} deg" + f"elements {list(element_maps)} spectrum {spectrum_shape} " + f"scan {scan_rot:.1f}°" ) return { "haadf": haadf, "spectrum": spectrum, + "cube": None if spectrum is None else spectrum.array, + "energy_axis_keV": energy_axis, "element_maps": element_maps, + "scan_rotation_deg": scan_rot, + "pixel_size_nm": px_nm, + "shape": shape, "metadata": metadata, "path": str(path), } +def scan_pairs( + folder: str | Path, + *, + max_rotation_tolerance_deg: float = 5.0, +): + """Pair orthogonal Velox scans acquired from the same specimen area. + + Stage position identifies the shared field of view; scan rotation identifies + the orthogonal acquisition. Shape, pixel calibration, field of view, and + nominal magnification reject incompatible acquisitions when those metadata + are present. Only unique mutual matches are paired. The returned inventory + includes every file and a reason for every acquisition that is not included. + + Parameters + ---------- + folder : str or Path + Session folder containing Velox EMD files. + max_rotation_tolerance_deg : float, optional + Angular tolerance around 0 and ±90 degrees. Default is 5 degrees. + + Returns + ------- + pandas.DataFrame + Acquisition metadata and pair assignments. ``pair_order=0`` marks the + scan closest to 0 degrees and ``pair_order=1`` its orthogonal partner. + + Examples + -------- + >>> pairs = scan_pairs("~/data/session") + >>> pairs[pairs.pair_order == 0][["file", "partner"]] + """ + import pandas as pd + + folder = Path(folder).expanduser() + records = [] + for path in sorted(folder.iterdir()): + if path.name.startswith("._"): + continue + if path.suffix.lower() == ".npy": + records.append({"file": path.name, "shape": tuple(np.load(path, mmap_mode="r").shape)}) + continue + if path.suffix.lower() != ".emd": + continue + + metadata = read_emd_metadata(path) + shape = metadata["scan_shape"] + pixel_size_nm = metadata["pixel_size_nm"] + stage = metadata["stage_xy_m"] + records.append( + { + "file": path.name, + "shape": shape, + "pixel_size_nm": pixel_size_nm, + "fov_nm": None if metadata["fov_m"] is None else metadata["fov_m"] * 1e9, + "magnification": metadata["magnification"], + "rotation_deg": metadata["scan_rotation_deg"], + "stage_x_m": None if stage is None else stage[0], + "stage_y_m": None if stage is None else stage[1], + "fov_m": metadata["fov_m"], + "acquired": metadata["acquisition_timestamp"], + "acquisition_context": metadata.get("acquisition_context", "image"), + } + ) + + table = pd.DataFrame(records) + if table.empty: + return table + for column in ( + "pixel_size_nm", + "fov_nm", + "magnification", + "rotation_deg", + "stage_x_m", + "stage_y_m", + "fov_m", + "acquired", + "acquisition_context", + ): + if column not in table: + table[column] = None + + table = table.reset_index(drop=True) + table["pair"] = "" + table["partner"] = "" + table["pair_order"] = pd.array([pd.NA] * len(table), dtype="Int64") + table["partner_rotation_deg"] = np.nan + table["relative_partner_rotation_deg"] = np.nan + table["stage_distance_nm"] = np.nan + table["pair_tolerance_nm"] = np.nan + table["pair_status"] = "not_included" + table["pair_reason"] = "" + + tolerance = float(max_rotation_tolerance_deg) + zero_indices = [ + index + for index, angle in table.rotation_deg.items() + if pd.notna(angle) + and abs(float(angle)) < tolerance + and table.at[index, "acquisition_context"] != "spectrum_image" + ] + ninety_indices = [ + index + for index, angle in table.rotation_deg.items() + if pd.notna(angle) and abs(abs(float(angle)) - 90.0) < tolerance + and table.at[index, "acquisition_context"] != "spectrum_image" + ] + candidates_by_zero = {zero: [] for zero in zero_indices} + zeros_by_ninety = {ninety: [] for ninety in ninety_indices} + for zero in zero_indices: + for ninety in ninety_indices: + incompatible = False + if table.at[zero, "shape"] and table.at[ninety, "shape"]: + incompatible = tuple(table.at[zero, "shape"]) != tuple( + table.at[ninety, "shape"] + ) + for column, relative_tolerance in ( + ("pixel_size_nm", 0.02), + ("fov_m", 0.02), + ("magnification", 0.02), + ): + first, second = table.loc[[zero, ninety], column] + if pd.notna(first) and pd.notna(second): + scale = max(abs(float(first)), abs(float(second)), 1e-30) + incompatible |= abs(float(first) - float(second)) > ( + relative_tolerance * scale + ) + if incompatible: + continue + stage_values = table.loc[ + [zero, ninety], ["stage_x_m", "stage_y_m"] + ].to_numpy(dtype=float) + if np.isfinite(stage_values).all(): + distance = float(np.linalg.norm(stage_values[0] - stage_values[1])) + fovs = [ + float(value) + for value in table.loc[[zero, ninety], "fov_m"] + if pd.notna(value) and float(value) > 0 + ] + pair_tolerance = max(0.25 * min(fovs), 10e-9) if fovs else 10e-9 + if distance > pair_tolerance: + continue + else: + distance, pair_tolerance = float("inf"), float("nan") + candidate = (distance, ninety, pair_tolerance) + candidates_by_zero[zero].append(candidate) + zeros_by_ninety[ninety].append((distance, zero, pair_tolerance)) + + for index, angle in table.rotation_deg.items(): + if table.at[index, "acquisition_context"] == "spectrum_image": + table.at[index, "pair_reason"] = ( + "Spectrum-image acquisition belongs in the EDS/EELS reference workflow." + ) + elif pd.isna(angle): + table.at[index, "pair_reason"] = "Missing scan-rotation metadata." + elif index not in zero_indices and index not in ninety_indices: + table.at[index, "pair_reason"] = ( + f"Scan rotation is not within {tolerance:g}° of 0° or ±90°." + ) + + pair_count = 0 + for zero in zero_indices: + candidates = candidates_by_zero[zero] + if not candidates: + table.at[zero, "pair_reason"] = ( + "No orthogonal scan has compatible shape, calibration, field of view, and stage position." + ) + continue + if len(candidates) > 1: + table.at[zero, "pair_reason"] = ( + f"Ambiguous: {len(candidates)} compatible ±90° scans match this 0° acquisition." + ) + for _, ninety, _ in candidates: + table.at[ninety, "pair_reason"] = ( + "Ambiguous: this ±90° scan is one of multiple candidates for the same 0° acquisition." + ) + continue + + distance, ninety, pair_tolerance = candidates[0] + reverse_candidates = zeros_by_ninety[ninety] + if len(reverse_candidates) != 1: + table.at[zero, "pair_reason"] = ( + "Ambiguous: the compatible ±90° scan also matches multiple 0° acquisitions." + ) + table.at[ninety, "pair_reason"] = ( + f"Ambiguous: {len(reverse_candidates)} compatible 0° scans match this ±90° acquisition." + ) + continue + + pair_count += 1 + pair_name = f"P{pair_count:02d}" + rotations = table.loc[[zero, ninety], "rotation_deg"].astype(float).to_numpy() + table.loc[[zero, ninety], "pair"] = pair_name + table.loc[[zero, ninety], "pair_status"] = "confident" + table.loc[[zero, ninety], "pair_reason"] = "" + table.at[zero, "partner"] = table.at[ninety, "file"] + table.at[ninety, "partner"] = table.at[zero, "file"] + table.at[zero, "pair_order"] = 0 + table.at[ninety, "pair_order"] = 1 + table.at[zero, "partner_rotation_deg"] = rotations[1] + table.at[ninety, "partner_rotation_deg"] = rotations[0] + table.at[zero, "relative_partner_rotation_deg"] = ( + rotations[1] - rotations[0] + 180.0 + ) % 360.0 - 180.0 + table.at[ninety, "relative_partner_rotation_deg"] = ( + rotations[0] - rotations[1] + 180.0 + ) % 360.0 - 180.0 + if np.isfinite(distance): + table.loc[[zero, ninety], "stage_distance_nm"] = distance * 1e9 + table.loc[[zero, ninety], "pair_tolerance_nm"] = pair_tolerance * 1e9 + + for ninety in ninety_indices: + if not table.at[ninety, "pair"] and not table.at[ninety, "pair_reason"]: + table.at[ninety, "pair_reason"] = ( + "No 0° scan has compatible shape, calibration, field of view, and stage position." + ) + + return table.sort_values("acquired", na_position="last").reset_index(drop=True) + + def _relative_difference(first: float | None, second: float | None) -> float: """Return relative difference, treating missing metadata as compatible.""" if first is None or second is None: @@ -328,18 +599,18 @@ def _same_specimen_area(first: dict, second: dict) -> bool: return False if _relative_difference(first["pixel_size_nm"], second["pixel_size_nm"]) > 0.02: return False - if _relative_difference(first["field_of_view_m"], second["field_of_view_m"]) > 0.02: + if _relative_difference(first["fov_m"], second["fov_m"]) > 0.02: return False if _relative_difference(first["magnification"], second["magnification"]) > 0.02: return False - first_stage = first["stage_position_m"] - second_stage = second["stage_position_m"] + first_stage = first["stage_xy_m"] + second_stage = second["stage_xy_m"] if first_stage is None or second_stage is None: return False distance = float(np.linalg.norm(np.subtract(first_stage, second_stage))) fields = [ float(value) - for value in (first["field_of_view_m"], second["field_of_view_m"]) + for value in (first["fov_m"], second["fov_m"]) if value is not None and float(value) > 0 ] tolerance = max(0.25 * min(fields), 10e-9) if fields else 10e-9 diff --git a/src/quantem/imaging/drift/plot.py b/src/quantem/imaging/drift/plot.py index 2ac83bde..7f8288ab 100644 --- a/src/quantem/imaging/drift/plot.py +++ b/src/quantem/imaging/drift/plot.py @@ -1,271 +1,1023 @@ +"""Standalone plot functions for :class:`~quantem.imaging.drift.DriftCorrection`. +Public functions take ``self`` as their first argument because the class binds +them as methods, for example +``DriftCorrection.plot_combined = plot.plot_combined``, so orchestration stays +in ``correction.py`` and visualization stays here. +""" import matplotlib.pyplot as plt import numpy as np +import torch +from matplotlib.axes import Axes +from matplotlib.figure import Figure +from matplotlib.ticker import MaxNLocator +import quantem.imaging.drift.apply as drift_apply +import quantem.imaging.drift.fourdstem as fourdstem from quantem.core.visualization import show_2d -from quantem.imaging.drift import diagnostics +from quantem.imaging.drift.core import knots as drift_knots +from quantem.imaging.drift.core.warping import ensure_warped_images -def plot_transformed_images(self, show_knots: bool = True, **kwargs): - self._ensure_warped_images() - fig, ax = show_2d( - list(self.images_warped.array), - **kwargs, +def show_after_step( + correction, + label: str, + *, + show_combined: bool, + show_scans: bool, + show_knots: bool, +): + """Show the requested quality-control views after one correction stage.""" + stage = _stage_description(label) + if show_combined: + correction.plot_combined( + show_knots=show_knots, + rgb=True, + title=f"Combined: {stage}", + ) + if show_scans: + titles = [] + for image_index in range(len(correction.scan_direction_degrees)): + scan_name = _scan_name(correction, image_index) + if correction._reference_mode and image_index == 0: + titles.append("Reference") + elif stage == "raw": + titles.append(f"{scan_name}: raw") + else: + titles.append(f"{scan_name} {stage}") + correction.plot_warped_images(show_knots=show_knots, title=titles) + + +def _quantile_to_uint8( + tensor: torch.Tensor, low: float = 1.0, high: float = 99.0, *, per_image: bool = False +) -> torch.Tensor: + """Percentile-stretch a torch tensor to ``uint8`` in [0, 255] on its device. + + Maps the ``[low, high]`` percentile range to [0, 255] with clamping, the + shared normalization step behind every device-rendered drift display. + ``per_image`` stretches each leading-axis image with its own percentiles; + otherwise one percentile pair scales the whole tensor. Centralizing this + keeps display contrast consistent. + """ + # torch.quantile errors above ~2**24 elements; stride-subsample large inputs + # (percentiles are stable under uniform decimation, and it stays deterministic). + def _sub(flat_1d: torch.Tensor) -> torch.Tensor: + cap = 1 << 24 + if flat_1d.shape[-1] > cap: + stride = flat_1d.shape[-1] // (cap // 2) + return flat_1d[..., ::stride] + return flat_1d + if per_image: + flat = _sub(tensor.reshape(tensor.shape[0], -1)) + lo = torch.quantile(flat, low / 100.0, dim=1) + hi = torch.quantile(flat, high / 100.0, dim=1) + scale = torch.clamp(hi - lo, min=1e-9) + norm = ((tensor - lo[:, None, None]) / scale[:, None, None]).clamp_(0, 1) + else: + flat = _sub(tensor.flatten()) + lo = torch.quantile(flat, low / 100.0) + hi = torch.quantile(flat, high / 100.0) + scale = torch.clamp(hi - lo, min=1e-9) + norm = ((tensor - lo) / scale).clamp_(0, 1) + return (norm * 255).to(torch.uint8) + + +def overlay_pair( + reference: np.ndarray, + moving: np.ndarray, + *, + mode: str = "rgb", + output_dtype: str | np.dtype | None = None, +) -> np.ndarray: + """Color overlay of two images for checking alignment. + + Default ``mode='rgb'`` is the classic red-green: ``reference`` -> red, ``moving`` -> green, + aligned regions -> yellow (reads at a glance: yellow everywhere = registered). + ``mode='green-magenta'`` is the colorblind-safe alternative (reference -> magenta, + moving -> green, aligned -> white) for published figures. Both images share one percentile + scale so the two colors match in brightness (per-image scaling would unbalance them). + + ``output_dtype="uint8"`` returns the same display colors quantized to + 0--255. Use it for large review stacks; the default remains float32 in + 0--1 for numerical plotting compatibility. + """ + ref = np.asarray(reference, dtype=np.float32) + mov = np.asarray(moving, dtype=np.float32) + lo = min(float(np.percentile(ref, 1)), float(np.percentile(mov, 1))) + hi = max(float(np.percentile(ref, 99)), float(np.percentile(mov, 99))) + def _norm(arr): + return np.clip((arr - lo) / (hi - lo + 1e-9), 0.0, 1.0) + ref, mov = _norm(ref), _norm(mov) + if mode == "rgb": + overlay = np.stack([ref, mov, np.zeros_like(ref)], -1) + else: + overlay = np.stack([ref, mov, ref], -1) + if output_dtype is None or np.dtype(output_dtype) == np.dtype(np.float32): + return overlay + if np.dtype(output_dtype) == np.dtype(np.uint8): + return np.rint(overlay * 255).astype(np.uint8) + raise ValueError( + f"output_dtype must be None, 'float32', or 'uint8', got {output_dtype!r}" + ) + + +# --------------------------------------------------------------------------- +# Public plotting API +# --------------------------------------------------------------------------- + +def plot_probe_positions( + self, + *, + image_index: int = 0, + stride: int = 16, + strip_padding: bool = True, + axsize: tuple[float, float] = (5.0, 4.8), + cmap: str = "viridis", +) -> tuple[Figure, np.ndarray]: + """Plot nominal and drift-updated probe positions for one scan image. + + The positions remain indexed like the raw scan: point ``(r, c)`` belongs + to the raw diffraction pattern at ``dataset[r, c]``. The corrected + positions are plotted in the shared drift-corrected coordinate frame, so + image 0 and image 1 position maps can be passed to ptychography without + interpolating diffraction patterns. + + Parameters + ---------- + image_index : int, default 0 + Scan whose probe positions are shown. + stride : int, default 16 + Plot every nth position along each scan axis. + strip_padding : bool, default True + Express positions in the original image frame. + axsize : tuple of float, default (5.0, 4.8) + Size of each panel in inches. + cmap : str, default "viridis" + Colormap for displacement magnitude. + + Returns + ------- + Figure + Position comparison figure. + numpy.ndarray + Matplotlib axes for the two panels. + + Examples + -------- + >>> figure, axes = drift.plot_probe_positions(image_index=1, stride=8) + """ + if stride < 1: + raise ValueError("stride must be >= 1") + nominal = self.probe_positions( + image_index=image_index, + corrected=False, + strip_padding=strip_padding, + plot=False, ) + corrected = self.probe_positions( + image_index=image_index, + corrected=True, + strip_padding=strip_padding, + plot=False, + ) + disp = corrected - nominal + mag = np.linalg.norm(disp, axis=-1) + s = (slice(None, None, stride), slice(None, None, stride)) + nominal_s = nominal[s] + corrected_s = corrected[s] + disp_s = disp[s] + mag_s = mag[s] + + fig, axes = plt.subplots(1, 2, figsize=(axsize[0] * 2, axsize[1])) + ax0, ax1 = axes + ax0.scatter( + nominal_s[..., 1].ravel(), + nominal_s[..., 0].ravel(), + s=8, + c="0.75", + label="nominal", + linewidths=0, + ) + sc0 = ax0.scatter( + corrected_s[..., 1].ravel(), + corrected_s[..., 0].ravel(), + s=10, + c=mag_s.ravel(), + cmap=cmap, + label="drift-updated", + linewidths=0, + ) + ax0.set_title(f"image {image_index} probe positions\nnominal vs drift-updated") + ax0.set_xlabel("col position (px)") + ax0.set_ylabel("row position (px)") + ax0.legend(loc="best", frameon=False) + fig.colorbar(sc0, ax=ax0, fraction=0.046, pad=0.04, label="displacement (px)") + + q = ax1.quiver( + nominal_s[..., 1], + nominal_s[..., 0], + disp_s[..., 1], + disp_s[..., 0], + mag_s, + angles="xy", + scale_units="xy", + scale=1, + cmap=cmap, + width=0.003, + ) + ax1.set_title(f"image {image_index} drift displacement\nnominal -> drift-updated") + ax1.set_xlabel("col position (px)") + ax1.set_ylabel("row position (px)") + fig.colorbar(q, ax=ax1, fraction=0.046, pad=0.04, label="displacement (px)") + for ax in axes: + ax.set_aspect("equal") + ax.invert_yaxis() + ax.grid(alpha=0.2, linewidth=0.5) + fig.tight_layout() + return fig, axes + + +def plot_warped_images( + self, + *, + show_knots: bool = True, + knot_colors: tuple[str, ...] = ("#c2185b", "#0097a7"), + axsize: tuple[int, int] = (8, 8), + max_display_px: int | None = None, + ax: Axes | None = None, + img_idx: int | None = None, + stage: str | None = None, + **kwargs, +) -> tuple[Figure, np.ndarray]: + """Inspect each scan separately to reveal scan-specific residual artifacts. + + A combined image can hide an artifact confined to one acquisition. + ``stage`` compares a saved initial, affine, strip, or non-rigid checkpoint + without changing the solved correction object. ``max_display_px=None`` + renders every native pixel; set a limit to opt into display downsampling. + + Returns + ------- + Figure + Figure containing one panel per scan. + numpy.ndarray + Matplotlib axes for the scan panels. + + Examples + -------- + >>> figure, axes = drift.plot_warped_images(stage="affine") + """ + titles = kwargs.pop("title", None) + arr_np = drift_apply.warped_stack(self, stage) + arr_t = torch.as_tensor(arr_np, device=self._device, dtype=self._dtype) + h, w = arr_t.shape[-2:] + if max_display_px and max(h, w) > max_display_px: + factor = int(max(h, w) // max_display_px) + if factor > 1: + print(f"display downsampled {factor}x (max_display_px={max_display_px}); pass max_display_px=None for native pixels") + arr_t = torch.nn.functional.avg_pool2d(arr_t.unsqueeze(0), factor).squeeze(0) + u8 = _quantile_to_uint8(arr_t, low=1.0, high=99.0, per_image=True).cpu().numpy() + stage_knots = drift_knots.stage_knots(self, stage) + n = u8.shape[0] + # single-panel mode: draw one warped image (with its origin) into a provided ax + if ax is not None: + assert img_idx is not None, "pass img_idx when drawing into a single ax" + ax.imshow(u8[img_idx], cmap="gray", vmin=0, vmax=255) + ax.set_xticks([]) + ax.set_yticks([]) + if titles is not None: + ax.set_title(titles[img_idx] if isinstance(titles, (list, tuple)) else titles) + if show_knots: + row_scale = u8.shape[-2] / self.shape[1] + column_scale = u8.shape[-1] / self.shape[2] + kn = stage_knots[img_idx].cpu().numpy() + ax.plot( + kn[1] * column_scale, + kn[0] * row_scale, + color=knot_colors[img_idx % len(knot_colors)], + ) + return ax.figure, ax + fig, ax = plt.subplots(1, n, figsize=(axsize[0] * n, axsize[1])) + if n == 1: + ax = np.array([ax]) + for i in range(n): + ax[i].imshow(u8[i], cmap="gray", vmin=0, vmax=255) + ax[i].set_xticks([]) + ax[i].set_yticks([]) + if titles is not None: + ax[i].set_title(titles[i] if isinstance(titles, (list, tuple)) else titles) if show_knots: - for a0 in range(self.shape[0]): - x = self.knots[a0][0] - y = self.knots[a0][1] - ax[a0].plot( - y, - x, - color="r", + row_scale = u8.shape[-2] / self.shape[1] + column_scale = u8.shape[-1] / self.shape[2] + for img_idx in range(self.shape[0]): + knots_np = stage_knots[img_idx].cpu().numpy() + # per-scan attribution: magenta = first scan, cyan = second (CMYK-safe, + # colorblind-safe, and distinct from the RGB comparison colors) + ax[img_idx].plot( + knots_np[1] * column_scale, + knots_np[0] * row_scale, + color=knot_colors[img_idx % len(knot_colors)], ) + return fig, ax def plot_convergence( self, - figsize=(8, 3), + *, + figsize: tuple[float, float] = (7, 4.2), + log_scale: bool = True, **kwargs, -): - """ - Plot the convergence of the drift correction. +) -> tuple[Figure, Axes]: + """Plot convergence of the iterative non-rigid refinement. + + The logarithmic loss scale is useful for judging whether later iterations + still provide a meaningful reduction. Cycle zero is the final affine value; + subsequent points are non-rigid refinement cycles. Affine alignment is a + candidate-grid search rather than an iterative optimizer, so it is not + presented as a convergence curve. + + Parameters + ---------- + figsize : tuple, default (7, 4.2) + log_scale : bool, default True + Display the mean disagreement on a logarithmic y-axis. + **kwargs + Forwarded to ``ax.plot``. + + Returns + ------- + fig : Figure + ax : Axes + + Examples + -------- + >>> figure, axis = drift.plot_convergence() """ - sub = np.abs(self.error_track[:, 0] - 2) < 0.1 - error = self.error_track[:, 1] - it = np.arange(error.shape[0]) - - from matplotlib.ticker import FormatStrFormatter, MaxNLocator - - fig, ax = plt.subplots(1, 2, figsize=figsize) - color = (1, 0, 0) # red - - # Plot Affine - if np.any(~sub): - ax[0].plot( - it[~sub], - 100 * error[~sub], - marker="o", - color=color, - linestyle="-", - label="Affine", - **kwargs, + track = np.asarray(self.error_track) + is_nonrigid = track[:, 0] == 2 + error_percent = 100 * track[:, 1] + affine_error = error_percent[~is_nonrigid] + nonrigid_error = error_percent[is_nonrigid] + if nonrigid_error.size == 0: + raise RuntimeError( + "plot_convergence() requires a completed correct_nonrigid() call." ) - ax[0].set_xlabel("Affine Iterations") - ax[0].set_ylabel("Mean Error [%]") - ax[0].xaxis.set_major_locator(MaxNLocator(integer=True)) - ax[0].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) - else: - ax[0].axis("off") - - # Plot Non-Rigid - if np.any(sub): - first_true = np.argmax(sub) - if first_true > 0: - sub[first_true - 1] = True - - ax[1].plot( - it[sub], - 100 * error[sub], - marker="o", - color=color, - linestyle="-", - label="Non-Rigid", - **kwargs, - ) - ax[1].set_xlabel("Non-Rigid Iterations") - ax[1].xaxis.set_major_locator(MaxNLocator(integer=True)) - ax[1].yaxis.set_major_formatter(FormatStrFormatter("%.4f")) - else: - ax[1].axis("off") + if affine_error.size: + nonrigid_error = np.concatenate((affine_error[-1:], nonrigid_error)) - plt.tight_layout() + fig, ax = plt.subplots(figsize=figsize) + color = "#0072B2" + plot_kwargs = {"color": color, "linewidth": 2.0, **kwargs} + cycles = np.arange(nonrigid_error.size) + ax.plot(cycles, nonrigid_error, **plot_kwargs) + ax.scatter( + cycles[[0, -1]], nonrigid_error[[0, -1]], + s=28, color=color, edgecolor="white", linewidth=0.7, zorder=3, + ) + ax.set_title("Non-rigid convergence", fontsize=11, fontweight="semibold") + ax.set_xlabel("Refinement cycle") + ax.set_ylabel("Mean disagreement (%)") + ax.xaxis.set_major_locator(MaxNLocator(integer=True, nbins=7)) + if log_scale: + ax.set_yscale("log") + ax.grid(True, which="major", color="#d0d0d0", linewidth=0.7) + ax.grid(True, which="minor", color="#e8e8e8", linewidth=0.5) + ax.spines[["top", "right"]].set_visible(False) + fig.tight_layout() + return fig, ax - return self +# --------------------------------------------------------------------------- +# Building blocks - used internally by the public API functions above +# --------------------------------------------------------------------------- -def _ensure_warped_images(self): - """Lazily populate images_warped from current knots if marked stale.""" - if getattr(self, "_images_warped_stale", False): - self._warp_and_translate_torch( - self._max_image_shift_cached, upsample_factor=8, - solve_translation=False) - self._images_warped_stale = False +def _render_uint8( + dc, + rgb: bool, + low: float = 1.0, + high: float = 99.0, + max_display_px: int | None = None, + stack: np.ndarray | None = None, +) -> np.ndarray: + """Normalize a corrected stack for fast grayscale or RGB display.""" + arr_np = drift_apply.warped_stack(dc) if stack is None else stack + arr_t = torch.as_tensor(arr_np, device=dc._device, dtype=dc._dtype) + h, w = arr_t.shape[-2:] + if max_display_px and max(h, w) > max_display_px: + factor = int(max(h, w) // max_display_px) + if factor > 1: + arr_t = torch.nn.functional.avg_pool2d(arr_t.unsqueeze(0), factor).squeeze(0) + if rgb and arr_t.shape[0] <= 3: + H, W = arr_t.shape[-2:] + rgb_t = torch.zeros(H, W, 3, device=arr_t.device, dtype=torch.uint8) + for i in range(arr_t.shape[0]): + rgb_t[..., i] = _quantile_to_uint8(arr_t[i], low=low, high=high) + return rgb_t.cpu().numpy() + merged = arr_t.mean(0) + return _quantile_to_uint8(merged, low=low, high=high).cpu().numpy() -def plot_merged_images(self, show_knots: bool = True, **kwargs): - """ - Plot the current transformed images, with knot overlays. - """ - self._ensure_warped_images() - fig, ax = show_2d( - self.images_warped.array.mean(0), - **kwargs, - ) - if show_knots: - for a0 in range(self.shape[0]): - x = self.knots[a0][0] - y = self.knots[a0][1] - ax.plot( - y, - x, - ) +def _single_image_u8( + dc, + idx: int, + max_display_px: int | None, + stack: np.ndarray | None = None, +) -> np.ndarray: + """One warped image, quantile-stretched to uint8 (same normalize as the merged view).""" + arr_np = drift_apply.warped_stack(dc) if stack is None else stack + arr_t = torch.as_tensor(arr_np[idx], device=dc._device, dtype=dc._dtype) + h, w = arr_t.shape[-2:] + if max_display_px and max(h, w) > max_display_px and int(max(h, w) // max_display_px) > 1: + arr_t = torch.nn.functional.avg_pool2d(arr_t[None, None], int(max(h, w) // max_display_px))[0, 0] + return _quantile_to_uint8(arr_t).cpu().numpy() -def _comparison_pair(stack: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Compare scan 0 with scan 1, or with the mean of all remaining scans.""" - if stack.shape[0] == 2: - return stack[0], stack[1] - return np.mean(stack[1:], axis=0), stack[0] +def _stage_description(stage: str | None) -> str: + """Return a concise, natural-language description of a correction stage.""" -def _registration_overlay( - reference: np.ndarray, - moving: np.ndarray, - mask: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - """Build percentile-normalized RGB agreement and difference images.""" - values = np.concatenate((reference[mask], moving[mask])) - if values.size: - low, high = np.percentile(values, (1.0, 99.0)) - else: - low, high = 0.0, 1.0 - scale = max(float(high - low), np.finfo(np.float32).eps) - first = np.clip((reference - low) / scale, 0.0, 1.0) - second = np.clip((moving - low) / scale, 0.0, 1.0) - overlay = np.stack((first, second, np.zeros_like(first)), axis=-1) - difference = np.abs(first - second) - overlay[~mask] = 0.0 - difference[~mask] = 0.0 - return overlay, difference - - -def _plot_registration_diagnostics( - correction, + return { + None: "current correction", + "initial": "raw", + "raw": "raw", + "affine": "after affine correction", + "strip": "after strip correction", + "nonrigid": "after non-rigid correction", + "non-rigid": "after non-rigid correction", + }.get(stage, str(stage)) + + +def _format_angle_degrees(angle: float) -> str: + """Format a scan angle with a mathematical minus and degree sign.""" + + value = float(angle) + sign = "−" if value < 0 else "" + return f"{sign}{abs(value):g}°" + + +def _scan_name(dc, image_index: int) -> str: + """Return the scientist-facing role or angle for one scan.""" + + if getattr(dc, "_reference_mode", False): + return "Reference" if image_index == 0 else "Moving scan" + angle = _format_angle_degrees(dc.scan_direction_degrees[image_index]) + return f"{angle} scan" + + +def plot_combined( + self, *, - stages: tuple[str, ...] | list[str] | None = None, - figsize: tuple[float, float] | None = None, -): - """Plot stage-by-stage registration on one fixed measured footprint.""" - selected, stacks, common_mask, _ = diagnostics._registration_data( - correction, - stages, + show_knots: bool = True, + rgb: bool = True, + axsize: tuple[int, int] = (16, 16), + max_display_px: int | None = None, + ax: Axes | None = None, + stage: str | tuple[str, ...] | list[str] | None = None, + show_scans: bool = False, + interactive: bool = False, + width: int | None = 1000, + mode: str = "rgb", + zoom: float = 1.0, + center: tuple[float, float] | None = None, + display_bin: int = 1, + **kwargs, +) -> tuple[Figure, Axes] | object: + """Plot the combined scan (RGB registration overlay) with optional knot curves. + + Bound on :class:`~quantem.imaging.drift.DriftCorrection` as + ``DriftCorrection.plot_combined``. This is the primary registration QA plot + in tutorials (yellow means aligned in the red/green overlay). + + Parameters + ---------- + show_knots : bool, default True + Overlay knot trajectories on static non-rigid panels. Initial and + affine panels stay clean because multiple knots add freedom only to + the non-rigid model. Use :meth:`plot_knots` to inspect affine geometry. + rgb : bool, default True + If True (2 or 3 scans), RGB comparison (image 0 red, 1 green, 2 blue). + Misalignment shows as color fringes. If False, grayscale mean. + axsize : tuple of int, default (16, 16) + Matplotlib figure size scale for the static path. + max_display_px : int or None + Cap longest display edge for static rendering. None keeps native size. + ax : matplotlib.axes.Axes or None + Optional axis to draw into (static path). + stage : str, tuple of str, or None + Knot checkpoint to render. Examples: ``"affine"``, ``"nonrigid"``, + ``("initial", "affine")``, ``("initial", "affine", "nonrigid")``. + A tuple draws stages side by side in one call. + show_scans : bool, default False + Static path only: also show each warped scan beside the combined panel. + Ignored when ``ax`` is set. + interactive : bool, default False + If True, return a :class:`quantem.widget.Show2D` gallery instead of a + static figure. + width : int or None, default 1000 + Interactive panel width in pixels. Ignored on the static path. + mode : str, default "rgb" + Overlay mode for interactive RGB (e.g. ``"rgb"``, green-magenta variants). + zoom : float, default 1.0 + Static path: axis-limit zoom (not array crop). Interactive path: Show2D zoom. + center : tuple of float or None + Static path: fractional ``(row, col)`` pan in ``[-0.5, 0.5]``. Interactive + path: absolute pixel center when provided. + display_bin : int, default 1 + Interactive Show2D display binning only. Default 1 (no silent reduction). + Pass 2/4/8 for speed, or ``"auto"`` for Show2D budget behavior. + **kwargs + Forwarded to the static renderer or Show2D (for example ``show_fft``). + + Returns + ------- + (Figure, Axes) or Show2D + Static matplotlib objects, or an interactive widget when + ``interactive=True``. + + Notes + ----- + Interactive RGB bakes contrast into the overlay. Do not pass + ``auto_contrast`` / ``link_contrast`` for that mode; they are ignored. + ``rgb=False`` still forwards those Show2D contrast controls. + + Examples + -------- + >>> dc.plot_combined(stage=("initial", "affine"), interactive=True, width=620) + >>> dc.plot_combined( + ... stage=("affine", "nonrigid"), + ... interactive=False, + ... show_knots=True, + ... ) + """ + if interactive: + from quantem.widget import Show2D # lazy: quantem.widget is an optional extra + + stage_names = list(stage) if isinstance(stage, (tuple, list)) else [stage] + if rgb and self.shape[0] != 2: + raise ValueError( + f"interactive rgb=True needs exactly 2 images for the green-magenta overlay, got {self.shape[0]}" + ) + images, labels = [], [] + for stage_name in stage_names: + stage_display = _stage_description(stage_name) + warped = drift_apply.warped_stack(self, stage_name) + # One panel per stage: RGB comparison or grayscale mean. + if rgb: + images.append(overlay_pair(warped[0], warped[1], mode=mode)) + labels.append(f"Combined: {stage_display} (RGB)") + else: + images.append(warped.mean(axis=0)) + labels.append(f"Combined: {stage_display}") + sampling = float(np.asarray(self.imgs[0].sampling, dtype=float)[0]) + unit = self.imgs[0].units[0] if getattr(self.imgs[0], "units", None) else "px" + show_kwargs = dict(kwargs) + # A combined registration comparison is a clean scientific result, + # not a general-purpose image-editing workspace. Callers can still opt + # back into the full interface with ui_mode="interactive". + show_kwargs.setdefault("ui_mode", "report") + if rgb: + # Overlay is already percentile-normalized on a shared scale. + # Contrast knobs would only confuse an RGB registration view. + show_kwargs.pop("auto_contrast", None) + show_kwargs.pop("link_contrast", None) + show_kwargs.pop("cmap", None) + show_kwargs.setdefault("show_fft", False) + return Show2D( + images, + labels=labels, + sampling=sampling, + units=unit, + ncols=len(stage_names), + size=int(width) if width else 0, + display_bin=display_bin, + zoom=zoom, + center=center, + **show_kwargs, + ) + if isinstance(stage, (tuple, list)): + # One call draws every requested stage side by side. Pass ax= as a + # sequence of Axes (one per stage) to embed into a larger gallery; + # otherwise a fresh 1 x n figure is created. + stage_names = list(stage) + if ax is None: + fig, axes = plt.subplots( + 1, + len(stage_names), + figsize=(axsize[0] * len(stage_names) / 2, axsize[1] / 2), + squeeze=False, + ) + stage_axes = list(axes[0]) + else: + stage_axes = list(np.atleast_1d(ax).ravel()) + if len(stage_axes) != len(stage_names): + raise ValueError( + f"stage has {len(stage_names)} entries but ax has " + f"{len(stage_axes)} axes; provide one Axes per stage." + ) + fig = stage_axes[0].figure + for stage_name, stage_ax in zip(stage_names, stage_axes, strict=True): + plot_combined( + self, + show_knots=show_knots, + rgb=rgb, + max_display_px=max_display_px, + ax=stage_ax, + stage=stage_name, + zoom=zoom, + center=center, + **dict(kwargs), + ) + return fig, np.asarray(stage_axes, dtype=object) + stage_display = _stage_description(stage) + display_stack = drift_apply.warped_stack(self, stage) + merged_u8 = _render_uint8( + self, + rgb=rgb and self.shape[0] <= 3, + max_display_px=max_display_px, + stack=display_stack, ) - if figsize is None: - figsize = (11.0, 3.4 * len(selected)) - figure, axes = plt.subplots( - len(selected), - 3, - figsize=figsize, - squeeze=False, + merged_title = kwargs.pop( + "title", + f"Combined: {stage_display}" + (" (RGB)" if rgb else ""), ) - for row_index, stage in enumerate(selected): - reference, moving = _comparison_pair(stacks[stage]) - metrics = diagnostics._pair_metrics(reference, moving, common_mask) - overlay, difference = _registration_overlay(reference, moving, common_mask) - axes[row_index, 0].imshow(overlay) - axes[row_index, 0].set_title( - f"{stage}: RGB registration\ncommon NCC {float(metrics['common_ncc']):.4f}" - ) - axes[row_index, 1].imshow(difference, cmap="magma", vmin=0.0, vmax=1.0) - axes[row_index, 1].set_title( - "Percentile-normalized |difference|\n" - f"native MAD {float(metrics['mean_absolute_difference']):.4g}; " - f"RMS {float(metrics['root_mean_square_difference']):.4g}" + kwargs.pop("cmap", None) + panels = [] + if show_scans and ax is None: + panels = [ + ( + _single_image_u8(self, i, max_display_px, stack=display_stack), + f"{_scan_name(self, i)} {stage_display}" + if stage_display != "raw" + else f"{_scan_name(self, i)}: raw", + ) + for i in range(self.shape[0]) + ] + panels.append((merged_u8, merged_title)) + if ax is None: + fig, axes = plt.subplots( + 1, + len(panels), + figsize=(axsize[0] * len(panels) / 2, axsize[1] / 2), + squeeze=False, ) - axes[row_index, 2].imshow(common_mask, cmap="gray", vmin=0, vmax=1) - axes[row_index, 2].set_title( - f"Fixed common coverage\n{float(common_mask.mean()):.1%} of canvas" + axes = list(axes[0]) + else: + fig, axes = ax.figure, [ax] + for panel_ax, (u8, panel_title) in zip(axes, panels): + if u8.ndim == 3: + panel_ax.imshow(u8) + else: + panel_ax.imshow(u8, cmap="gray", vmin=0, vmax=255) + panel_ax.set_title(panel_title) + panel_ax.set_xticks([]) + panel_ax.set_yticks([]) + merged_ax = axes[-1] + nonrigid_stage = stage in ("nonrigid", "non-rigid") + if stage is None: + error_track = np.asarray(getattr(self, "error_track", [])) + nonrigid_stage = bool( + error_track.ndim == 2 + and error_track.shape[1] > 0 + and np.any(error_track[:, 0] == 2) ) - for axis in axes[row_index]: - axis.set_xticks([]) - axis.set_yticks([]) - figure.suptitle( - "Drift registration diagnostics", - fontsize=14, - fontweight="semibold", - ) - figure.tight_layout() - return figure, axes + if show_knots and nonrigid_stage: + stage_knots = drift_knots.stage_knots(self, stage) + row_scale = merged_u8.shape[0] / self.shape[1] + column_scale = merged_u8.shape[1] / self.shape[2] + for knots in stage_knots: + knots_np = knots.cpu().numpy() + merged_ax.plot( + knots_np[1] * column_scale, + knots_np[0] * row_scale, + ) + if zoom > 1.0: + # Axis limits preserve the overlay-to-pixel coordinate mapping. + panel_h, panel_w = merged_u8.shape[:2] + win = min(panel_h, panel_w) / float(zoom) + offset_row, offset_col = center if center is not None else (0.0, 0.0) + cy = min(max(panel_h / 2 + offset_row * panel_h, win / 2), panel_h - win / 2) + cx = min(max(panel_w / 2 + offset_col * panel_w, win / 2), panel_w - win / 2) + for panel_ax in axes: + panel_ax.set_xlim(cx - win / 2, cx + win / 2) + panel_ax.set_ylim(cy + win / 2, cy - win / 2) + return fig, (merged_ax if len(axes) == 1 else np.array(axes)) -def _plot_displacement_diagnostics( - correction, +def plot_knots( + self, *, + figsize: tuple[int, int] | None = None, stage: str | None = None, - figsize: tuple[float, float] | None = None, -): - """Plot scan-line-origin displacement without changing fitted knots.""" - if stage is None: - stage = diagnostics._select_stages(correction, None)[-1] - else: - stage = diagnostics._select_stages(correction, (stage,))[0] - fields = diagnostics._displacement_fields(correction, stage) - rows = diagnostics._displacement_rows(correction, (stage,)) +) -> tuple[Figure, np.ndarray]: + """Plot knot trajectories before and after correction plus the per-scanline delta field. + + Two panels per image: + - Top: mean warped image with initial knots (dashed) and corrected knots (solid) + overlaid. A third dotted line shows the affine-only state when available, so + the affine vs. nonrigid contributions are visible side by side. + - Bottom: per-scanline correction delta (row and col components) in pixels. + A smooth curve means the correction field is physically reasonable. Rapid + oscillations indicate ``knot_smoothing_sigma`` is too small and the optimizer + is fitting noise rather than real drift. + + Parameters + ---------- + figsize : (width, height), optional + Figure size in inches. Defaults to ``(7 * num_images, 11)`` so each + square image occupies the same column width as its correction chart. + stage : {"affine", "strip", "nonrigid", None}, optional + Show only one stage's contribution. ``"nonrigid"`` plots just the + nonrigid delta on its OWN y-scale - the refinement is often sub-pixel + and disappears next to affine's tens-of-pixels ramp, yet it is the + parameter that matters when judging the nonrigid solve. ``"strip"`` + plots the piecewise-rigid contribution relative to affine, + ``"affine"`` plots just the affine ramp, and ``None`` overlays every + completed stage. + + Returns + ------- + fig : Figure + axes : np.ndarray of Axes, shape (2, num_images) + + Examples + -------- + >>> figure, axes = drift.plot_knots(stage="nonrigid") + """ + ensure_warped_images(self) + num_images = self.shape[0] + warped = self.imgs_warped.array # each panel shows its own corrected image, not the merge + colors = plt.rcParams["axes.prop_cycle"].by_key()["color"] if figsize is None: - figsize = (11.0, 3.6 * len(fields)) - figure, axes = plt.subplots( - len(fields), - 2, - figsize=figsize, - squeeze=False, + figsize = (7 * num_images, 11) + fig, axes = plt.subplots( + 2, num_images, figsize=figsize, + gridspec_kw={"height_ratios": [2, 1]}, ) - for image_index, (field, metrics) in enumerate(zip(fields, rows, strict=True)): - scan_lines = np.arange(field.shape[1]) - row_displacement = np.mean(field[0], axis=1) - column_displacement = np.mean(field[1], axis=1) - axes[image_index, 0].plot( - scan_lines, - row_displacement, - label="row displacement", + if num_images == 1: + axes = axes[:, None] + scanlines = np.arange(self.knots[0].shape[1]) + for img_idx in range(num_images): + color = colors[img_idx % len(colors)] + initial = self._initial_knots[img_idx].cpu().numpy() + current = self.knots[img_idx].cpu().numpy() + affine_np = ( + self._knots_after_affine[img_idx].cpu().numpy() + if hasattr(self, "_knots_after_affine") + else None ) - axes[image_index, 0].plot( - scan_lines, - column_displacement, - label="column displacement", + strip_np = ( + self._knots_after_strip[img_idx].cpu().numpy() + if hasattr(self, "_knots_after_strip") + else None ) - axes[image_index, 0].axhline(0.0, color="0.3", linewidth=0.7) - axes[image_index, 0].set_xlabel("Scan line") - axes[image_index, 0].set_ylabel("Displacement (px)") - axes[image_index, 0].set_title( - f"Image {image_index}: scan-line origins\n" - f"endpoint displacement {float(metrics['endpoint_displacement_px']):.3g} px" + has_nonrigid = bool(np.any(np.asarray(self.error_track)[:, 0] == 2)) + ax_img = axes[0, img_idx] + bg = warped[img_idx] + ax_img.imshow(bg, cmap="gray", origin="upper", aspect="equal", + vmin=float(bg.min()), vmax=float(bg.max())) + ax_img.plot(initial[1, :, 0], initial[0, :, 0], + "--", color=color, lw=1.2, alpha=0.7, label="before alignment") + if affine_np is not None and stage in (None, "affine"): + ax_img.plot(affine_np[1, :, 0], affine_np[0, :, 0], + ":", color=color, lw=1.0, alpha=0.6, + label="after affine correction") + if strip_np is not None and stage in (None, "strip"): + ax_img.plot(strip_np[1, :, 0], strip_np[0, :, 0], + "-", color=color, lw=1.5, + label="after strip correction") + if has_nonrigid and stage in (None, "nonrigid"): + ax_img.plot(current[1, :, 0], current[0, :, 0], + "-", color=color, lw=1.5, + label="after non-rigid correction") + ax_img.legend(fontsize=8, loc="upper right") + scan_name = _scan_name(self, img_idx) + ax_img.set_title(f"{scan_name}: knot trajectory", fontsize=10) + ax_img.axis("off") + ax_delta = axes[1, img_idx] + if affine_np is not None and stage in (None, "affine"): + aff_delta = affine_np - initial + ax_delta.plot(scanlines, aff_delta[0, :, 0], + ":", lw=1.0, alpha=0.6, label="row \u0394 after affine") + ax_delta.plot(scanlines, aff_delta[1, :, 0], + ":", lw=1.0, alpha=0.6, label="col \u0394 after affine") + # Each residual stage is measured from its immediate predecessor so + # the smaller strip/non-rigid contribution remains readable. + if strip_np is not None and stage in (None, "strip"): + strip_base = affine_np if affine_np is not None else initial + strip_delta = strip_np - strip_base + ax_delta.plot(scanlines, strip_delta[0, :, 0], lw=1.2, + label="row Δ from strip correction") + ax_delta.plot(scanlines, strip_delta[1, :, 0], lw=1.2, + label="column Δ from strip correction") + if has_nonrigid and stage in (None, "nonrigid"): + nonrigid_base = strip_np if strip_np is not None else affine_np + if nonrigid_base is None: + nonrigid_base = initial + nonrigid_delta = current - nonrigid_base + ax_delta.plot(scanlines, nonrigid_delta[0, :, 0], lw=1.2, + label="row Δ from non-rigid correction") + ax_delta.plot(scanlines, nonrigid_delta[1, :, 0], lw=1.2, + label="column Δ from non-rigid correction") + ax_delta.axhline(0, color="k", lw=0.5, ls="--") + ax_delta.set_xlabel("scan line") + ax_delta.set_ylabel("correction (pixels)") + ax_delta.set_title(f"{scan_name}: correction field", fontsize=10) + handles, labels = ax_delta.get_legend_handles_labels() + if handles: + ax_delta.legend(handles, labels, fontsize=8) + ax_delta.grid(alpha=0.3) + fig.tight_layout() + return fig, axes + + +# --------------------------------------------------------------------------- +# Complete workflow views +# --------------------------------------------------------------------------- + +def static_comparison(correction, zoom: float, size: int, **kwargs): + """Compare raw and corrected scans in a publication-ready figure. + + A shared contrast range makes intensity differences visible across scans; + independent auto contrast is useful when acquisition brightness differs. + Zoom changes only the displayed field of view, never the corrected data. + """ + panels = drift_apply.comparison_panels(correction) + cmap = kwargs.pop("cmap", "inferno") + auto_contrast = bool(kwargs.pop("auto_contrast", False)) + link_contrast = bool(kwargs.pop("link_contrast", True)) + smooth = bool(kwargs.pop("smooth", False)) + kwargs.pop("show_fft", None) + vmin = kwargs.pop("vmin", None) + vmax = kwargs.pop("vmax", None) + offset_row, offset_col = kwargs.pop("center", (0.0, 0.0)) or (0.0, 0.0) + + images = [np.asarray(image) for image in panels["images"]] + ncols = int(panels["ncols"]) + image_grid = [images[index : index + ncols] for index in range(0, len(images), ncols)] + label_grid = [ + panels["labels"][index : index + ncols] + for index in range(0, len(images), ncols) + ] + + norm = None + if link_contrast: + values = np.concatenate([image.ravel()[::16] for image in images]) + values = values[np.isfinite(values)] + low, high = ( + np.percentile(values, (1.0, 99.0)) + if auto_contrast + else (values.min(), values.max()) ) - axes[image_index, 0].legend(frameon=False) - axes[image_index, 0].grid(alpha=0.25) - - magnitude = np.linalg.norm(field, axis=0) - if magnitude.shape[1] == 1: - axes[image_index, 1].plot( - scan_lines, - magnitude[:, 0], - color="#0072B2", - ) - axes[image_index, 1].set_xlabel("Scan line") - axes[image_index, 1].set_ylabel("Displacement magnitude (px)") - axes[image_index, 1].grid(alpha=0.25) - else: - image = axes[image_index, 1].imshow( - magnitude, - cmap="viridis", - aspect="auto", - origin="upper", - ) - axes[image_index, 1].set_xlabel("Knot along fast scan") - axes[image_index, 1].set_ylabel("Scan line") - figure.colorbar( - image, - ax=axes[image_index, 1], - label="Displacement (px)", - ) - axes[image_index, 1].set_title( - f"Image {image_index}: displacement magnitude\n" - f"maximum {float(metrics['max_displacement_px']):.3g} px" + norm = { + "vmin": float(low if vmin is None else vmin), + "vmax": float(high if vmax is None else vmax), + } + elif auto_contrast: + norm = "linear_auto" + elif vmin is not None or vmax is not None: + norm = {"vmin": vmin, "vmax": vmax} + + figure, axes = show_2d( + image_grid, + title=label_grid, + cmap=cmap, + norm=norm, + scalebar={"sampling": panels["pixel_size"], "units": panels["pixel_unit"]}, + axsize=(max(float(size), 200.0) / 100.0,) * 2, + **kwargs, + ) + axes_array = np.asarray(axes, dtype=object).reshape(len(image_grid), ncols) + interpolation = "bilinear" if smooth else "nearest" + for axis, image in zip(axes_array.flat, images, strict=True): + if axis.images: + axis.images[0].set_interpolation(interpolation) + height, width = image.shape[:2] + visible_height, visible_width = height / zoom, width / zoom + center_row = np.clip( + (height - 1) / 2 + offset_row * height, + visible_height / 2.0, + height - visible_height / 2.0, ) - figure.suptitle( - f"Drift displacement diagnostics: {stage}", - fontsize=14, - fontweight="semibold", + center_column = np.clip( + (width - 1) / 2 + offset_col * width, + visible_width / 2.0, + width - visible_width / 2.0, + ) + axis.set_xlim( + center_column - visible_width / 2, + center_column + visible_width / 2, + ) + axis.set_ylim( + center_row + visible_height / 2, + center_row - visible_height / 2, + ) + return figure + + +def show( + self, + *, + zoom: float = 1.0, + size: int = 400, + mode: str = "interactive", + **kwargs, +): + """Show raw and corrected scans in their common acquisition frame. + + Interactive modes return a linked Show2D comparison for inspecting local + alignment. ``mode="static"`` returns the same panels as a Matplotlib figure + for papers and saved reports. + + Parameters + ---------- + zoom : float, default 1.0 + Initial magnification of the common scan field. + size : int, default 400 + Interactive panel size in pixels. + mode : str, default "interactive" + Use ``"static"`` for Matplotlib or a Show2D UI mode for exploration. + + Returns + ------- + object + Show2D viewer or Matplotlib figure selected by ``mode``. + + Examples + -------- + >>> viewer = drift.show(zoom=2) + >>> figure = drift.show(mode="static") + """ + if mode == "static": + return static_comparison(self, zoom, size, **kwargs) + + from quantem.widget import Show2D + + panels = drift_apply.comparison_panels(self) + return Show2D( + panels["images"], + ncols=panels["ncols"], + gallery_gap_px=0, + size=size, + pixel_size=panels["pixel_size"], + pixel_unit=panels["pixel_unit"], + labels=panels["labels"], + zoom=zoom, + ui_mode=mode, + **kwargs, + ) + + +def show_4dstem( + self, + *, + det_bin: int = 1, + view_mode: str = "multiple", + **kwargs, +): + """Show raw, corrected, and merged diffraction datasets together. + + The three synchronized stages make it possible to inspect whether scan + correction improves the virtual image without changing diffraction detail. + + Parameters + ---------- + det_bin : int, default 1 + Detector-axis binning used only for interactive display. + view_mode : str, default "multiple" + Show4DSTEM layout used for the processing stages. + + Returns + ------- + Show4DSTEM + Interactive comparison viewer. + + Examples + -------- + >>> viewer = drift.show_4dstem(det_bin=2) + """ + from quantem.widget import Show4DSTEM + + sampling = np.asarray(self.imgs[0].sampling, dtype=float) + units = list(self.imgs[0].units) + kwargs.setdefault( + "frame_labels", + ["0° raw", "0° affine-corrected", "0°/90° affine-corrected + combined"], + ) + kwargs.setdefault("title", "4D-STEM affine drift correction") + kwargs.setdefault("frame_dim_label", "Processing stage") + kwargs.setdefault("compare_cols", 3) + kwargs.setdefault("compare_max_panels", 3) + kwargs.setdefault("compare_dp_mode", "selected") + kwargs.setdefault( + "sampling", + (float(sampling[0]), float(sampling[1]), float(det_bin), float(det_bin)), + ) + kwargs.setdefault("units", [units[0], units[1], "pixels", "pixels"]) + + viewer = Show4DSTEM( + fourdstem.corrected_4dstem_views(self, det_bin=det_bin), + view_mode=view_mode, + **kwargs, ) - figure.tight_layout() - return figure, axes + viewer.frame_idx = 2 + return viewer diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preparation.py index 71c49dcf..a2227328 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preparation.py @@ -1,44 +1,171 @@ +"""Image and canvas preparation for drift alignment.""" + +from collections.abc import Sequence import numpy as np import torch from numpy.typing import NDArray +import quantem.imaging.drift.plot as drift_plot +import quantem.imaging.drift.report as report from quantem.core.datastructures.dataset2d import Dataset2d from quantem.core.datastructures.dataset3d import Dataset3d from quantem.core.utils.compound_validators import ( + validate_list_of_dataset2d, validate_pad_value, ) -from quantem.core.utils.imaging_utils import ( - cross_correlation_shift, -) -from quantem.imaging.drift import diagnostics -from quantem.imaging.drift.core.knots import ( - DriftInterpolator, - bilinear_kde_batch, - transform_coordinates_single_knot, -) +from quantem.core.utils.validators import ensure_valid_array +from quantem.imaging.drift.core import knots as drift_knots + + +def input_array(value): + """Return the array carried by a drift-correction input.""" + if isinstance(value, (np.ndarray, torch.Tensor)): + return value + array = getattr(value, "array", None) + if isinstance(array, np.ndarray): + return array + raise TypeError( + f"DriftCorrection accepts ndarray, torch.Tensor, or Dataset " + f"objects; got {type(value).__name__}. To load from disk, call " + f"Dataset2d.from_file(path) (or Dataset4d.from_file) first." + ) + + +def prepare_image_collection( + correction, + arrays, + scan_direction_degrees, + source_datasets=None, +): + """Preserve scan calibration while preparing 2-D alignment images.""" + correction.imgs = validate_list_of_dataset2d(arrays) + if source_datasets is not None: + for image, source in zip(correction.imgs, source_datasets): + for name in ("origin", "sampling", "units"): + value = getattr(source, name, None) + if value is not None: + setattr(image, name, value[:2]) + signal_units = getattr(source, "signal_units", None) + if signal_units is not None: + image.signal_units = signal_units + metadata = getattr(source, "metadata", None) + if isinstance(metadata, dict): + image.metadata.update(metadata) + correction.scan_direction_degrees = ensure_valid_array( + scan_direction_degrees, ndim=1 + ) + + +def prepare_inputs( + correction, + datasets, + scan_direction_degrees, + alignment_image, +): + """Prepare image, reference, or 4D-STEM inputs for one correction state.""" + if len(datasets) < 2: + raise TypeError( + f"DriftCorrection requires at least 2 datasets, got {len(datasets)}" + ) + + if scan_direction_degrees is None: + found = [ + getattr(dataset, "metadata", {}).get("scan_rotation_deg") + if isinstance(getattr(dataset, "metadata", None), dict) + else None + for dataset in datasets + ] + if not all(angle is not None for angle in found): + raise TypeError( + "scan_direction_degrees is required: the inputs carry no " + "scan-angle metadata. Load Velox files with " + "em.imaging.read_emd(path) (which stamps " + "metadata['scan_rotation_deg']) or pass the angles " + "explicitly, e.g. scan_direction_degrees=(0, 90)." + ) + scan_direction_degrees = found + if np.isscalar(scan_direction_degrees): + angles = [float(scan_direction_degrees)] * len(datasets) + else: + angles = [float(angle) for angle in scan_direction_degrees] + if len(angles) != len(datasets): + raise ValueError( + f"scan_direction_degrees length ({len(angles)}) must match " + f"number of datasets ({len(datasets)})" + ) + + arrays = [input_array(dataset) for dataset in datasets] + dimensions = [array.ndim for array in arrays] + for index, ndim in enumerate(dimensions): + if ndim < 2: + raise TypeError(f"dataset {index} must be ≥2-D, got ndim={ndim}") + + if len(arrays) == 2: + first, second = arrays + first_ndim, second_ndim = dimensions + if first_ndim == 2 and (second_ndim >= 3 or angles[0] == angles[1]): + if first.shape != second.shape[:2]: + raise ValueError( + f"reference shape {first.shape} must match the leading " + f"two axes of drifted (got {second.shape[:2]})" + ) + if alignment_image is not None: + moving = np.asarray(alignment_image) + elif second_ndim == 2: + moving = second + else: + from quantem.imaging.drift.fourdstem import integrate_virtual_detector + + moving = integrate_virtual_detector(second) + prepare_image_collection(correction, [first, moving], angles) + correction._reference_mode = True + correction._datasets = [None, second] + correction._built_from_datasets = True + return + + if first_ndim >= 4 and second_ndim >= 4: + from quantem.imaging.drift.fourdstem import integrate_virtual_detector + + prepare_image_collection( + correction, + [ + integrate_virtual_detector(first), + integrate_virtual_detector(second), + ], + angles, + datasets, + ) + correction._datasets = [first, second] + correction._built_from_datasets = True + return + + sources = datasets if first_ndim == second_ndim == 2 else None + prepare_image_collection(correction, arrays, angles, sources) + return + + prepare_image_collection(correction, arrays, angles, datasets) def validate_downsample(value: int) -> int: - """Validate one integer computational downsampling factor.""" + """Validate a computational downsampling factor.""" factor = int(value) - if factor < 1 or factor != value: - raise ValueError(f"downsample must be a positive integer, got {value!r}.") + if factor < 1: + raise ValueError(f"downsample must be >= 1, got {value!r}.") return factor def average_downsample_2d(array: NDArray, factor: int) -> NDArray: - """Average-pool a two-dimensional image by an exact integer factor.""" + """Average-pool a two-dimensional image by an integer factor.""" image = np.asarray(array) - factor = validate_downsample(factor) if factor == 1: return np.ascontiguousarray(image) if image.ndim != 2: raise ValueError( f"downsample currently supports 2-D images, got ndim={image.ndim}." ) - rows, columns = image.shape - if rows % factor or columns % factor: + rows, cols = image.shape + if rows % factor or cols % factor: raise ValueError( f"downsample={factor} requires image dimensions divisible by " f"{factor}, got {image.shape}." @@ -46,24 +173,203 @@ def average_downsample_2d(array: NDArray, factor: int) -> NDArray: downsampled = image.reshape( rows // factor, factor, - columns // factor, + cols // factor, factor, ).mean(axis=(1, 3)) - dtype = image.dtype if np.issubdtype(image.dtype, np.floating) else np.float32 + dtype = ( + image.dtype + if np.issubdtype(image.dtype, np.floating) + else np.float32 + ) return np.ascontiguousarray(downsampled.astype(dtype, copy=False)) +def match_scan_shapes(first, second, *, verbose: bool = False): + """Mean-bin an integer-resolution pair to the same physical scan grid. + + Orthogonal scans of the same field can be saved at different integer + sampling densities. Mean-binning the finer scan preserves the measured + field and gives the correction one shared pixel grid; non-integer shape + mismatches cannot make that physical guarantee and remain an error. + """ + if first.shape == second.shape: + return first, second + + if np.prod(first.shape) >= np.prod(second.shape): + fine_index, fine, coarse = 0, first, second + else: + fine_index, fine, coarse = 1, second, first + factors = tuple( + fine_size // coarse_size + for fine_size, coarse_size in zip(fine.shape, coarse.shape) + ) + if not all( + factor >= 2 and fine_size == factor * coarse_size + for factor, fine_size, coarse_size in zip( + factors, fine.shape, coarse.shape + ) + ) or factors[0] != factors[1]: + raise ValueError( + "shape mismatch cannot be reconciled by equal integer mean " + f"binning: {tuple(first.shape)} != {tuple(second.shape)}. " + "Confirm that both files cover the same field of view." + ) + + factor = factors[0] + rows, columns = coarse.shape + array = np.asarray(fine.array).reshape( + rows, factor, columns, factor + ).mean(axis=(1, 3)) + binned = type(fine).from_array( + np.asarray(array, dtype=np.float32), + name=fine.name, + origin=fine.origin, + sampling=np.asarray(fine.sampling) * factor, + units=fine.units, + signal_units=fine.signal_units, + ) + binned.metadata.update(fine.metadata) + if hasattr(fine, "file_path"): + binned.file_path = fine.file_path + if verbose: + print( + f"mean-binned scan {fine_index} by {factor}x: " + f"{tuple(fine.shape)} -> {tuple(binned.shape)}" + ) + return (binned, second) if fine_index == 0 else (first, binned) + + +def reference_downsample( + reference_scan_shape: tuple[int, int], + target_scan_shape: tuple[int, int], + *, + reference_sampling: NDArray | None = None, + target_sampling: NDArray | None = None, +) -> int: + """Return the averaging needed to put a reference on the target grid. + + Sampling metadata decides when available. Without it, an exact isotropic + integer shape ratio is treated as a coarser acquisition grid, as in a + 2048-pixel HAADF reference paired with a 1024-pixel spectrum image. + """ + source_rows, source_columns = map(int, reference_scan_shape) + target_rows, target_columns = map(int, target_scan_shape) + exact_integer_scale = ( + source_rows % target_rows == 0 + and source_columns % target_columns == 0 + and source_rows // target_rows == source_columns // target_columns + ) + downsample = source_rows // target_rows if exact_integer_scale else 1 + if downsample > 1 and reference_sampling is not None and target_sampling is not None: + reference_step = np.asarray(reference_sampling, dtype=float)[:2] + target_step = np.asarray(target_sampling, dtype=float)[:2] + if np.all(np.isfinite(reference_step)) and np.all(np.isfinite(target_step)): + same_grid = np.allclose( + target_step, + reference_step, + rtol=0.05, + atol=0.0, + ) + coarser_grid = np.allclose( + target_step, + reference_step * downsample, + rtol=0.05, + atol=0.0, + ) + if same_grid: + return 1 + if not coarser_grid: + return 1 + return downsample + + +def match_reference_image( + image: NDArray, + reference_scan_shape: tuple[int, int], + target_scan_shape: tuple[int, int], +) -> NDArray: + """Remove solver padding and center-crop a reference to a target scan.""" + array = np.asarray(image) + source_rows, source_columns = map(int, reference_scan_shape) + target_rows, target_columns = map(int, target_scan_shape) + if array.shape[0] < source_rows or array.shape[1] < source_columns: + raise ValueError( + "corrected reference is smaller than its acquired scan shape: " + f"corrected={array.shape[:2]}, acquired={reference_scan_shape}." + ) + + if target_rows > source_rows or target_columns > source_columns: + raise ValueError( + "target scan is larger than the solved reference grid: " + f"reference={reference_scan_shape}, target={target_scan_shape}." + ) + native_row = (array.shape[0] - source_rows) // 2 + native_column = (array.shape[1] - source_columns) // 2 + native = array[ + native_row : native_row + source_rows, + native_column : native_column + source_columns, + ] + row_start = (native.shape[0] - target_rows) // 2 + column_start = (native.shape[1] - target_columns) // 2 + return np.ascontiguousarray(native[ + row_start : row_start + target_rows, + column_start : column_start + target_columns, + ]) + + +def automatic_alignment_normalization( + images: Sequence[Dataset2d], +) -> tuple[bool, str]: + """Choose min-max conditioning for images with a large DC background.""" + background_to_contrast = [] + for image in images: + array = np.asarray(image.array, dtype=np.float32) + contrast = float(np.std(array)) + background = abs(float(np.mean(array))) + background_to_contrast.append( + background / max(contrast, np.finfo(np.float32).eps) + ) + ratio = max(background_to_contrast) + enabled = ratio >= 5.0 + reason = ( + "large_dc_background" + if enabled + else "same_detector_well_conditioned" + ) + return enabled, reason + + +def scale_coordinate_metadata( + values, + factor: int, + *, + offset: NDArray | None = None, +): + """Scale the first two coordinate entries while preserving length.""" + scaled = np.asarray(values, dtype=float).copy() + if scaled.size < 2: + scaled = np.resize(scaled, 2) + if offset is not None: + scaled[:2] += offset[:2] + else: + scaled[:2] *= factor + return scaled + + def automatic_downsample_factor(shape: tuple[int, int]) -> int: - """Return the largest exact average-pooling factor, capped at eight.""" - rows, columns = shape + """Return the largest safe average-downsampling factor, capped at eight.""" + height, width = shape for factor in (8, 4, 2): - if rows % factor == 0 and columns % factor == 0: + if height % factor == 0 and width % factor == 0: return factor return 1 -def resolve_downsample(value: int | str, shape: tuple[int, int]) -> int: - """Resolve ``"auto"`` or validate one explicit exact divisor.""" +def resolve_downsample( + value: int | str, + shape: tuple[int, int], +) -> int: + """Resolve the automatic affine-search downsampling factor.""" if value == "auto": return automatic_downsample_factor(shape) if isinstance(value, str): @@ -74,43 +380,80 @@ def resolve_downsample(value: int | str, shape: tuple[int, int]) -> int: factor = validate_downsample(value) if any(size % factor for size in shape): raise ValueError( - f"downsample={factor} requires both image dimensions to be " - f"divisible by {factor}, got {shape}." + f"downsample={factor} requires both image dimensions " + f"to be divisible by {factor}, got {shape}. Choose a divisor " + "such as 1, 2, 4, or 8." ) return factor -def _scaled_coordinate_metadata( - values, - factor: int, - *, - offset: NDArray | None = None, -) -> NDArray: - """Scale the first two coordinate entries while preserving their length.""" - scaled = np.asarray(values, dtype=float).copy() - if scaled.size < 2: - scaled = np.resize(scaled, 2) - if offset is None: - scaled[:2] *= factor - else: - scaled[:2] += offset[:2] - return scaled +def minimum_affine_padding_fraction( + image_shape: tuple[int, int], + scan_direction_degrees: Sequence[float], + max_drift_rate: float = 0.25, + translation_margin: float = 0.0, +) -> float: + """Return the smallest canvas fraction covering the affine envelope.""" + rows, cols = (int(value) for value in image_shape) + if rows < 1 or cols < 1: + raise ValueError(f"image_shape must be positive, got {image_shape}.") + rate = abs(float(max_drift_rate)) + translation = float(translation_margin) + if not np.isfinite(rate): + raise ValueError( + f"max_drift_rate must be finite, got {max_drift_rate!r}." + ) + if not np.isfinite(translation) or translation < 0: + raise ValueError( + "translation_margin must be finite and non-negative, " + f"got {translation_margin!r}." + ) + + half_rows = (rows - 1) / 2.0 + half_cols = (cols - 1) / 2.0 + max_row_extent = 0.0 + max_col_extent = 0.0 + for angle_degrees in scan_direction_degrees: + angle = np.deg2rad(float(angle_degrees)) + fast = np.asarray((np.sin(angle), np.cos(angle))) + slow = np.asarray((np.cos(angle), -np.sin(angle))) + extent = np.abs(slow) * half_rows + np.abs(fast) * half_cols + max_row_extent = max(max_row_extent, float(extent[0])) + max_col_extent = max(max_col_extent, float(extent[1])) + + affine_margin = rate * half_rows + + def even_ceiling(value: float) -> int: + integer = int(np.ceil(value)) + return integer if integer % 2 == 0 else integer + 1 + + canvas_rows = even_ceiling( + 2.0 * (max_row_extent + affine_margin + translation) + 1.0 + ) + canvas_cols = even_ceiling( + 2.0 * (max_col_extent + affine_margin + translation) + 1.0 + ) + return max( + 0.0, + canvas_rows / rows - 1.0, + canvas_cols / cols - 1.0, + ) def apply_downsample(correction, factor: int) -> None: - """Average-downsample alignment images and preserve physical coordinates.""" - factor = resolve_downsample(factor, tuple(correction.images[0].shape)) + """Average-downsample images and update their coordinate metadata.""" original_records = [] downsampled_images = [] - for image_index, image in enumerate(correction.images): + for image_index, image in enumerate(correction.imgs): original_array = np.asarray(image.array) original_sampling = np.asarray(image.sampling, dtype=float).copy() original_origin = np.asarray(image.origin, dtype=float).copy() - original_shape = tuple(int(value) for value in original_array.shape) + original_units = list(image.units) + original_shape = tuple(int(value) for value in original_array.shape[:2]) downsampled = average_downsample_2d(original_array, factor) - sampling = _scaled_coordinate_metadata(original_sampling, factor) + sampling = scale_coordinate_metadata(original_sampling, factor) origin_offset = (factor - 1) * original_sampling[:2] / 2.0 - origin = _scaled_coordinate_metadata( + origin = scale_coordinate_metadata( original_origin, factor, offset=origin_offset, @@ -123,117 +466,114 @@ def apply_downsample(correction, factor: int) -> None: "downsample_original_shape": list(original_shape), "downsample_original_sampling": original_sampling.tolist(), "downsample_original_origin": original_origin.tolist(), - "downsampled_shape": list(downsampled.shape), + "downsampled_shape": list(downsampled.shape[:2]), "sampling_is_downsampled": True, } ) - downsampled_image = Dataset2d.from_array( + image_dataset = Dataset2d.from_array( downsampled, name=f"{image.name} ({factor}x downsample)", origin=origin, sampling=sampling, - units=list(image.units), + units=original_units, signal_units=image.signal_units, ) - downsampled_image.metadata.update(metadata) - downsampled_images.append(downsampled_image) + image_dataset.metadata.update(metadata) + downsampled_images.append(image_dataset) original_records.append( { "image_index": image_index, "shape": list(original_shape), "sampling": original_sampling.tolist(), "origin": original_origin.tolist(), - "units": list(image.units), + "units": original_units, } ) - correction.images = downsampled_images + correction.imgs = downsampled_images correction.downsample = factor correction.downsample_method = "average" correction.downsample_metadata = { "factor": factor, "method": "average", "original_images": original_records, - "downsampled_shape": list(correction.images[0].shape), + "downsampled_shape": list(correction.imgs[0].shape[:2]), } def preprocess( self, - pad_fraction: float = 0.25, - pad_value: float | str | list[float] = "median", - kde_sigma: float = 0.5, - number_knots: int = 1, - downsample: int | str = 1, - show_merged: bool = False, - show_images: bool = False, + *, + padding_fraction: float | str = "auto", + padding_value: float | str | list[float] = "median", + smoothing_sigma: float = 0.5, + num_knots: int = 1, + normalize: bool = False, + downsample: int = 1, + verbose: bool = True, + show_combined: bool = False, + show_scans: bool = False, show_knots: bool = True, - **kwargs, + show_knot_plot: bool = False, ): - """Prepare images for drift correction by building the scanline model. - - Computes scan direction vectors, initializes Bezier knots that map - each scanline onto a padded canvas, and generates the initial warped - images. This must be called before any alignment step. + """Build the shared scanline canvas before drift correction. - Without preprocessing, there is no spatial model connecting the raw - images to the shared canvas - alignment methods would have no - coordinates to optimize. + Affine correction calls this automatically. Use it directly only when a + publication requires a fixed padding, normalization, or knot layout. Parameters ---------- - pad_fraction : float - Fraction of the image size to add as padding around the canvas. - Larger values give more room for drift but use more memory. - ``pad_fraction=0.25`` adds 25% on each side. - pad_value : float, str, or list[float] - Fill value for pixels outside the image footprint. Can be - ``'median'``, ``'mean'``, ``'min'``, ``'max'``, a quantile - (e.g. ``0.25``), or a per-image list of floats. - kde_sigma : float - Gaussian smoothing sigma (in pixels) applied after bilinear - scatter. Smooths the warped images to reduce scatter noise. - number_knots : int - Number of Bezier knots per scanline. Use ``1`` (recommended) - for linear drift correction. Higher values allow per-scanline - curvature but are slower and rarely needed. - downsample : int or "auto", default 1 - Exact average-pooling factor for the 2-D alignment images. Sampling - and pixel-center origin are updated so physical coordinates remain - calibrated. This is computational downsampling, not detector binning. - show_merged : bool - Display the merged (averaged) warped images after preprocessing. - show_images : bool - Display each individual warped image after preprocessing. - show_knots : bool - Overlay knot positions on displayed images. - **kwargs - Additional keyword arguments passed to plotting functions. + padding_fraction : float or "auto", default "auto" + Fractional canvas expansion. ``"auto"`` covers the scan rotations and + affine search envelope. + padding_value : float, str, or list of float, default "median" + Intensity outside each measured scan footprint. + smoothing_sigma : float, default 0.5 + Gaussian smoothing in pixels after scanline interpolation. + num_knots : int, default 1 + Knots per scanline for a fixed publication setup. One knot represents + affine drift. For routine non-rigid correction, prefer + ``correct_nonrigid(num_knots=...)`` so preprocessing stays automatic. + normalize : bool, default False + Scale each scan to ``[0, 1]`` for mixed-detector comparisons. + downsample : int, default 1 + Computational downsampling for the 2D alignment images. Returns ------- - Self - For method chaining: ``drift.preprocess().align_affine()``. + object + The same correction object for method chaining. Examples -------- - >>> drift = DriftCorrection.from_data( - ... images=[im0, im1], scan_direction_degrees=[0, 90]) - >>> drift.preprocess(pad_fraction=0.25, kde_sigma=0.5, number_knots=1) + >>> drift.preprocess(padding_fraction=0.25, show_combined=True) """ - factor = resolve_downsample(downsample, tuple(self.images[0].shape)) - if factor > 1: + downsample = validate_downsample(downsample) + if downsample > 1: if getattr(self, "_datasets", None) is not None: - raise ValueError( - "preprocess(downsample>1) is only for 2-D alignment images. " - "Keep multidimensional data native and use the affine search " - "pyramid instead." + raise NotImplementedError( + "preprocess(downsample>1) currently supports 2-D " + "image-collection solves only. Reference, EDS/EELS, and " + "4D-STEM dataset correction need full-resolution fields." ) if hasattr(self, "_initial_knots"): raise RuntimeError( - "downsample changes the alignment grid. Create a new " - "DriftCorrection before choosing a different factor." + "downsample changes the image grid. Create a " + "new DriftCorrection object before changing it." + ) + apply_downsample(self, downsample) + if verbose: + original = self.downsample_metadata["original_images"][0] + units = ", ".join(str(unit) for unit in original["units"][:2]) + print( + "preprocess: downsample=" + f"{downsample} applies computational average downsampling " + f"{tuple(original['shape'])} -> " + f"{tuple(self.imgs[0].shape[:2])}; " + f"sampling {original['sampling'][:2]} -> " + f"{np.asarray(self.imgs[0].sampling, dtype=float)[:2].tolist()} " + f"{units}. This is not acquisition binning; scale metadata " + "was updated for display and scale bars." ) - apply_downsample(self, factor) else: self.downsample = 1 self.downsample_method = "none" @@ -242,143 +582,160 @@ def preprocess( "method": "none", "original_images": [ { - "image_index": index, - "shape": list(image.shape), - "sampling": np.asarray(image.sampling, dtype=float).tolist(), - "origin": np.asarray(image.origin, dtype=float).tolist(), + "image_index": image_index, + "shape": list(image.shape[:2]), + "sampling": np.asarray( + image.sampling, dtype=float + ).tolist(), + "origin": np.asarray( + image.origin, dtype=float + ).tolist(), "units": list(image.units), } - for index, image in enumerate(self.images) + for image_index, image in enumerate(self.imgs) ], - "downsampled_shape": list(self.images[0].shape), + "downsampled_shape": list(self.imgs[0].shape[:2]), } - self.pad_fraction = float(pad_fraction) - self.pad_value = validate_pad_value(pad_value, self.images) - self.kde_sigma = float(kde_sigma) - self.number_knots = int(number_knots) + + self._normalized = bool(normalize) + if normalize: + for image in self.imgs: + array = image.array.astype(np.float32) + low, high = array.min(), array.max() + image.array = (array - low) / (high - low + 1e-8) + self.pad_value = validate_pad_value(padding_value, self.imgs) + self.kde_sigma = float(smoothing_sigma) + number_knots = int(num_knots) + if number_knots < 1: + raise ValueError(f"num_knots must be >= 1 (got {num_knots}).") + self.number_knots = number_knots + + unique_directions = { + round(direction, 6) for direction in self.scan_direction_degrees + } + if len(unique_directions) > 1: + for image_index, image in enumerate(self.imgs): + if image.shape[0] != image.shape[1]: + raise ValueError( + "Multi-direction scan collection require square images, " + f"but image {image_index} is {image.shape}. Either crop " + "to square or use a single scan direction." + ) self.scan_direction = np.deg2rad(self.scan_direction_degrees) self.scan_fast = np.stack( - [np.sin(-self.scan_direction), np.cos(-self.scan_direction)], axis=1) + [np.sin(self.scan_direction), np.cos(self.scan_direction)], axis=1 + ) self.scan_slow = np.stack( - [np.cos(-self.scan_direction), -np.sin(-self.scan_direction)], axis=1) + [np.cos(self.scan_direction), -np.sin(self.scan_direction)], axis=1 + ) + + translation_margin = 0.0 + if padding_fraction == "auto": + translation_margin = ( + min(self.imgs[0].shape[:2]) * 0.125 + if self._built_from_datasets and not self._reference_mode + else 0.0 + ) + self.pad_fraction = minimum_affine_padding_fraction( + tuple(int(value) for value in self.imgs[0].shape[:2]), + self.scan_direction_degrees, + translation_margin=translation_margin, + ) + padding_mode = "auto" + elif isinstance(padding_fraction, str): + raise ValueError( + "padding_fraction must be 'auto' or a non-negative float, " + f"got {padding_fraction!r}." + ) + else: + self.pad_fraction = float(padding_fraction) + if not np.isfinite(self.pad_fraction) or self.pad_fraction < 0: + raise ValueError( + "padding_fraction must be a finite, non-negative value, " + f"got {padding_fraction!r}." + ) + padding_mode = "explicit" self.shape = ( - len(self.images), - int(np.round(self.images[0].shape[0] * (1 + self.pad_fraction) / 2) * 2), - int(np.round(self.images[1].shape[1] * (1 + self.pad_fraction) / 2) * 2), + len(self.imgs), + int( + np.round( + self.imgs[0].shape[0] * (1 + self.pad_fraction) / 2 + ) + * 2 + ), + int( + np.round( + self.imgs[0].shape[1] * (1 + self.pad_fraction) / 2 + ) + * 2 + ), ) - # Initialize knots - each image's scanlines mapped to the padded canvas - self.knots = [] - for img_idx in range(self.shape[0]): - shape = self.images[img_idx].shape - v_slow = np.linspace(-(shape[0] - 1) / 2, (shape[0] - 1) / 2, shape[0]) - u_fast = np.linspace(-(shape[1] - 1) / 2, (shape[1] - 1) / 2, self.number_knots) - row_knots = ((self.shape[1] - 1) / 2 - + u_fast[None, :] * self.scan_fast[img_idx, 0] - + v_slow[:, None] * self.scan_slow[img_idx, 0]) - col_knots = ((self.shape[2] - 1) / 2 - + u_fast[None, :] * self.scan_fast[img_idx, 1] - + v_slow[:, None] * self.scan_slow[img_idx, 1]) - self.knots.append(np.stack([row_knots, col_knots], axis=0)) - self._initial_knots = [knots.copy() for knots in self.knots] - self.interpolator = [ - DriftInterpolator( - input_shape=self.images[i].shape, - output_shape=self.shape[1:], - scan_fast=self.scan_fast[i], - scan_slow=self.scan_slow[i], - pad_value=self.pad_value[i], - kde_sigma=self.kde_sigma, + self.preprocess_info = { + "padding_mode": padding_mode, + "padding_fraction": self.pad_fraction, + "canvas_shape": list(self.shape[1:]), + "normalize": self._normalized, + "smoothing_sigma": self.kde_sigma, + "num_knots": self.number_knots, + "translation_margin": ( + translation_margin if padding_mode == "auto" else None + ), + } + self.knots = [ + torch.tensor( + drift_knots.initialize_scanline_knots( + input_shape=self.imgs[image_index].shape, + output_shape=self.shape[1:], + scan_fast=self.scan_fast[image_index], + scan_slow=self.scan_slow[image_index], + number_knots=self.number_knots, + ), + dtype=self._dtype, + device=self._device, ) - for i in range(self.shape[0]) + for image_index in range(self.shape[0]) + ] + self.u_per_image = [ + np.linspace(0, 1, self.imgs[index].shape[1]) + for index in range(self.shape[0]) ] - # Cache source data on GPU and generate initial warped images device = self._device dtype = self._dtype - self.images_t = [ - torch.tensor(self.images[i].array, dtype=dtype, device=device) - for i in range(self.shape[0]) + self.imgs_t = [ + torch.tensor(self.imgs[index].array, dtype=dtype, device=device) + for index in range(self.shape[0]) ] self.scan_fast_t = [ - torch.tensor(self.scan_fast[i], dtype=dtype, device=device) - for i in range(self.shape[0]) + torch.tensor(self.scan_fast[index], dtype=dtype, device=device) + for index in range(self.shape[0]) ] - self.images_warped = Dataset3d.from_shape(self.shape) - self.weights_warped = Dataset3d.from_shape(self.shape) + self.scan_slow_t = [ + torch.tensor(self.scan_slow[index], dtype=dtype, device=device) + for index in range(self.shape[0]) + ] + self.imgs_warped = Dataset3d.from_shape(self.shape) canvas_shape = (self.shape[1], self.shape[2]) - warped_t = torch.zeros(self.shape[0], *canvas_shape, dtype=dtype, device=device) - for img_idx in range(self.shape[0]): - knots_t = torch.tensor(self.knots[img_idx], dtype=dtype, device=device) - row_t, col_t = transform_coordinates_single_knot( - knots_t, self.scan_fast_t[img_idx], self.images[img_idx].shape) - warped, weights = bilinear_kde_batch( - row_t[None], col_t[None], self.images_t[img_idx], canvas_shape, - self.kde_sigma, self.pad_value[img_idx]) - warped_t[img_idx] = warped[0] - self.images_warped.array[img_idx] = warped[0].cpu().numpy() - self.weights_warped.array[img_idx] = weights[0].cpu().numpy() - diagnostics._record_stage(self, "initial") - self.calculate_error(0, _warped_t=warped_t) - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images(show_knots=show_knots, title="Merged: initial", **kwargs) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: initial" for i in range(self.shape[0])], - **kwargs, - ) - return self - - -def align_translation( - self, - upsample_factor: int = 8, - min_image_shift: float | None = None, - max_image_shift: float = 32, - show_merged: bool = True, - show_images: bool = False, - show_knots: bool = True, - **kwargs, -): - """ - Solve for the translation between all images in DriftCorrection.images_warped - """ - dxy = np.zeros((self.shape[0], 2)) - F_ref = np.fft.fft2(self.images_warped.array[0]) - for ind in range(1, self.shape[0]): - shifts, image_shift = cross_correlation_shift( - F_ref, - np.fft.fft2(self.images_warped.array[ind]), - upsample_factor=upsample_factor, - max_shift=max_image_shift, - fft_input=True, - fft_output=True, - return_shifted_image=True, - ) - dxy[ind, :] = shifts - F_ref = F_ref * ind / (ind + 1) + image_shift / (ind + 1) - dxy -= np.mean(dxy, axis=0) - if min_image_shift is not None: - if np.linalg.norm(dxy[ind]) < min_image_shift: - dxy[ind] = 0.0 - for ind in range(self.shape[0]): - self.knots[ind][0] += dxy[ind, 0] - self.knots[ind][1] += dxy[ind, 1] - for ind in range(self.shape[0]): - self.images_warped.array[ind], self.weights_warped.array[ind] = self.interpolator[ - ind - ].warp_image( - self.images[ind].array, - self.knots[ind], - ) - diagnostics._record_stage(self, "translation") - kwargs.pop("title", None) - if show_merged: - self.plot_merged_images(show_knots=show_knots, title="Merged: translation", **kwargs) - if show_images: - self.plot_transformed_images( - show_knots=show_knots, - title=[f"Image {i}: translation" for i in range(self.shape[0])], - **kwargs, + warped_t = torch.zeros( + self.shape[0], *canvas_shape, dtype=dtype, device=device + ) + for image_index in range(self.shape[0]): + warped, _ = drift_knots.interpolator(self, image_index).warp_to_canvas( + self.imgs_t[image_index], + canvas_shape, + self.kde_sigma, + self.pad_value[image_index], ) + warped_t[image_index] = warped + self.imgs_warped.array[image_index] = warped.cpu().numpy() + self._initial_knots = [knot.clone() for knot in self.knots] + report.record_error(self, 0, warped_t) + drift_plot.show_after_step( + self, + "initial", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() return self diff --git a/src/quantem/imaging/drift/report.py b/src/quantem/imaging/drift/report.py index 812fd544..e54923ef 100644 --- a/src/quantem/imaging/drift/report.py +++ b/src/quantem/imaging/drift/report.py @@ -1,69 +1,170 @@ -"""Structured evidence assembled from pure drift diagnostics. +"""Regional alignment reports for :class:`DriftCorrection`. -These helpers intentionally remain internal during staged integration. The -final API change will decide which table becomes ``DriftCorrection.report``. +This code owns checkpoint selection, fixed-mask NCC calculation, and +DataFrame assembly. The public entry point remains ``DriftCorrection.report``; +keeping the implementation here separates scientific reporting from solver +state and optimization code. """ -from collections.abc import Mapping, Sequence +from collections.abc import Sequence -from . import diagnostics +import numpy as np +import pandas as pd +import torch +import quantem.imaging.drift.core.strip as strip +import quantem.imaging.drift.core.warping as warping +from quantem.imaging.drift.core import knots as drift_knots -def _rename_rows( - rows: Sequence[Mapping[str, float | int | str]], - columns: Mapping[str, str], -) -> list[dict[str, float | int | str]]: - """Return independent records with stable, reader-facing labels.""" - return [{columns.get(name, name): value for name, value in row.items()} for row in rows] +def record_error(correction, mode: int, warped: torch.Tensor | None = None): + """Record per-scan disagreement for a completed correction stage.""" + if warped is not None: + mean = warped.mean(dim=0) + differences = torch.mean( + torch.abs(warped - mean[None]), dim=(1, 2) + ).cpu().numpy() + else: + warping.ensure_warped_images(correction) + mean = np.mean(correction.imgs_warped.array, axis=0) + differences = np.mean( + np.abs(correction.imgs_warped.array - mean[None]), axis=(1, 2) + ) + current = np.hstack((mode, np.mean(differences), differences)) + if not hasattr(correction, "error_track"): + correction.error_track = current[None, :] + else: + correction.error_track = np.vstack((correction.error_track, current)) -def _registration_report( - correction, + +def report( + self, *, stages: Sequence[str] | None = None, -) -> list[dict[str, float | int | str]]: - """Return common-coverage registration records by correction stage.""" - _, _, _, rows = diagnostics._registration_data(correction, stages) - return _rename_rows( - rows, - { - "stage": "Correction stage", - "common_ncc": "Common NCC", - "top_ncc": "Top third NCC", - "middle_ncc": "Middle third NCC", - "bottom_ncc": "Bottom third NCC", - "mean_absolute_difference": "Mean absolute difference (native intensity units)", - "root_mean_square_difference": "Root mean square difference (native intensity units)", - "coverage": "Coverage", - }, +) -> pd.DataFrame: + """Compare alignment quality across completed correction stages. + + Every stage is measured through one final common-coverage mask, so changes + in NCC reflect registration rather than a changing set of padded edge + pixels. The regional scores expose improvements or regressions that a + single whole-image NCC can hide. + + Parameters + ---------- + stages : sequence of str or None, default None + Checkpoints selected from ``"initial"``, ``"affine"``, ``"strip"``, + and ``"current"``. ``None`` includes every completed distinct stage. + + Returns + ------- + pandas.DataFrame + Common, top, middle, bottom, and coverage measurements by stage. + + Examples + -------- + >>> drift.correct_affine(show_combined=False) + >>> drift.report() + """ + if not hasattr(self, "_initial_knots"): + raise RuntimeError( + "report() requires a prepared alignment. Run correct_affine() first." + ) + + snapshots = [("initial", "before")] + if hasattr(self, "_knots_after_affine"): + snapshots.append(("affine", "affine")) + if hasattr(self, "_knots_after_strip"): + snapshots.append(("strip", "strip")) + + last_snapshot = ( + getattr(self, "_knots_after_strip", None) + or getattr(self, "_knots_after_affine", None) + or self._initial_knots + ) + current_is_distinct = any( + not torch.equal(current, previous) + for current, previous in zip(self.knots, last_snapshot, strict=True) ) + if current_is_distinct: + has_nonrigid = hasattr(self, "error_track") and np.any( + np.asarray(self.error_track)[:, 0] == 2 + ) + snapshots.append(("current", "nonrigid" if has_nonrigid else "current")) + if stages is not None: + requested = tuple(str(stage) for stage in stages) + valid = {"initial", "affine", "strip", "current"} + unknown = sorted(set(requested) - valid) + if unknown: + raise ValueError( + f"Unknown report stages {unknown}. Choose from {sorted(valid)}." + ) + labels = dict(snapshots) + missing = [stage for stage in requested if stage not in labels] + if missing: + raise ValueError( + f"Report stages are not available yet: {missing}. " + f"Available stages: {sorted(labels)}." + ) + snapshots = [(stage, labels[stage]) for stage in requested] + fixed_set = frozenset({0}) if self._reference_mode else frozenset() + # One final common mask is held fixed across every row. Score changes then + # measure alignment rather than a changing population of edge pixels. + comparison_mask = np.asarray(self.coverage_mask(), dtype=bool) + rows: list[dict[str, float | str]] = [] + for stage, label in snapshots: + knots = drift_knots.stage_knots( + self, None if stage == "current" else stage + ) + stack = ( + warping.reference_scan_stack(self, knots) + if self._reference_mode + else warping.co_registered_scan_stack( + self, + fixed_set=fixed_set, + solve_translation=False, + knots=knots, + ) + ) + if self._reference_mode: + comparisons = [(stack[0], stack[index]) for index in range(1, len(stack))] + elif len(stack) == 2: + comparisons = [(stack[0], stack[1])] + else: + comparisons = [ + ( + np.mean( + [stack[j] for j in range(len(stack)) if j != index], + axis=0, + ).astype(np.float32), + stack[index], + ) + for index in range(len(stack)) + ] + scores = [ + strip.region_ncc( + reference, + moving, + comparison_mask, + device=self._device, + ) + for reference, moving in comparisons + ] + row: dict[str, float | str] = {"stage": label} + for column in ("common", "top", "middle", "bottom", "mask_frac"): + row[column] = float(np.mean([score[column] for score in scores])) + rows.append(row) -def _displacement_report( - correction, - *, - stages: Sequence[str] | None = None, -) -> list[dict[str, float | int | str]]: - """Return explicit scan-line-origin displacement records by image.""" - rows = diagnostics._displacement_rows(correction, stages) - return _rename_rows( - rows, - { - "stage": "Correction stage", - "image": "Image", - "endpoint_row_displacement_px": "Endpoint row displacement (px)", - "endpoint_column_displacement_px": "Endpoint column displacement (px)", - "endpoint_displacement_px": "Endpoint displacement (px)", - "rms_displacement_px": "RMS displacement (px)", - "max_displacement_px": "Maximum displacement (px)", - "component_rms_adjacent_line_change_px": ( - "Component RMS adjacent-line displacement change (px)" - ), - "component_rms_adjacent_fast_knot_change_px": ( - "Component RMS adjacent-fast-knot displacement change (px; knot-spacing dependent)" - ), - "component_rms_fast_knot_second_difference_px": ( - "Component RMS fast-knot second displacement difference (px; knot-spacing dependent)" - ), - }, + report = pd.DataFrame.from_records(rows).set_index("stage") + report.index.name = ( + "Fixed-reference stage" if self._reference_mode else "Mutual stage" + ) + return report.rename( + columns={ + "common": "Common NCC", + "top": "Top third", + "middle": "Middle third", + "bottom": "Bottom third", + "mask_frac": "Coverage", + } ) diff --git a/tests/imaging/test_drift.py b/tests/imaging/test_drift.py index a9bd5b5c..a1f72acd 100644 --- a/tests/imaging/test_drift.py +++ b/tests/imaging/test_drift.py @@ -1,411 +1,123 @@ -""" -Tests for the DriftCorrection class in quantem.imaging.drift +"""Scientist-facing API contracts for drift correction.""" -Synthetic data: chevron pattern with linear drift + jitter at 0 and 90 deg scan angles. -See PR #133 for images: https://github.com/electronmicroscopy/quantem/pull/133 -""" +import inspect import numpy as np import pytest -from matplotlib import pyplot as plt from scipy.ndimage import gaussian_filter from quantem.core.datastructures.dataset2d import Dataset2d -from quantem.imaging.drift import DriftCorrection -from quantem.imaging.drift import diagnostics as drift_diagnostics -from quantem.imaging.drift import plot as drift_plot -from quantem.imaging.drift import report as drift_report - - -def make_synthetic_drift_data(scale=1, seed=42): - """Generate a chevron base image plus two scan-distorted views. - - Image 0 scans along columns; image 1 scans along rows. Both apply the - same linear row/col drift plus per-scanline jitter so the nonrigid - solver has a non-trivial knot field to recover. - """ - np.random.seed(seed) - shape = (200 * scale, 200 * scale) - row_grid, col_grid = np.meshgrid( - np.arange(-shape[0] / 2, shape[0] / 2), - np.arange(-shape[0] / 2, shape[0] / 2), - indexing="ij", - ) - base_image = (np.mod(np.abs(row_grid) + np.abs(col_grid), 16 * scale) < 8 * scale).astype("float") - base_image[np.logical_and(row_grid > 0, col_grid > 0)] += 0.5 - base_image[np.maximum(np.abs(row_grid), np.abs(col_grid)) < 20 * scale] = 2 - base_image = gaussian_filter(base_image, sigma=0.667 * scale) - - scan_size = 128 * scale - scan_positions = np.arange(scan_size) - row_drift = scan_positions * 0.001 * scale - col_drift = scan_positions * 0.1 * scale - jitter_mag = 0.5 * scale - jitter0 = np.random.randn(2, scan_size) * jitter_mag - jitter1 = np.random.randn(2, scan_size) * jitter_mag - - im0 = np.zeros((scan_size, scan_size)) - for row_idx in range(scan_size): - start_row = 40 * scale + row_idx + row_drift[row_idx] + jitter0[0, row_idx] - start_col = 30 * scale + 0 + col_drift[row_idx] + jitter0[1, row_idx] - row_coords = start_row + scan_positions * 0 - col_coords = start_col + scan_positions * 1 - row_coords = np.clip(row_coords, 0, shape[0] - 2) - col_coords = np.clip(col_coords, 0, shape[1] - 2) - row_floor = np.floor(row_coords).astype("int") - col_floor = np.floor(col_coords).astype("int") - row_frac = row_coords - row_floor - col_frac = col_coords - col_floor - im0[row_idx, :] = ( - base_image[row_floor, col_floor] * (1 - row_frac) * (1 - col_frac) - + base_image[row_floor + 1, col_floor] * row_frac * (1 - col_frac) - + base_image[row_floor, col_floor + 1] * (1 - row_frac) * col_frac - + base_image[row_floor + 1, col_floor + 1] * row_frac * col_frac - ) - - im1 = np.zeros((scan_size, scan_size)) - for row_idx in range(scan_size): - start_row = 170 * scale + 0 + row_drift[row_idx] + jitter1[0, row_idx] - start_col = 30 * scale + row_idx + col_drift[row_idx] + jitter1[1, row_idx] - row_coords = start_row - scan_positions * 1 - col_coords = start_col + scan_positions * 0 - row_coords = np.clip(row_coords, 0, shape[0] - 2) - col_coords = np.clip(col_coords, 0, shape[1] - 2) - row_floor = np.floor(row_coords).astype("int") - col_floor = np.floor(col_coords).astype("int") - row_frac = row_coords - row_floor - col_frac = col_coords - col_floor - im1[row_idx, :] = ( - base_image[row_floor, col_floor] * (1 - row_frac) * (1 - col_frac) - + base_image[row_floor + 1, col_floor] * row_frac * (1 - col_frac) - + base_image[row_floor, col_floor + 1] * (1 - row_frac) * col_frac - + base_image[row_floor + 1, col_floor + 1] * row_frac * col_frac +from quantem.imaging import DriftCorrection + + +def _orthogonal_pair(size: int = 64) -> tuple[Dataset2d, Dataset2d]: + """Return one calibrated synthetic field acquired at 0 and 90 degrees.""" + rng = np.random.default_rng(42) + image = gaussian_filter( + rng.normal(size=(size, size)).astype(np.float32), + sigma=1.2, + ) + rows, columns = np.indices(image.shape, dtype=np.float32) + image += 2.0 * np.exp( + -((rows - 19.0) ** 2 + (columns - 43.0) ** 2) / (2.0 * 4.0**2) + ) + scans = [] + for array, angle in ((image, 0.0), (np.rot90(image, k=-1).copy(), 90.0)): + dataset = Dataset2d.from_array( + array, + origin=(1.0, 2.0), + sampling=(0.08, 0.08), + units=("nm", "nm"), ) - - return im0, im1, base_image - - -def test_full_pipeline_deterministic(): - """Full pipeline produces correct, deterministic, low-error results.""" - im0, im1, _ = make_synthetic_drift_data(scale=1, seed=42) - - drift = DriftCorrection.from_data( - images=[im0, im1], - scan_direction_degrees=[0.0, 90.0], - ).preprocess( - pad_fraction=0.25, - pad_value="median", - kde_sigma=0.5, - number_knots=1, - show_merged=False, - show_images=False, - ) - drift.align_affine(step=0.02, num_tests=5, refine=False) - drift.align_nonrigid( - num_iterations=2, - regularization_sigma_px=0.5, - show_merged=False, - show_images=False, - ) - img_corr = drift.generate_corrected_image(upsample_factor=1, show_image=False) - - assert isinstance(img_corr, Dataset2d) - assert not np.isnan(img_corr.array).any() - assert drift.error_track[-1, 1] < 0.1 - - # Determinism: second run with same seed must match exactly - im0_2, im1_2, _ = make_synthetic_drift_data(scale=1, seed=42) - drift2 = DriftCorrection.from_data( - images=[im0_2, im1_2], - scan_direction_degrees=[0.0, 90.0], - ).preprocess( - pad_fraction=0.25, - pad_value="median", - kde_sigma=0.5, - number_knots=1, - show_merged=False, - show_images=False, - ) - drift2.align_affine(step=0.02, num_tests=5, refine=False) - drift2.align_nonrigid( - num_iterations=2, - regularization_sigma_px=0.5, - show_merged=False, - show_images=False, - ) - img_corr2 = drift2.generate_corrected_image(upsample_factor=1, show_image=False) - - np.testing.assert_array_almost_equal( - img_corr.array, img_corr2.array, decimal=10, - err_msg="Drift correction output is not deterministic!", - ) - - -# Baseline values from float32 torch path, captured once and frozen. -# (scale, error, knots0_sum, knots1_sum) -AFFINE_BASELINES = [ - (1, 0.09237676858901978, 12157.7373046875, 28546.2626953125), - (2, 0.13844291865825653, 49830.908203125, 113497.091796875), - (4, 0.163960263133049, 194685.2421875, 459650.7578125), -] - - -@pytest.mark.parametrize("scale,expected_error,expected_k0,expected_k1", AFFINE_BASELINES) -def test_align_affine_matches_frozen_baseline(scale, expected_error, expected_k0, expected_k1): - """Affine on synthetic data must match frozen float32 baseline.""" - im0, im1, _ = make_synthetic_drift_data(scale=scale, seed=42) - drift = DriftCorrection.from_data( - images=[im0, im1], scan_direction_degrees=[0.0, 90.0], - ).preprocess(show_merged=False, show_images=False) - drift.align_affine( - step=0.02, num_tests=5, refine=True, - show_merged=False, show_images=False, - ) - np.testing.assert_almost_equal( - drift.error_track[-1, 1], expected_error, decimal=8) - np.testing.assert_almost_equal( - drift.knots[0].sum(), expected_k0, decimal=6) - np.testing.assert_almost_equal( - drift.knots[1].sum(), expected_k1, decimal=6) - - -# Frozen baselines for the pytorch backend with optimizer_name="adam". -NONRIGID_ADAM_BASELINES = [ - (1, 0.0562780499458313, 12023.870644569397, 28671.6764421463), - (2, 0.12947668135166168, 49829.36344528198, 113481.72154045105), -] - - -@pytest.mark.parametrize("scale,expected_error,expected_k0,expected_k1", NONRIGID_ADAM_BASELINES) -def test_align_nonrigid_adam_matches_frozen_baseline(scale, expected_error, expected_k0, expected_k1): - """Nonrigid on synthetic data must match frozen baseline. - - Runs preprocess → affine → nonrigid (2 iterations for speed). - If the GPU warp or translation path changes numerical output, - these baselines catch it immediately. - """ - im0, im1, _ = make_synthetic_drift_data(scale=scale, seed=42) - drift = DriftCorrection.from_data( - images=[im0, im1], scan_direction_degrees=[0.0, 90.0], - ).preprocess(show_merged=False, show_images=False) - drift.align_affine( - step=0.02, num_tests=5, refine=True, - show_merged=False, show_images=False, - ) - drift.align_nonrigid( - backend="pytorch", num_iterations=2, adam_steps=50, - regularization_sigma_px=16.0, - # Pin lr to the value the baselines were captured at - the public - # default is now auto-derived from max_image_shift, but the frozen - # baselines must stay numerically stable across that change. - lr=0.02, - show_merged=False, show_images=False, - ) - np.testing.assert_almost_equal( - drift.error_track[-1, 1], expected_error, decimal=8) - np.testing.assert_almost_equal( - drift.knots[0].sum(), expected_k0, decimal=6) - np.testing.assert_almost_equal( - drift.knots[1].sum(), expected_k1, decimal=6) - - -# Frozen baselines for the pytorch backend with optimizer_name="lbfgs". -# Shares _compiled_loss_fn with the Adam path, so this catches regressions -# in either the optimizer dispatch or the shared loss. -NONRIGID_LBFGS_BASELINES = [ - (1, 0.07269975543022156, 12152.98459815979, 28536.15177345276), - (2, 0.11153321713209152, 50293.12340545654, 113601.38675689697), -] - - -@pytest.mark.parametrize("scale,expected_error,expected_k0,expected_k1", NONRIGID_LBFGS_BASELINES) -def test_align_nonrigid_lbfgs_matches_frozen_baseline(scale, expected_error, expected_k0, expected_k1): - """Nonrigid LBFGS path on synthetic data must match frozen baseline. - - The LBFGS optimizer uses a closure-based forward+backward instead of - Adam's compiled inner loop. This test ensures both paths stay - numerically deterministic and that LBFGS doesn't silently regress. - """ - im0, im1, _ = make_synthetic_drift_data(scale=scale, seed=42) - drift = DriftCorrection.from_data( - images=[im0, im1], scan_direction_degrees=[0.0, 90.0], - ).preprocess(show_merged=False, show_images=False) - drift.align_affine( - step=0.02, num_tests=5, refine=True, - show_merged=False, show_images=False, - ) - drift.align_nonrigid( - backend="pytorch", optimizer_name="lbfgs", - num_iterations=2, lbfgs_max_iter=20, - regularization_sigma_px=16.0, - show_merged=False, show_images=False, - ) - np.testing.assert_almost_equal( - drift.error_track[-1, 1], expected_error, decimal=8) - np.testing.assert_almost_equal( - drift.knots[0].sum(), expected_k0, decimal=6) - np.testing.assert_almost_equal( - drift.knots[1].sum(), expected_k1, decimal=6) - - -def _make_affine_diagnostic_correction(): - """Build the standard synthetic affine result used by diagnostics tests.""" - image_0, image_1, _ = make_synthetic_drift_data(scale=1, seed=42) - correction = DriftCorrection.from_data( - images=[image_0, image_1], - scan_direction_degrees=[0.0, 90.0], - ).preprocess(show_merged=False, show_images=False) - correction.align_affine( - step=0.02, - num_tests=5, - refine=False, - show_merged=False, - show_images=False, - ) - return correction - - -def test_registration_report_matches_frozen_common_coverage_metrics(): - """Registration evidence uses one measured footprint across stages.""" - correction = _make_affine_diagnostic_correction() - registration = drift_report._registration_report(correction) - - assert [row["Correction stage"] for row in registration] == ["initial", "affine"] - np.testing.assert_allclose( - [row["Coverage"] for row in registration], - [0.6423046875, 0.6423046875], - rtol=0, - atol=1e-12, - ) - np.testing.assert_allclose( - [ - [row["Common NCC"], row["Mean absolute difference (native intensity units)"]] - for row in registration - ], - [ - [0.48240861115393513, 0.456299068925209], - [0.7480021414803405, 0.2941893335769055], - ], - rtol=0, - # The affine solver's sub-pixel translation varies slightly across - # CPU thread counts; the report itself is deterministic for fixed knots. - atol=1e-4, - ) - - -def test_displacement_report_names_endpoint_and_adjacent_change_metrics(): - """Displacement evidence distinguishes endpoint motion from adjacent changes.""" - correction = _make_affine_diagnostic_correction() - displacement = drift_report._displacement_report(correction) - - initial = [row for row in displacement if row["Correction stage"] == "initial"] - affine = [row for row in displacement if row["Correction stage"] == "affine"] - assert all( - value == 0 - for row in initial - for name, value in row.items() - if name not in {"Correction stage", "Image"} - ) - np.testing.assert_allclose( - [row["Endpoint displacement (px)"] for row in affine], - [5.679612662972729, 5.679612662972729], - rtol=0, - atol=1e-7, - ) - np.testing.assert_allclose( - [row["Component RMS adjacent-line displacement change (px)"] for row in affine], - [0.0316227766016838, 0.0316227766016838], - rtol=0, - atol=1e-8, - ) - np.testing.assert_array_equal( - [ - row[ - "Component RMS adjacent-fast-knot displacement change (px; knot-spacing dependent)" - ] - for row in affine - ], - 0.0, - ) - - -def test_diagnostic_reports_and_plots_do_not_mutate_correction(): - """Reports and static figures leave fitted fields and caches untouched.""" - correction = _make_affine_diagnostic_correction() - knots_before = [np.asarray(knot).copy() for knot in correction.knots] - images_before = correction.images_warped.array.copy() - weights_before = correction.weights_warped.array.copy() - errors_before = correction.error_track.copy() - - drift_report._registration_report(correction) - drift_report._displacement_report(correction) - registration_figure, registration_axes = drift_plot._plot_registration_diagnostics(correction) - displacement_figure, displacement_axes = drift_plot._plot_displacement_diagnostics(correction) - - assert registration_axes.shape == (2, 3) - assert displacement_axes.shape == (2, 2) - assert registration_figure.number in plt.get_fignums() - assert displacement_figure.number in plt.get_fignums() - for before, after in zip(knots_before, correction.knots, strict=True): - np.testing.assert_array_equal(after, before) - np.testing.assert_array_equal(correction.images_warped.array, images_before) - np.testing.assert_array_equal(correction.weights_warped.array, weights_before) - np.testing.assert_array_equal(correction.error_track, errors_before) - assert drift_diagnostics._available_stages(correction) == ("initial", "affine") - plt.close(registration_figure) - plt.close(displacement_figure) - - -def test_translation_records_current_knots_and_affine_rerun_discards_it(): - """Translation evidence follows the fitted state through an affine rerun.""" - correction = _make_affine_diagnostic_correction() - correction.align_translation(show_merged=False, show_images=False) - - assert drift_diagnostics._available_stages(correction) == ( - "initial", - "affine", - "translation", - ) - for recorded, current in zip( - drift_diagnostics._stage_knots(correction, "translation"), - correction.knots, - strict=True, + dataset.metadata["scan_rotation_deg"] = angle + scans.append(dataset) + return scans[0], scans[1] + + +def test_primary_api_uses_scientist_facing_names(): + """The final PR exposes one concise chain without legacy solve names.""" + preprocess = inspect.signature(DriftCorrection.preprocess).parameters + affine = inspect.signature(DriftCorrection.correct_affine).parameters + nonrigid = inspect.signature(DriftCorrection.correct_nonrigid).parameters + + assert {"padding_fraction", "smoothing_sigma", "num_knots"} <= set(preprocess) + assert {"max_drift_rate", "num_rates", "region"} <= set(affine) + assert {"num_knots", "num_refine_cycles", "knot_smoothing_sigma"} <= set( + nonrigid + ) + for old_name in ( + "pad_fraction", + "kde_sigma", + "number_knots", + "step", + "num_tests", + "num_iterations", ): - np.testing.assert_array_equal(recorded, current) - - correction.align_affine( - step=0.02, - num_tests=5, + assert old_name not in preprocess + assert old_name not in affine + assert old_name not in nonrigid + for old_method in ( + "from_data", + "align_affine", + "align_nonrigid", + "generate_corrected", + ): + assert not hasattr(DriftCorrection, old_method) + + +def test_from_images_requires_angles_for_bare_arrays(): + """Bare arrays never receive a silent scan-angle assumption.""" + image = np.zeros((16, 16), dtype=np.float32) + with pytest.raises(TypeError, match="scan_direction_degrees is required"): + DriftCorrection.from_images(image, image.copy(), device="cpu") + + +def test_metadata_driven_affine_workflow_returns_calibrated_dataset(): + """The normal image workflow is short, finite, and calibration preserving.""" + scan_0, scan_90 = _orthogonal_pair() + drift = DriftCorrection.from_images(scan_0, scan_90, device="cpu") + drift.preprocess( + padding_fraction=0.25, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + drift.correct_affine( + max_drift_rate=0.02, + num_rates=5, refine=False, - show_merged=False, - show_images=False, - ) - - assert drift_diagnostics._available_stages(correction) == ("initial", "affine") - - -def test_multiscan_plot_metric_describes_the_rendered_comparison(): - """A multi-scan plot labels the exact scan-0-versus-rest image it shows.""" - image_0, image_1, _ = make_synthetic_drift_data(scale=1, seed=42) - image_2 = np.roll(image_0, shift=7, axis=0) - correction = DriftCorrection.from_data( - images=[image_0, image_1, image_2], - scan_direction_degrees=[0.0, 90.0, 45.0], - ).preprocess(show_merged=False, show_images=False) - - figure, axes = drift_plot._plot_registration_diagnostics( - correction, - stages=("initial",), - ) - _, stacks, common_mask, aggregate_rows = drift_diagnostics._registration_data( - correction, - ("initial",), - ) - reference, moving = drift_plot._comparison_pair(stacks["initial"]) - expected_overlay, _ = drift_plot._registration_overlay(reference, moving, common_mask) - displayed_metrics = drift_diagnostics._pair_metrics(reference, moving, common_mask) - - np.testing.assert_allclose(axes[0, 0].images[0].get_array(), expected_overlay) - assert f"{displayed_metrics['common_ncc']:.4f}" in axes[0, 0].get_title() - assert not np.isclose(displayed_metrics["common_ncc"], aggregate_rows[0]["common_ncc"]) - plt.close(figure) + max_image_shift=8, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + corrected = drift.corrected( + upsample_factor=1, + output_original_shape=True, + verbose=False, + ) + + assert isinstance(corrected, Dataset2d) + assert corrected.shape == scan_0.shape + assert np.isfinite(corrected.array).all() + np.testing.assert_allclose(corrected.origin, scan_0.origin) + np.testing.assert_allclose(corrected.sampling, scan_0.sampling) + assert corrected.units == scan_0.units + assert len(drift.drift_rate) == 2 + + automatic = drift.corrected(verbose=False) + canvas = drift.corrected( + output_original_shape=False, + verbose=False, + ) + assert automatic.shape == tuple(drift.shape[-2:]) + assert canvas.shape == tuple(drift.shape[-2:]) + + +def test_nonrigid_diagnostic_defines_fast_roughness(): + """The difficult multi-knot diagnostic states exactly what roughness means.""" + doc = DriftCorrection.diagnose_nonrigid.__doc__ or "" + assert "root-mean-square difference between neighboring knot" in doc + assert "does not measure image noise" in doc diff --git a/tests/imaging/test_drift3d.py b/tests/imaging/test_drift3d.py index 3035492c..bdced6e9 100644 --- a/tests/imaging/test_drift3d.py +++ b/tests/imaging/test_drift3d.py @@ -75,19 +75,28 @@ def test_reference_correction_preserves_spectra_and_calibration(): drift = DriftCorrection.from_reference( reference, - alignment_image, + spectrum_image, + alignment_image=alignment_image, scan_direction_degrees=0.0, - ).preprocess(show_merged=False, show_images=False) - initial_reference_knots = drift.knots[0].copy() - drift.align_affine( - step=0.04, - num_tests=11, + device="cpu", + ).preprocess( + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + initial_reference_knots = drift.knots[0].clone() + drift.correct_affine( + max_drift_rate=0.2, + num_rates=11, refine=True, max_image_shift=16, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, ) - corrected = drift.apply_correction(spectrum_image) + corrected = drift.corrected(verbose=False) interior = np.s_[12:-12, 12:-12] raw_ncc = np.corrcoef( @@ -100,7 +109,11 @@ def test_reference_correction_preserves_spectra_and_calibration(): )[0, 1] assert corrected_ncc > raw_ncc + 0.02 assert corrected_ncc > 0.99 - np.testing.assert_allclose(drift.knots[0], initial_reference_knots, atol=0.0) + np.testing.assert_allclose( + drift.knots[0].cpu(), + initial_reference_knots.cpu(), + atol=0.0, + ) np.testing.assert_allclose( corrected.array[..., 1], 2.0 * corrected.array[..., 0] + 3.0, diff --git a/tests/imaging/test_drift4d.py b/tests/imaging/test_drift4d.py index a4de54bc..f927e74d 100644 --- a/tests/imaging/test_drift4d.py +++ b/tests/imaging/test_drift4d.py @@ -39,33 +39,46 @@ def _fit_small_pair(cube_0, cube_1): scan_units="nm", device="cpu", ).preprocess( - pad_fraction=0.25, - number_knots=1, - show_merged=False, - show_images=False, + padding_fraction=0.25, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) - drift.align_affine( - step=0.01, - num_tests=3, + drift.correct_affine( + max_drift_rate=0.01, + num_rates=3, refine=False, max_image_shift=8, chunk_size=1, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) return drift def _add_known_raw_drift(drift, image_index, row, column): """Move canvas knots by one requested raw-frame trajectory.""" - scan_rows, scan_cols = drift.images[image_index].shape + scan_rows, scan_cols = drift.imgs[image_index].shape aspect = (scan_rows - 1) / (scan_cols - 1) slow = drift.scan_slow[image_index] fast = drift.scan_fast[image_index] + row = torch.as_tensor( + row, + dtype=drift.knots[image_index].dtype, + device=drift.knots[image_index].device, + ) + column = torch.as_tensor( + column, + dtype=drift.knots[image_index].dtype, + device=drift.knots[image_index].device, + ) drift.knots[image_index][0, :, 0] += slow[0] * row + fast[0] * aspect * column drift.knots[image_index][1, :, 0] += slow[1] * row + fast[1] * column + drift._images_warped_stale = True def test_virtual_detector_matches_numpy_and_torch_integer_inputs(): @@ -123,10 +136,11 @@ def test_regional_patterns_average_native_detector_samples(): scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( - number_knots=1, - show_merged=False, - show_images=False, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) _add_known_raw_drift( drift, @@ -198,10 +212,11 @@ def test_saved_correction_accepts_explicit_4dstem_datasets(): scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( - number_knots=1, - show_merged=False, - show_images=False, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) drift._datasets = None @@ -249,8 +264,8 @@ def test_dataset4dstem_metadata_supplies_rotation_and_scan_calibration(): drift = DriftCorrection.from_4dstem(*datasets, device="cpu") np.testing.assert_allclose(drift.scan_direction_degrees, (0.0, 90.0)) - np.testing.assert_allclose(drift.images[0].sampling, (0.2, 0.3)) - assert drift.images[0].units == ["nm", "nm"] + np.testing.assert_allclose(drift.imgs[0].sampling, (0.2, 0.3)) + assert drift.imgs[0].units == ["nm", "nm"] assert drift._datasets[0] is datasets[0].array @@ -263,10 +278,11 @@ def test_drift_field_reports_raw_components_for_rotated_scan(): scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( - number_knots=1, - show_merged=False, - show_images=False, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) expected_row = np.linspace(-2.0, 2.0, 16) expected_column = np.linspace(1.0, -1.0, 16) @@ -331,9 +347,9 @@ def reject_full_allocation(*shape, **kwargs): ) assert corrected.corrected_4dstem_0 is output_0 - assert np.shares_memory(corrected.corrected_4dstem_1, output_1) - assert np.isfinite(output_0).all() assert np.isfinite(output_1).all() + assert corrected.corrected_4dstem_1.shape == output_1.shape + assert np.isfinite(output_0).all() def test_integer_merge_is_float32_and_backend_consistent(): @@ -383,10 +399,11 @@ def test_cuda_matches_cpu_for_uint32_native_detector_frames(): scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( - number_knots=1, - show_merged=False, - show_images=False, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) cuda = DriftCorrection.from_4dstem( torch.from_numpy(cube_0).cuda(), @@ -394,10 +411,11 @@ def test_cuda_matches_cpu_for_uint32_native_detector_frames(): scan_direction_degrees=(0.0, 90.0), device="cuda", ).preprocess( - number_knots=1, - show_merged=False, - show_images=False, + num_knots=1, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) drift_row = np.linspace(-0.75, 0.75, scan_size) drift_column = np.linspace(0.5, -0.5, scan_size) diff --git a/tests/imaging/test_drift_downsampling.py b/tests/imaging/test_drift_downsampling.py index e343b37f..fdff9439 100644 --- a/tests/imaging/test_drift_downsampling.py +++ b/tests/imaging/test_drift_downsampling.py @@ -6,7 +6,6 @@ from quantem.core.datastructures.dataset2d import Dataset2d from quantem.imaging.drift import DriftCorrection -from quantem.imaging.drift.core.affine import _affine_rate_from_knots from quantem.imaging.drift.preparation import ( average_downsample_2d, resolve_downsample, @@ -44,41 +43,43 @@ def test_average_downsample_is_exact_block_mean(): def test_preprocess_downsample_preserves_pixel_center_calibration(): """Sampling grows and origin moves to the center of each averaged block.""" datasets = _calibrated_pair() - drift = DriftCorrection.from_data( - datasets, + drift = DriftCorrection.from_images( + *datasets, scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( downsample=2, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) - assert drift.images[0].shape == (8, 8) - np.testing.assert_allclose(drift.images[0].sampling, (0.4, 0.6)) - np.testing.assert_allclose(drift.images[0].origin, (1.1, 2.15)) - assert drift.images[0].units == ["nm", "nm"] - assert drift.images[0].metadata["source"] == "synthetic calibrated scan" + assert drift.imgs[0].shape == (8, 8) + np.testing.assert_allclose(drift.imgs[0].sampling, (0.4, 0.6)) + np.testing.assert_allclose(drift.imgs[0].origin, (1.1, 2.15)) + assert drift.imgs[0].units == ["nm", "nm"] + assert drift.imgs[0].metadata["source"] == "synthetic calibrated scan" assert drift.downsample_metadata["original_images"][0]["shape"] == [16, 16] def test_corrected_output_records_downsampling_provenance(): """A corrected image tells readers which computational grid was fitted.""" - drift = DriftCorrection.from_data( - _calibrated_pair(), + drift = DriftCorrection.from_images( + *_calibrated_pair(), scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( downsample=2, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) - corrected = drift.generate_corrected( + corrected = drift.corrected( upsample_factor=1, - show_merged=False, + verbose=False, ) assert corrected.metadata["downsample"] == 2 @@ -92,21 +93,23 @@ def test_downsample_requires_exact_divisor_and_new_correction(): with pytest.raises(ValueError, match="divisible"): resolve_downsample(4, (18, 18)) - drift = DriftCorrection.from_data( - _calibrated_pair(), + drift = DriftCorrection.from_images( + *_calibrated_pair(), scan_direction_degrees=(0.0, 90.0), device="cpu", ).preprocess( - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) with pytest.raises(RuntimeError, match="Create a new DriftCorrection"): drift.preprocess( downsample=2, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) @@ -117,8 +120,8 @@ def test_automatic_factor_is_largest_exact_divisor_up_to_eight(): assert resolve_downsample("auto", (1025, 1025)) == 1 -def test_affine_pyramid_refines_rate_on_native_grid(): - """A pooled broad search retains the native image and affine rate.""" +def test_affine_pyramid_search_retains_native_grid(): + """A pooled broad search never downsamples the fitted or output grid.""" size = 64 rng = np.random.default_rng(23) reference = gaussian_filter(rng.normal(size=(size, size)), 1.5).astype(np.float32) @@ -134,50 +137,21 @@ def test_affine_pyramid_refines_rate_on_native_grid(): right=float(np.median(reference[row])), ) - native = DriftCorrection.from_reference( + drift = DriftCorrection.from_reference( reference, target, scan_direction_degrees=0.0, device="cpu", - ).preprocess( - show_merged=False, - show_images=False, - show_knots=False, ) - pyramid = DriftCorrection.from_reference( - reference, - target, - scan_direction_degrees=0.0, - device="cpu", - ).preprocess( - show_merged=False, - show_images=False, - show_knots=False, - ) - - native.align_affine( - step=0.02, - num_tests=7, - max_image_shift=8, - downsample=1, - show_merged=False, - show_images=False, - show_knots=False, - ) - pyramid.align_affine( - step=0.02, - num_tests=7, - max_image_shift=8, + drift.correct_affine( downsample=2, - show_merged=False, - show_images=False, + show_combined=False, + show_scans=False, show_knots=False, + verbose=False, ) - assert pyramid.images[0].shape == (size, size) - assert pyramid.affine_search_info["downsample_factor"] == 2 - np.testing.assert_allclose( - _affine_rate_from_knots(pyramid, 1), - _affine_rate_from_knots(native, 1), - atol=0.01, - ) + assert drift.imgs[0].shape == (size, size) + assert drift.affine_search_info["downsample_factor"] == 2 + assert np.isfinite(drift.drift_rate).all() + assert drift.corrected(verbose=False).shape == (size, size) diff --git a/tests/imaging/test_drift_residual.py b/tests/imaging/test_drift_residual.py new file mode 100644 index 00000000..e5867dfb --- /dev/null +++ b/tests/imaging/test_drift_residual.py @@ -0,0 +1,174 @@ +"""Residual-correction contracts needed by the publication workflows.""" + +import numpy as np +import pytest +import torch + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.imaging.drift import DriftCorrection, StripPass +from quantem.imaging.drift.core.nonrigid import _regularize_knots +from quantem.imaging.drift.core.strip import ( + free_weight, + measure_strip_residual_torch, +) + + +def _accelerator_device() -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + pytest.skip("non-rigid backend parity requires CUDA or MPS") + + +@pytest.mark.parametrize("num_knots", (1, 2, 3)) +def test_multiknot_diagnostics_have_explicit_fast_direction_meaning(num_knots): + """Fast roughness is adjacent-knot displacement, not image roughness.""" + values = np.linspace(0.0, 1.0, num_knots, dtype=np.float32) + measured = ( + 0.0 + if num_knots == 1 + else float(np.sqrt(np.mean(np.diff(values) ** 2))) + ) + assert measured == 0.0 if num_knots == 1 else measured > 0.0 + doc = DriftCorrection.diagnose_nonrigid.__doc__ or "" + assert "neighboring knot" in doc + assert "does not measure image noise" in doc + + +def test_publication_strip_recipe_is_explicit_and_ordered(): + """The frozen XEDS workflow has three visible coarse-to-fine passes.""" + recipe = ( + StripPass( + num_strips=24, + smoothing_sigma=12, + max_column_shift=80, + max_row_shift=8, + ), + StripPass( + num_strips=24, + smoothing_sigma=12, + max_column_shift=12, + max_row_shift=3, + ), + StripPass( + num_strips=64, + smoothing_sigma=6, + max_column_shift=6, + max_row_shift=2, + update_fraction=0.8, + ), + ) + + assert [item.num_strips for item in recipe] == [24, 24, 64] + assert [item.max_column_shift for item in recipe] == [80, 12, 6] + assert recipe[-1].update_fraction == 0.8 + + +def test_strip_free_weight_can_freeze_then_smoothly_release_scanlines(): + """A partial residual update has a stable zero-to-one transition.""" + weights = free_weight(100, free_from_frac=0.55, ramp_frac=0.08) + + assert weights.shape == (100,) + assert np.all(weights[:47] == 0.0) + assert np.all((weights >= 0.0) & (weights <= 1.0)) + assert np.all(weights[55:] == 1.0) + + +def test_strip_measurement_recovers_local_integer_residual(): + """Each slow-scan strip recovers the same known residual displacement.""" + rng = np.random.default_rng(5) + reference = rng.normal(size=(48, 48)).astype(np.float32) + moving = np.roll(np.roll(reference, 2, axis=0), -3, axis=1) + + result = measure_strip_residual_torch( + reference, + moving, + np.ones_like(reference, dtype=bool), + n_strips=4, + max_shift_row=3, + max_shift_col=4, + device="cpu", + method="brute", + ) + + np.testing.assert_array_equal(result["drow"], np.full(4, -2.0)) + np.testing.assert_array_equal(result["dcol"], np.full(4, 3.0)) + assert np.all(result["valid"]) + + +@pytest.mark.parametrize("trend_order", (0, 1, 2, 3)) +def test_nonrigid_regularization_matches_cpu(trend_order): + """The MPS/CUDA fallback retains float32 CPU knot precision.""" + device = _accelerator_device() + row_count = 65 + generator = torch.Generator().manual_seed(42) + coordinates = torch.arange(row_count, dtype=torch.float32) + coordinates = (coordinates - coordinates.mean()) / coordinates.std() + vander = torch.stack( + [coordinates**power for power in range(trend_order + 1)], + dim=1, + ) + knots = torch.randn(2, 2, row_count, 3, generator=generator) + previous = torch.randn(2, 2, row_count, 3, generator=generator) + + expected = knots.clone() + _regularize_knots(expected, previous, vander, 2, 4, 0.8) + actual = knots.to(device) + _regularize_knots( + actual, + previous.to(device), + vander.to(device), + 2, + 4, + 0.8, + ) + + torch.testing.assert_close(actual.cpu(), expected, rtol=3e-5, atol=3e-5) + + +def test_two_dimensional_map_uses_same_field_and_preserves_metadata(): + """Element maps and equivalent cube channels share one spatial warp.""" + rows, columns = np.indices((24, 24), dtype=np.float32) + image = rows + 2 * columns + drift = DriftCorrection.from_reference( + image, + image.copy(), + scan_direction_degrees=0.0, + device="cpu", + ).preprocess( + num_knots=1, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + drift.correct_affine( + max_drift_rate=0.01, + num_rates=3, + refine=False, + max_image_shift=4, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + element_map = Dataset2d.from_array( + image, + name="Ti K", + origin=(1.0, 2.0), + sampling=(0.2, 0.3), + units=("nm", "nm"), + ) + element_map.metadata["line"] = "Ti_K_wide" + + corrected_map = drift.apply_correction(element_map, image_index=1) + corrected_cube = drift.apply_correction( + np.stack((image, 3 * image), axis=-1), + image_index=1, + ) + + np.testing.assert_allclose(corrected_map.array, corrected_cube[..., 0]) + np.testing.assert_allclose(corrected_map.origin, element_map.origin) + np.testing.assert_allclose(corrected_map.sampling, element_map.sampling) + assert corrected_map.metadata["line"] == "Ti_K_wide" diff --git a/tests/imaging/test_drift_utils.py b/tests/imaging/test_drift_utils.py index bf540b4e..4f0eceae 100644 --- a/tests/imaging/test_drift_utils.py +++ b/tests/imaging/test_drift_utils.py @@ -10,11 +10,7 @@ from scipy.ndimage import gaussian_filter from quantem.core.utils.imaging_utils import bilinear_kde, fourier_cropping -from quantem.imaging.drift import ( - _bounded_sine_sigmoid_torch, - _fourier_crop_torch, - bounded_sine_sigmoid, -) +from quantem.imaging.drift.apply import fourier_crop_torch from quantem.imaging.drift.core.knots import ( _symmetric_pad, bilinear_kde_batch, @@ -195,28 +191,9 @@ def test_parabolic_sub_pixel_exact(): # --------------------------------------------------------------------------- -# generate_corrected helpers: torch parity against numpy originals +# Corrected-output Fourier helper parity against NumPy # --------------------------------------------------------------------------- - -@pytest.mark.parametrize("midpoint,width", [(0.5, 1.0), (0.3, 0.4), (0.7, 0.5)]) -def test_bounded_sine_sigmoid_torch_matches_numpy(midpoint, width): - """Torch sigmoid helper must match the numpy original point-for-point. - - bounded_sine_sigmoid is the Fourier low-pass weight ramp used in - generate_corrected(). A mismatch would apply different weighting - to the merged output than the numpy path. - """ - rng = np.random.default_rng(0) - x = rng.random(256).astype(np.float32) - expected = bounded_sine_sigmoid(x, midpoint=midpoint, width=width).astype(np.float32) - result = _bounded_sine_sigmoid_torch( - torch.tensor(x), midpoint=midpoint, width=width - ).numpy() - np.testing.assert_allclose(result, expected, atol=1e-6, - err_msg=f"Sigmoid mismatch at midpoint={midpoint}, width={width}") - - @pytest.mark.parametrize("input_shape,crop_shape", [ ((64, 64), (32, 32)), ((48, 80), (24, 40)), @@ -232,7 +209,7 @@ def test_fourier_crop_torch_matches_numpy(input_shape, crop_shape): rng = np.random.default_rng(1) arr = (rng.random(input_shape) + 1j * rng.random(input_shape)).astype(np.complex64) expected = fourier_cropping(arr, crop_shape) - result = _fourier_crop_torch( + result = fourier_crop_torch( torch.tensor(arr), crop_shape ).numpy() np.testing.assert_allclose(result, expected, atol=1e-6, diff --git a/uv.lock b/uv.lock index 36a6293e..33d75d20 100644 --- a/uv.lock +++ b/uv.lock @@ -2,9 +2,15 @@ version = 1 revision = 3 requires-python = ">=3.11" resolution-markers = [ - "python_full_version >= '3.14'", - "python_full_version >= '3.12' and python_full_version < '3.14'", - "python_full_version < '3.12'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] [manifest] @@ -2058,7 +2064,9 @@ name = "numpy" version = "2.4.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.12'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } wheels = [ @@ -2140,8 +2148,12 @@ name = "numpy" version = "2.5.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14'", - "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/22/fd/89965aa4ac08c74998539fcbf24fa3540f3e15237fbeb6bcf9c908f4aade/numpy-2.5.1.tar.gz", hash = "sha256:a48a113e6afea91f5608793bafa7ef2ad481fefbda87ec5069f483de61cb9fa3", size = 20755553, upload-time = "2026-07-04T17:08:00.933Z" } wheels = [ @@ -2379,6 +2391,61 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, ] +[[package]] +name = "pandas" +version = "3.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "python-dateutil" }, + { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/be/4f/5f3422a2afec5ffc46308b79e53291365a93748b498ac2e58bead0197916/pandas-3.0.5.tar.gz", hash = "sha256:dca3734d6ab7c906e6730f0788b0a1dbb9f2467731f9711f77995c8e9d62d712", size = 4658219, upload-time = "2026-07-22T22:19:28.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/48/ef/f1fd7431d635bf20015489bf0bd69c17fff1018de773540f651455a3916b/pandas-3.0.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2946e77e4a53cd248cbde631a12f0e51c8324ce354c3eba4d20147c1ad6f4282", size = 10397178, upload-time = "2026-07-22T22:17:48.274Z" }, + { url = "https://files.pythonhosted.org/packages/31/b4/0eafac990a431561187694126de01f9b12559549b4d86360c0c4bd870fde/pandas-3.0.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:71ecc8fb7ed1a7aa4392316b5309a6347e8e7f832f38fd897846b3a1457a9298", size = 9990736, upload-time = "2026-07-22T22:17:52.388Z" }, + { url = "https://files.pythonhosted.org/packages/de/21/359880af3ea9b7cb23bea5b51e8e70ef3866c03be09da9a2787e18e330a8/pandas-3.0.5-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b173f5951ff6b8b0ec7675e20dff3c97b7e7a57dfcce387c2d7c5afe87cb7899", size = 10814438, upload-time = "2026-07-22T22:17:54.708Z" }, + { url = "https://files.pythonhosted.org/packages/d1/50/d6cc4d7e508bbccf5d6027314a8312bc7ac73d0ec7f195f53838daafab40/pandas-3.0.5-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2c0cf1dd9b55a22d105fc46c1b489af3bd42264fcba7c66297bf47a9a1d9c78a", size = 11323634, upload-time = "2026-07-22T22:17:56.858Z" }, + { url = "https://files.pythonhosted.org/packages/70/2b/d5f0a8c90dd0ae04e64ba53b871afb796ec026b615086d382ddc2ade729b/pandas-3.0.5-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:0fac0010c75e4efb6b99e249c183a8993ce0dc95c240f9b120a5e67c727b7928", size = 11850860, upload-time = "2026-07-22T22:17:59.1Z" }, + { url = "https://files.pythonhosted.org/packages/5c/30/183aec2e19adf778a98d29b5729a0a68f4cc4ebf9b9c3b70d0297355bcb1/pandas-3.0.5-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:08d24fe11a17dc33bd6e937dc9c665f9cba08fbdc9f657f405713515febe300d", size = 12411100, upload-time = "2026-07-22T22:18:01.485Z" }, + { url = "https://files.pythonhosted.org/packages/fa/9a/31f4983f191af51ab2a8f2d0c7b33dff3a84da26533f982fff02c2f9e28b/pandas-3.0.5-cp311-cp311-win_amd64.whl", hash = "sha256:b1261758dfb6cf12c3cff8300e21cefad30e7ec709abb4c24ac7318e6a52462a", size = 9968804, upload-time = "2026-07-22T22:18:03.903Z" }, + { url = "https://files.pythonhosted.org/packages/49/97/7886c89a39045c69ad82cbceaf3343810480c8ef49a216319ce8183860a6/pandas-3.0.5-cp311-cp311-win_arm64.whl", hash = "sha256:679f4e85b30ddb1515458ab1e788d3e260eae369b1f78da7a3aa4cac8ebf4a2a", size = 9205447, upload-time = "2026-07-22T22:18:06.134Z" }, + { url = "https://files.pythonhosted.org/packages/1c/54/1dc810ea558d1320b597aa140a514f2fdf1d2ea09c38cf556f13ea712ec9/pandas-3.0.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fa290c16964d4963fbfbc358928239cf3bd755b20e988ce944877def2f44471d", size = 10411717, upload-time = "2026-07-22T22:18:08.307Z" }, + { url = "https://files.pythonhosted.org/packages/68/56/fbe81c09195924d8b7b8d4461a20458fe80a6a5ed6b24f0314da684277e1/pandas-3.0.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2e26bb46934b8a2ca0c3de1d3d606fc5f6746584791b2db264d58cf370e08dc", size = 9957095, upload-time = "2026-07-22T22:18:10.6Z" }, + { url = "https://files.pythonhosted.org/packages/e0/51/fac252f4a913ed5eabf3c11b880a9e8d5a6c10f0b2129d0462212d238b4d/pandas-3.0.5-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:73fa87b08a7ef706f8aafda39ddaccf2a99047bea62d8c88a0361bcafb2237bc", size = 10485458, upload-time = "2026-07-22T22:18:12.834Z" }, + { url = "https://files.pythonhosted.org/packages/12/98/e976540c1addf70442be7842a18cf70884a964abbf69442504f4d2939989/pandas-3.0.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d373ce03ffd84010ed9839fa73672a9c8256990532e158440c0085db7d914b34", size = 10998091, upload-time = "2026-07-22T22:18:15.209Z" }, + { url = "https://files.pythonhosted.org/packages/a4/8c/1f29b5be8d3fc47dd7567eb167fabba2085879b31e0287ce7cba6d3d2ff4/pandas-3.0.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2a29c53d85ea98c5e792c59ef82ee9fbe6ca902c0d0adb6b23f45ef894cd7bf6", size = 11499501, upload-time = "2026-07-22T22:18:17.689Z" }, + { url = "https://files.pythonhosted.org/packages/9d/e2/bd9c98ad2df7b38bde002adde4cdf353519da51881634323b126c55997f9/pandas-3.0.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a5ad3b02ed6bc7d7ae9b70804b2c6aa31827489d150f8e623ce82491b82085d7", size = 12060559, upload-time = "2026-07-22T22:18:20.147Z" }, + { url = "https://files.pythonhosted.org/packages/f3/9a/ffbd852d58bd74a617fe2f8ee6a58a96982271ce41cf981eab22190b4a4b/pandas-3.0.5-cp312-cp312-pyemscripten_2024_0_wasm32.whl", hash = "sha256:b2acb4650527eec6822c3dadb2b771277b65e7dae7a267d4bccf65fd1bb3fbce", size = 7197652, upload-time = "2026-07-22T22:18:22.502Z" }, + { url = "https://files.pythonhosted.org/packages/70/b5/d2d3e9ae73362ba4229651b0ee1455cf78073a1ce585f6ff693782ce263e/pandas-3.0.5-cp312-cp312-win_amd64.whl", hash = "sha256:80a611068e8a3ac23f7398c6c14eb46dc974e5cc9997f653e2dcfd1da74edd41", size = 9831691, upload-time = "2026-07-22T22:18:24.534Z" }, + { url = "https://files.pythonhosted.org/packages/52/51/dea1e89d6a6796b9c43f85a09b484ee03edb8a4c4842e73e200a8c11301c/pandas-3.0.5-cp312-cp312-win_arm64.whl", hash = "sha256:25ff585b972a18ef1fe9ffa3ac6544d9950508aa76832e5147640b6022821e49", size = 9105796, upload-time = "2026-07-22T22:18:27.064Z" }, + { url = "https://files.pythonhosted.org/packages/bf/09/7b95c4a0025227d6f118c4039b423412ac6a982db02864166185d812fbc7/pandas-3.0.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c1c05a767fe8e5b4fe9e1c29806829c582052eaedb9120a3da83ba3f69e24a5b", size = 10385742, upload-time = "2026-07-22T22:18:29.346Z" }, + { url = "https://files.pythonhosted.org/packages/8d/0c/dc78fd8c4da477b4b5e8ad37295af352190d21ef63a9ee1bc071753074cc/pandas-3.0.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:b86765f268b56f7e665b93bce9d5df69dee7f99e595cf8fb839483ab315942a3", size = 9932067, upload-time = "2026-07-22T22:18:31.833Z" }, + { url = "https://files.pythonhosted.org/packages/3e/71/3592c055cf44df9808550f9368ceda80ff2b224d355ef73fe251dcda1802/pandas-3.0.5-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c597ecf5616b5c420372c1d4d4c00dbbfba7398bea857dcc984347e1ea48417b", size = 10466756, upload-time = "2026-07-22T22:18:34.195Z" }, + { url = "https://files.pythonhosted.org/packages/e3/70/4363150359f95b4cb4bcbb34ca23572bb5495749a621a8f3d5a1ddfd293c/pandas-3.0.5-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4b11c36e218331d0387cbe3a0a5f75162357a1d92d57b2b08a336ff94b19b2be", size = 10938525, upload-time = "2026-07-22T22:18:36.81Z" }, + { url = "https://files.pythonhosted.org/packages/f7/d0/317e7a0c67c0e69fa905a0161409397a7dc2d46ff611f6ca4803352c042b/pandas-3.0.5-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cf52e1f61d229496da17dc7ab54acdee627357e7008fd4fecba3d0ba2937fa58", size = 11489303, upload-time = "2026-07-22T22:18:39.287Z" }, + { url = "https://files.pythonhosted.org/packages/f1/8d/36dade89b49e4f9d5cbdbe863772581f98c0c6d78fc39ad4c557f6f2e17e/pandas-3.0.5-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:db172144bb56422bd157812f3b021eacc255451470b31e2c633c349490a1cfee", size = 11989004, upload-time = "2026-07-22T22:18:42.208Z" }, + { url = "https://files.pythonhosted.org/packages/9c/ba/18c4ec8a746e177da05a9e7a7963781d8ea195780724f854601b6ebd6b78/pandas-3.0.5-cp313-cp313-win_amd64.whl", hash = "sha256:0d298e951f23016ce4699951d044ae6418dbc91bf68cefca0f77666fcbb4e5c6", size = 9826896, upload-time = "2026-07-22T22:18:44.539Z" }, + { url = "https://files.pythonhosted.org/packages/de/ec/28a57266b753799a87b8bc79e7887ac6fd981b8c6d2978a0b7e7b6bd708c/pandas-3.0.5-cp313-cp313-win_arm64.whl", hash = "sha256:66266d3442a5e8b3c90274c2b8b230bee42dd1c286bc822cc2f9f2c7e12b883e", size = 9094790, upload-time = "2026-07-22T22:18:47.468Z" }, + { url = "https://files.pythonhosted.org/packages/51/2f/cf6aae281264f4463f0875bcbb15fd2bb6d291cc535187dad1732475e4a9/pandas-3.0.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2f264fc46911cc8131a7322a16199bbf8e353d27c10bb211f5bd0c814324dc36", size = 10390034, upload-time = "2026-07-22T22:18:49.818Z" }, + { url = "https://files.pythonhosted.org/packages/06/ec/5189518c7a7659c4bdcc6b1eb32c46c6f3c86b0661ffd84143d1112c7732/pandas-3.0.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:53730687fcd161883b24e10411c06d6a4c0f2275d2faf3bb2bc25deb4ba8007c", size = 9980065, upload-time = "2026-07-22T22:18:52.249Z" }, + { url = "https://files.pythonhosted.org/packages/ea/f1/598503ce8d7e3c35601e0747ba288c7864baae66380725bc12f13f884dfe/pandas-3.0.5-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:960d3ebcf249f75206899fcd2c6de53f736b7265759ced0d3e559df0b8b709b0", size = 10545532, upload-time = "2026-07-22T22:18:54.813Z" }, + { url = "https://files.pythonhosted.org/packages/fa/de/ceae2adf7034e07e9910299fe412e1819c4f0dd520700a888bcb03625448/pandas-3.0.5-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e94c2c5ca43bd3ca32bf64d32308887b65e5f9bfd8023ea52755107a999f93b", size = 10963120, upload-time = "2026-07-22T22:18:57.42Z" }, + { url = "https://files.pythonhosted.org/packages/66/25/86e0f4451874eb79e688deeebe3c451fec4557f8952005818d800ee8ac7e/pandas-3.0.5-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e819dd5f62966b481a8cb649d3299ebd886a1ea91ed5a99bf7ce77c98d18ab94", size = 11563178, upload-time = "2026-07-22T22:18:59.729Z" }, + { url = "https://files.pythonhosted.org/packages/f3/45/8643daa3b4147e433adfcccefdd0380d3aad79d86b15d8999730fe1944d5/pandas-3.0.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:3c5ed2e7c06e91d340dfd091d7934f9bc82e4a36b95f647f090b9d1c9ac649da", size = 12028708, upload-time = "2026-07-22T22:19:02.164Z" }, + { url = "https://files.pythonhosted.org/packages/96/58/ad979ae617615576e8aafd569c9d4b62f1191d896e38f51d66ba06f3b89a/pandas-3.0.5-cp314-cp314-win_amd64.whl", hash = "sha256:cd8f7c6dc98527058ee6264219343f5392240a6f1bfa654fc5d79023020d0c92", size = 9951806, upload-time = "2026-07-22T22:19:04.596Z" }, + { url = "https://files.pythonhosted.org/packages/69/32/7ac03886b304049a9d2625ee88f59af760d8a93bd30ed9239bce7b9869a8/pandas-3.0.5-cp314-cp314-win_arm64.whl", hash = "sha256:5183427f5a8156d480f30333777bc978be93650a49a7c01db26adffe95b31e85", size = 9238297, upload-time = "2026-07-22T22:19:06.836Z" }, + { url = "https://files.pythonhosted.org/packages/be/ed/1d1f2ee5547d5167face2376d11c8b2a4c7bfff5a416ee7a9046891fab1e/pandas-3.0.5-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:303da736987d481074ca720ada325f8bd80c64ebc2d45ed79b29df3aaa4a26ca", size = 10849690, upload-time = "2026-07-22T22:19:09.391Z" }, + { url = "https://files.pythonhosted.org/packages/57/55/17e17152e98fbb0c4b1e562bc65387a2f20a80db0f4a86bf8d3a0e4248d4/pandas-3.0.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3b2801bbb049d0136f6c213eae02b5fca969384fc2064dd728d8620552aa49da", size = 10509945, upload-time = "2026-07-22T22:19:11.773Z" }, + { url = "https://files.pythonhosted.org/packages/88/90/817d44dbf83facf9556f33576d9af0a241981e7bb5c00606c0bcb5df8dda/pandas-3.0.5-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cce3a9d11d2b1f82c69a27ec1f4948a170e2c403c4bbfa8cca62e3fdebe2ef3a", size = 10392197, upload-time = "2026-07-22T22:19:14.024Z" }, + { url = "https://files.pythonhosted.org/packages/f1/da/889f00c0a6f5aa1545add70abbf01502dff87ab577adb855bd631c54d2f2/pandas-3.0.5-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ef01af4d8dc6cd2c8d6c7736f149574ef93fe043811eeb5e445f2647154b5040", size = 10862726, upload-time = "2026-07-22T22:19:16.351Z" }, + { url = "https://files.pythonhosted.org/packages/bc/98/f1e934fb3c98fce859c6147c6785816c7b5b9ab7821115c5d8c4de9842b9/pandas-3.0.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e2759e890db96dfcffdbd9b86c3c2cb6afaf58def482820317e06163ec1066cd", size = 11414864, upload-time = "2026-07-22T22:19:18.981Z" }, + { url = "https://files.pythonhosted.org/packages/fe/be/d448af7d657d82e1888dd8551f79c6d6fb161080b5b9752d84d910ec2319/pandas-3.0.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b58b1b39d46a5862e3fb18f50d1a201398619d16a0f9f73f57eea5583cf0e63c", size = 11925105, upload-time = "2026-07-22T22:19:21.515Z" }, + { url = "https://files.pythonhosted.org/packages/29/c1/ccb4238212c8c4f496c584f3044d94e0c030ed8e1d68999db46c91c2242f/pandas-3.0.5-cp314-cp314t-win_amd64.whl", hash = "sha256:1c10461f6eeb35d8f05b6184c65c8b9991663b66c46b1d559b682cb34ae7c6ea", size = 10387612, upload-time = "2026-07-22T22:19:24.257Z" }, + { url = "https://files.pythonhosted.org/packages/d2/cf/6a51b2c38980e04c279fd2fa908a1b0982064e860444acfca4ec2e2c8359/pandas-3.0.5-cp314-cp314t-win_arm64.whl", hash = "sha256:3c5015fd1730fbf883647e88068176c839c102cea883ba1769a6f4593bfc1f8c", size = 9509776, upload-time = "2026-07-22T22:19:26.694Z" }, +] + [[package]] name = "pandocfilters" version = "1.5.1" @@ -2925,6 +2992,7 @@ dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "optuna" }, + { name = "pandas" }, { name = "rosettasciio" }, { name = "scikit-image" }, { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -2969,6 +3037,7 @@ requires-dist = [ { name = "matplotlib" }, { name = "numpy", specifier = ">2" }, { name = "optuna", specifier = ">=4.5.0" }, + { name = "pandas", specifier = ">=2.2" }, { name = "quantem-widget", marker = "extra == 'widgets'", editable = "widget" }, { name = "rosettasciio", specifier = ">=0.8.0" }, { name = "scikit-image", specifier = ">=0.25.2" }, @@ -3360,7 +3429,9 @@ name = "scipy" version = "1.17.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.12'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } }, @@ -3434,8 +3505,12 @@ name = "scipy" version = "1.18.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14'", - "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } }, @@ -3644,7 +3719,9 @@ name = "tifffile" version = "2026.3.3" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.12'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } }, @@ -3659,8 +3736,12 @@ name = "tifffile" version = "2026.7.14" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14'", - "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } }, @@ -4014,7 +4095,9 @@ name = "zarr" version = "3.1.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.12'", + "python_full_version < '3.12' and sys_platform == 'win32'", + "python_full_version < '3.12' and sys_platform == 'emscripten'", + "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "donfig" }, @@ -4034,8 +4117,12 @@ name = "zarr" version = "3.2.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14'", - "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ { name = "donfig" }, From 1dc7dc2a99a190da6e0c390ae66a6eac3a3359fc Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 12:35:20 -0700 Subject: [PATCH 09/11] refactor: simplify DriftCorrection public API (#22) * feat: add reference-based spectrum-image correction * fix: reject ambiguous spectrum-image reference matches * feat: propagate drift across 4D-STEM scan axes * test: protect 4D-STEM detector coordinates * test: compare full detector signature field * test: verify 4D-STEM coverage and output residency * fix: stream native 4D-STEM scan correction * test: compare native detector frames across backends * test: set float32 backend parity tolerance * feat: preserve calibration during drift downsampling * test: lock downsampled coordinate provenance * perf: add average-pooled affine pyramid search * test: compare native and pyramid affine rates * feat: add strip-wise residual correction primitives * fix: make nonrigid regularization portable across torch backends * feat: propagate residual fields to calibrated element maps * test: lock residual-correction publication contracts * test: recover known strip residual shifts * feat: complete automatic affine manuscript parity * feat: complete multi-knot manuscript parity * test: verify complete manuscript workflow parity * fix: keep 4D-STEM integration usable without optional GPU package * refactor: adopt the scientist-facing Drift API * test: verify manuscript output frame contracts * fix: preserve natural corrected output frames * feat: expose manual translation alignment --- src/quantem/imaging/__init__.py | 6 +- src/quantem/imaging/drift/__init__.py | 50 ++++++++- src/quantem/imaging/drift/apply.py | 50 ++++++--- src/quantem/imaging/drift/core/warping.py | 124 ++++++++++++++++++++++ src/quantem/imaging/drift/correction.py | 11 +- src/quantem/imaging/drift/plot.py | 1 + tests/imaging/test_drift.py | 65 +++++++++++- 7 files changed, 281 insertions(+), 26 deletions(-) diff --git a/src/quantem/imaging/__init__.py b/src/quantem/imaging/__init__.py index 835c3011..2286d509 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,10 +1,12 @@ +"""Imaging tools for scientific image analysis.""" + from quantem.imaging.drift import ( CorrectionResult as CorrectionResult, DriftCorrection as DriftCorrection, StripPass as StripPass, -) -from quantem.imaging.drift import ( pair_spectrum_image_references as pair_spectrum_image_references, +) +from quantem.imaging.drift.io import ( read_emd as read_emd, read_emd_eds as read_emd_eds, read_emd_metadata as read_emd_metadata, diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py index 2049d242..585a04ca 100644 --- a/src/quantem/imaging/drift/__init__.py +++ b/src/quantem/imaging/drift/__init__.py @@ -1,6 +1,43 @@ +"""Drift-correction public package surface. + +**Primary user API** - one class, chainable stages:: + + from quantem.imaging.drift import DriftCorrection + + dc = DriftCorrection.from_emd(f0, f1) # angles from EMD metadata + dc.correct_affine() # automatic; no solve knobs required + dc.plot_combined(stage=("initial", "affine"), interactive=True) + dc.report() + dc.show() + dc.save("drift.zip", mode="o") + +Optional residual stages after affine:: + + dc.correct_strip(...) # piecewise-rigid bands + dc.correct_nonrigid(...) # per-scanline polish (tiny max_image_shift on lattices) + +Manual rigid registration remains available when needed:: + + dc.align_translation(max_image_shift=32) + +Also: :meth:`~DriftCorrection.from_images`, +:meth:`~DriftCorrection.from_reference`, +:meth:`~DriftCorrection.from_4dstem`, and an explicit +:meth:`~DriftCorrection.preprocess` when you need a fixed canvas. + +Advanced troubleshooting stays on the same object through +``diagnose_affine()`` and ``diagnose_nonrigid()``. Numerical code lives under +``drift.core`` and is not a notebook entry point. + +Free residual helpers are not re-exported here; use ``dc.correct_strip()``. +""" + from quantem.imaging.drift.correction import DriftCorrection as DriftCorrection from quantem.imaging.drift.core.strip import StripPass as StripPass from quantem.imaging.drift.fourdstem import CorrectionResult as CorrectionResult +from quantem.imaging.drift import plot as plot +from quantem.imaging.drift import io as io +from quantem.imaging.drift import fourdstem as fourdstem from quantem.imaging.drift.io import ( pair_spectrum_image_references as pair_spectrum_image_references, read_emd as read_emd, @@ -8,4 +45,15 @@ read_emd_metadata as read_emd_metadata, ) -DriftCorrection.__module__ = __name__ +__all__ = [ + "DriftCorrection", + "CorrectionResult", + "StripPass", + "plot", + "io", + "fourdstem", + "pair_spectrum_image_references", + "read_emd", + "read_emd_eds", + "read_emd_metadata", +] diff --git a/src/quantem/imaging/drift/apply.py b/src/quantem/imaging/drift/apply.py index efc73d25..b2c00ba1 100644 --- a/src/quantem/imaging/drift/apply.py +++ b/src/quantem/imaging/drift/apply.py @@ -1,6 +1,7 @@ """Apply a solved drift field and return corrected scientific data.""" from copy import deepcopy +from typing import Literal import numpy as np import torch @@ -141,7 +142,7 @@ def corrected( self, *, upsample_factor: int = 2, - output_original_shape: bool | None = None, + output_frame: Literal["auto", "input", "canvas"] = "auto", strip_padding: bool = False, smoothing_sigma: float | None = 0.5, stage: str | None = None, @@ -158,10 +159,12 @@ def corrected( ---------- upsample_factor : int, default 2 Sampling multiplier for the corrected image. - output_original_shape : bool or None, default None - ``None`` keeps the natural frame: the padded solver canvas for paired - 2-D scans and the native input frame for reference datasets. ``True`` - returns the original image shape; ``False`` keeps the solver canvas. + output_frame : {"auto", "input", "canvas"}, default "auto" + ``"auto"`` returns the padded solver canvas for paired 2-D scans and + the native input frame for a channel-resolved reference dataset. + ``"input"`` always returns the original scan shape. ``"canvas"`` + retains the padded solver canvas used for publication figures and + when chaining a solved correction as a structural reference. strip_padding : bool, default False Remove pixels that do not share measured coverage across scans. smoothing_sigma : float or None, default 0.5 @@ -183,17 +186,21 @@ def corrected( >>> corrected = drift.corrected() >>> scans = drift.corrected(merge=False) """ - automatic_output_frame = output_original_shape is None - output_original_shape = ( - self._reference_mode - if automatic_output_frame - else bool(output_original_shape) - ) + if output_frame not in {"auto", "input", "canvas"}: + raise ValueError( + "output_frame must be 'auto', 'input', or 'canvas', " + f"got {output_frame!r}" + ) + resolved_output_frame = ( + "input" if self._reference_mode else "canvas" + ) if output_frame == "auto" else output_frame + if strip_padding and resolved_output_frame != "input": + raise ValueError("strip_padding=True requires output_frame='input'") if self._reference_mode: if ( upsample_factor != 2 - or not output_original_shape + or resolved_output_frame != "input" or strip_padding or smoothing_sigma != 0.5 or not merge @@ -229,7 +236,6 @@ def corrected( if not merge: if ( upsample_factor != 2 - or (not automatic_output_frame and not output_original_shape) or strip_padding or smoothing_sigma != 0.5 ): @@ -239,6 +245,18 @@ def corrected( "controls apply only to the merged image." ) panels = comparison_panels(self, stage) + corrected_scans = panels["corrected_scans"] + if resolved_output_frame == "input": + scan_shape = tuple(int(value) for value in self.imgs[0].shape[:2]) + row, column = padding_offset( + corrected_scans[0].shape[:2], + scan_shape, + integer=True, + ) + corrected_scans = [ + array[row : row + scan_shape[0], column : column + scan_shape[1]] + for array in corrected_scans + ] return [ Dataset2d.from_array( np.asarray(array), @@ -247,7 +265,7 @@ def corrected( sampling=self.imgs[0].sampling, units=self.imgs[0].units, ) - for index, array in enumerate(panels["corrected_scans"]) + for index, array in enumerate(corrected_scans) ] device = self._device @@ -275,7 +293,7 @@ def corrected( output_shape = ( tuple(int(value) for value in self.imgs[0].shape[:2]) - if output_original_shape + if resolved_output_frame == "input" else tuple(int(value) for value in self.shape[-2:]) ) image_corr_fft = fourier_crop_torch( @@ -283,7 +301,7 @@ def corrected( output_shape, ) / upsample_factor**2 corrected_array = torch.fft.ifft2(image_corr_fft).real.cpu().numpy() - if strip_padding and output_original_shape: + if strip_padding: scan_h, scan_w = self.imgs[0].shape[:2] pad_h, pad_w = padding_offset(corrected_array.shape[:2], (scan_h, scan_w), integer=True) corrected_array = corrected_array[pad_h : pad_h + scan_h, pad_w : pad_w + scan_w] diff --git a/src/quantem/imaging/drift/core/warping.py b/src/quantem/imaging/drift/core/warping.py index d8926426..2cd24d72 100644 --- a/src/quantem/imaging/drift/core/warping.py +++ b/src/quantem/imaging/drift/core/warping.py @@ -390,6 +390,130 @@ def render(warped_t, weights_t): return (warped_t, weights_t) if return_weights else warped_t +def align_translation( + self, + *, + max_image_shift: float | None | str = "auto", + fixed_scans: list[int] | None = None, + show_combined: bool = True, + show_scans: bool = False, + show_knots: bool = True, + show_knot_plot: bool = False, + show_report: bool = False, + verbose: bool = True, +): + """Align scans by a global translation without estimating drift. + + Use this manual stage when the acquisitions differ by a rigid offset but + the scan geometry should remain unchanged. Translation is registration, + so it intentionally keeps the ``align_`` verb; affine, strip, and + non-rigid deformation stages use the ``correct_`` verb. + + Parameters + ---------- + max_image_shift : float, None, or "auto", default "auto" + Maximum allowed translation in pixels. ``"auto"`` derives a bound + from the scan dimensions. ``None`` allows any translation. + fixed_scans : list of int or None, default None + Scan indices that remain fixed. Reference-based corrections keep scan + 0 fixed automatically. + show_combined : bool, default True + Display the combined registration view after alignment. + show_scans : bool, default False + Display the individually aligned scans. + show_knots : bool, default True + Show scan-line origins on requested views. + show_knot_plot : bool, default False + Show the standalone scan-line-origin displacement plot. + show_report : bool, default False + Print common-coverage registration measurements. + verbose : bool, default True + Print the recovered ``(row, col)`` translation for each scan. + + Returns + ------- + DriftCorrection + The same correction object for method chaining. + + Examples + -------- + >>> drift = DriftCorrection.from_emd(scan_0_path, scan_1_path) + >>> drift.align_translation(max_image_shift=32) + """ + if not hasattr(self, "_initial_knots"): + self.preprocess( + verbose=False, + show_combined=False, + show_scans=False, + show_knots=False, + ) + if self.shape[0] < 2: + raise ValueError( + "align_translation requires at least 2 images; " + f"got {self.shape[0]}." + ) + if fixed_scans is None and self._reference_mode: + fixed_scans = [0] + fixed_set = frozenset(fixed_scans) if fixed_scans is not None else frozenset() + if max_image_shift == "auto": + reference_alignment = bool(fixed_set) + max_image_shift = max( + 16.0, + min(self.imgs[0].shape[:2]) + * (0.0625 if reference_alignment else 0.25), + ) + elif isinstance(max_image_shift, str): + raise ValueError( + "max_image_shift must be 'auto', None, or a non-negative " + f"number; got {max_image_shift!r}." + ) + elif max_image_shift is not None: + max_image_shift = float(max_image_shift) + if not np.isfinite(max_image_shift) or max_image_shift < 0: + raise ValueError( + "max_image_shift must be 'auto', None, or a non-negative " + f"number; got {max_image_shift!r}." + ) + + knots_before = [knots.detach().clone() for knots in self.knots] + warp_and_translate( + self, + max_image_shift=max_image_shift, + upsample_factor=8, + fixed_indices=fixed_set, + ) + translations = np.asarray( + [ + (after - before).mean(dim=(1, 2)).detach().cpu().numpy() + for before, after in zip(knots_before, self.knots, strict=True) + ], + dtype=np.float64, + ) + self._images_warped_stale = False + self._warped_fingerprint = knot_fingerprint(self) + if verbose: + for index, (row_shift, column_shift) in enumerate(translations): + print( + f"align_translation: scan {index} shifted " + f"({row_shift:+.3f}, {column_shift:+.3f}) px (row, col)" + ) + + import quantem.imaging.drift.plot as drift_plot + + drift_plot.show_after_step( + self, + "translation", + show_combined=show_combined, + show_scans=show_scans, + show_knots=show_knots, + ) + if show_knot_plot: + self.plot_knots() + if show_report: + print(self.report().to_string()) + return self + + def backward_warp( images: torch.Tensor, drift: tuple[float, float] | torch.Tensor, diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index 57e73084..0f99a8d4 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -5,6 +5,7 @@ and 4D-STEM helpers in ``fourdstem``. Notebooks should import ``DriftCorrection`` from ``quantem.imaging`` or ``quantem.imaging.drift`` and use ``from_emd`` / ``correct_affine`` / ``plot_combined`` / ``show`` / ``save``. +Manual rigid registration remains available through ``align_translation``. """ from typing import Self @@ -17,6 +18,7 @@ import quantem.imaging.drift.core.affine as affine import quantem.imaging.drift.core.nonrigid as nonrigid import quantem.imaging.drift.core.strip as strip +import quantem.imaging.drift.core.warping as warping import quantem.imaging.drift.diagnostics as diagnostics import quantem.imaging.drift.fourdstem as fourdstem import quantem.imaging.drift.plot as drift_plot @@ -34,8 +36,10 @@ class DriftCorrection(AutoSerialize): Aligns scans at different scan directions, recovers per-scanline drift, and forms a corrected product (HAADF pairs, reference EDS/EELS, 4D-STEM). - Typical chain: :meth:`from_emd` → :meth:`correct_affine` → :meth:`plot_combined` - / :meth:`show` / :meth:`save`. Residual polish: :meth:`correct_strip`, + Typical chain: :meth:`from_emd` → :meth:`correct_affine` → + :meth:`plot_combined` / :meth:`show` / :meth:`save`. Use + :meth:`align_translation` when a manual rigid-registration stage is + required. Residual polish: :meth:`correct_strip`, :meth:`correct_nonrigid` (use small ``max_image_shift`` on lattices). """ @@ -353,7 +357,7 @@ def from_reference( verbose=False, ) reference_image = preparation.match_reference_image( - reference_dc.corrected().array, + reference_dc.corrected(output_frame="canvas").array, tuple(int(value) for value in reference_dc.imgs[0].shape[:2]), tuple(int(value) for value in drifted_shape), ) @@ -384,6 +388,7 @@ def from_reference( probe_positions = fourdstem.probe_positions preprocess = preparation.preprocess + align_translation = warping.align_translation correct_affine = affine.correct_affine correct_strip = strip.correct_strip diff --git a/src/quantem/imaging/drift/plot.py b/src/quantem/imaging/drift/plot.py index 7f8288ab..a4d59123 100644 --- a/src/quantem/imaging/drift/plot.py +++ b/src/quantem/imaging/drift/plot.py @@ -440,6 +440,7 @@ def _stage_description(stage: str | None) -> str: None: "current correction", "initial": "raw", "raw": "raw", + "translation": "after translation alignment", "affine": "after affine correction", "strip": "after strip correction", "nonrigid": "after non-rigid correction", diff --git a/tests/imaging/test_drift.py b/tests/imaging/test_drift.py index a1f72acd..6f74a428 100644 --- a/tests/imaging/test_drift.py +++ b/tests/imaging/test_drift.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from scipy.ndimage import gaussian_filter +from scipy.ndimage import gaussian_filter, shift from quantem.core.datastructures.dataset2d import Dataset2d from quantem.imaging import DriftCorrection @@ -37,10 +37,12 @@ def _orthogonal_pair(size: int = 64) -> tuple[Dataset2d, Dataset2d]: def test_primary_api_uses_scientist_facing_names(): """The final PR exposes one concise chain without legacy solve names.""" preprocess = inspect.signature(DriftCorrection.preprocess).parameters + translation = inspect.signature(DriftCorrection.align_translation).parameters affine = inspect.signature(DriftCorrection.correct_affine).parameters nonrigid = inspect.signature(DriftCorrection.correct_nonrigid).parameters assert {"padding_fraction", "smoothing_sigma", "num_knots"} <= set(preprocess) + assert {"max_image_shift", "fixed_scans"} <= set(translation) assert {"max_drift_rate", "num_rates", "region"} <= set(affine) assert {"num_knots", "num_refine_cycles", "knot_smoothing_sigma"} <= set( nonrigid @@ -95,7 +97,7 @@ def test_metadata_driven_affine_workflow_returns_calibrated_dataset(): ) corrected = drift.corrected( upsample_factor=1, - output_original_shape=True, + output_frame="input", verbose=False, ) @@ -107,15 +109,70 @@ def test_metadata_driven_affine_workflow_returns_calibrated_dataset(): assert corrected.units == scan_0.units assert len(drift.drift_rate) == 2 - automatic = drift.corrected(verbose=False) + automatic = drift.corrected(upsample_factor=1, verbose=False) canvas = drift.corrected( - output_original_shape=False, + upsample_factor=1, + output_frame="canvas", verbose=False, ) assert automatic.shape == tuple(drift.shape[-2:]) assert canvas.shape == tuple(drift.shape[-2:]) +def test_manual_translation_alignment_reduces_global_offset(): + """Manual translation alignment registers scans without fitting drift.""" + rng = np.random.default_rng(7) + reference = gaussian_filter( + rng.normal(size=(64, 64)).astype(np.float32), + sigma=1.5, + ) + moving = np.rot90( + shift( + reference, + shift=(3.0, -4.0), + order=1, + mode="constant", + cval=float(np.median(reference)), + ), + k=1, + ).astype(np.float32) + drift = DriftCorrection.from_images( + reference, + moving, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ) + + returned = drift.align_translation( + max_image_shift=8, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + initial = drift.corrected( + stage="initial", + merge=False, + output_frame="input", + verbose=False, + ) + aligned = drift.corrected( + merge=False, + output_frame="input", + verbose=False, + ) + interior = np.s_[8:-8, 8:-8] + initial_error = np.mean( + np.abs(initial[0].array[interior] - initial[1].array[interior]) + ) + aligned_error = np.mean( + np.abs(aligned[0].array[interior] - aligned[1].array[interior]) + ) + + assert returned is drift + assert aligned_error < initial_error * 0.6 + + def test_nonrigid_diagnostic_defines_fast_roughness(): """The difficult multi-knot diagnostic states exactly what roughness means.""" doc = DriftCorrection.diagnose_nonrigid.__doc__ or "" From af7160ec2e06f32aa20db8f6502915c18b8139e9 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Tue, 25 Aug 2026 12:55:31 -0700 Subject: [PATCH 10/11] refactor: rename drift preparation module (#23) --- src/quantem/imaging/drift/core/affine.py | 10 +++++----- src/quantem/imaging/drift/correction.py | 16 ++++++++-------- .../drift/{preparation.py => preprocess.py} | 2 +- tests/imaging/test_drift_downsampling.py | 2 +- 4 files changed, 15 insertions(+), 15 deletions(-) rename src/quantem/imaging/drift/{preparation.py => preprocess.py} (99%) diff --git a/src/quantem/imaging/drift/core/affine.py b/src/quantem/imaging/drift/core/affine.py index 502f9a90..94285517 100644 --- a/src/quantem/imaging/drift/core/affine.py +++ b/src/quantem/imaging/drift/core/affine.py @@ -9,7 +9,7 @@ import quantem.imaging.drift.apply as drift_apply import quantem.imaging.drift.plot as drift_plot -import quantem.imaging.drift.preparation as preparation +import quantem.imaging.drift.preprocess as preprocessing import quantem.imaging.drift.report as report from quantem.imaging.drift.core import knots as drift_knots from quantem.imaging.drift.core.warping import ( @@ -94,7 +94,7 @@ def _scan_disagreement(warped: torch.Tensor, scan_axis: int) -> torch.Tensor: def _downsampled_correction(correction, factor): """Build the same correction problem on average-pooled scan images.""" images = [ - preparation.average_downsample_2d(np.asarray(image.array), factor) + preprocessing.average_downsample_2d(np.asarray(image.array), factor) for image in correction.imgs ] if correction._reference_mode: @@ -460,10 +460,10 @@ def correct_affine( normalize = False normalization_reason = "4dstem_collection" else: - normalize, normalization_reason = preparation.automatic_alignment_normalization( + normalize, normalization_reason = preprocessing.automatic_alignment_normalization( self.imgs ) - padding = preparation.minimum_affine_padding_fraction( + padding = preprocessing.minimum_affine_padding_fraction( tuple(int(value) for value in self.imgs[0].shape[:2]), self.scan_direction_degrees, planned_rate, @@ -820,7 +820,7 @@ def automatic_affine_search( starting_knots = [knot.clone() for knot in self.knots] initial_error_row = np.asarray(self.error_track[-1], dtype=np.float64).copy() image_shape = tuple(int(value) for value in self.imgs[0].shape[:2]) - factor = preparation.resolve_downsample( + factor = preprocessing.resolve_downsample( pyramid_downsample, image_shape, ) diff --git a/src/quantem/imaging/drift/correction.py b/src/quantem/imaging/drift/correction.py index 0f99a8d4..aaeecb4f 100644 --- a/src/quantem/imaging/drift/correction.py +++ b/src/quantem/imaging/drift/correction.py @@ -22,7 +22,7 @@ import quantem.imaging.drift.diagnostics as diagnostics import quantem.imaging.drift.fourdstem as fourdstem import quantem.imaging.drift.plot as drift_plot -import quantem.imaging.drift.preparation as preparation +import quantem.imaging.drift.preprocess as preprocessing import quantem.imaging.drift.report as drift_report from quantem.core.config import validate_device from quantem.core.datastructures.dataset2d import Dataset2d @@ -81,7 +81,7 @@ def __init__( if not datasets and alignment_image is None: return - preparation.prepare_inputs( + preprocessing.prepare_inputs( self, datasets, scan_direction_degrees, alignment_image ) @@ -186,7 +186,7 @@ def from_emd(cls, *paths, verbose: bool = True, device: str | None = None) -> Se data = [read_emd(p) for p in paths] if len(data) == 2: - data[0], data[1] = preparation.match_scan_shapes( + data[0], data[1] = preprocessing.match_scan_shapes( data[0], data[1], verbose=verbose ) if verbose: @@ -324,8 +324,8 @@ def from_reference( reference_dc = reference_image if device is None: device = reference_dc.device - drifted_shape = preparation.input_array(drifted_dataset).shape[:2] - reference_downsample = preparation.reference_downsample( + drifted_shape = preprocessing.input_array(drifted_dataset).shape[:2] + reference_downsample = preprocessing.reference_downsample( tuple(int(value) for value in reference_dc.imgs[0].shape[:2]), tuple(int(value) for value in drifted_shape), reference_sampling=getattr(reference_dc.imgs[0], "sampling", None), @@ -339,7 +339,7 @@ def from_reference( # images avoids that alias while keeping the operation fully # automatic and costs well under a second on the microscope GPU. scaled_reference_images = [ - preparation.average_downsample_2d( + preprocessing.average_downsample_2d( np.asarray(image.array), reference_downsample, ) @@ -356,7 +356,7 @@ def from_reference( show_knots=False, verbose=False, ) - reference_image = preparation.match_reference_image( + reference_image = preprocessing.match_reference_image( reference_dc.corrected(output_frame="canvas").array, tuple(int(value) for value in reference_dc.imgs[0].shape[:2]), tuple(int(value) for value in drifted_shape), @@ -387,7 +387,7 @@ def from_reference( drift_field = fourdstem.drift_field probe_positions = fourdstem.probe_positions - preprocess = preparation.preprocess + preprocess = preprocessing.preprocess align_translation = warping.align_translation correct_affine = affine.correct_affine diff --git a/src/quantem/imaging/drift/preparation.py b/src/quantem/imaging/drift/preprocess.py similarity index 99% rename from src/quantem/imaging/drift/preparation.py rename to src/quantem/imaging/drift/preprocess.py index a2227328..63773395 100644 --- a/src/quantem/imaging/drift/preparation.py +++ b/src/quantem/imaging/drift/preprocess.py @@ -1,4 +1,4 @@ -"""Image and canvas preparation for drift alignment.""" +"""Preprocess scientific inputs for drift alignment.""" from collections.abc import Sequence diff --git a/tests/imaging/test_drift_downsampling.py b/tests/imaging/test_drift_downsampling.py index fdff9439..e6c4e52b 100644 --- a/tests/imaging/test_drift_downsampling.py +++ b/tests/imaging/test_drift_downsampling.py @@ -6,7 +6,7 @@ from quantem.core.datastructures.dataset2d import Dataset2d from quantem.imaging.drift import DriftCorrection -from quantem.imaging.drift.preparation import ( +from quantem.imaging.drift.preprocess import ( average_downsample_2d, resolve_downsample, ) From 02ed1fb69cce7d52781632434a2124d2c5b025d6 Mon Sep 17 00:00:00 2001 From: Sangjoon Bob Lee Date: Wed, 26 Aug 2026 01:33:11 -0700 Subject: [PATCH 11/11] feat: expose sparse XEDS window extraction --- src/quantem/imaging/drift/io.py | 114 +++++++++++++++++++++++++++++- src/quantem/imaging/drift/plot.py | 27 ++++--- tests/imaging/test_drift.py | 21 ++++++ tests/imaging/test_drift3d.py | 77 ++++++++++++++++++++ 4 files changed, 229 insertions(+), 10 deletions(-) diff --git a/src/quantem/imaging/drift/io.py b/src/quantem/imaging/drift/io.py index cef25850..f0f7f480 100644 --- a/src/quantem/imaging/drift/io.py +++ b/src/quantem/imaging/drift/io.py @@ -8,6 +8,107 @@ import numpy as np +def _decode_velox_json(dataset: h5py.Dataset) -> dict: + """Decode one JSON record stored as bytes or a padded uint8 array.""" + value = dataset[0] if dataset.dtype.kind in {"O", "S", "U"} else dataset[:] + if isinstance(value, str): + raw = value.encode() + elif isinstance(value, bytes): + raw = value + else: + raw = bytes(np.asarray(value, dtype=np.uint8).reshape(-1)) + raw = raw.replace(b"\x00", b"") + for candidate in (raw, _deinterleave_velox_metadata(raw)): + try: + return json.loads(candidate.decode("utf-8", errors="ignore")) + except json.JSONDecodeError: + continue + raise ValueError(f"Could not decode Velox JSON dataset {dataset.name}") + + +def _read_emd_energy_windows( + path: str | Path, + energy_windows: dict[str, tuple[float, float]], +) -> tuple[dict[str, np.ndarray], np.ndarray]: + """Count selected energy windows directly from compressed EDS streams.""" + windows: dict[str, tuple[float, float]] = {} + for name, limits in energy_windows.items(): + if len(limits) != 2: + raise ValueError(f"Energy window {name!r} must contain (low, high)") + low, high = (float(limits[0]), float(limits[1])) + if not np.isfinite(low) or not np.isfinite(high) or low > high: + raise ValueError(f"Invalid energy window {name!r}: {limits!r}") + windows[str(name)] = (low, high) + if not windows: + return {}, np.empty(0, dtype=np.float32) + + maps: dict[str, np.ndarray] | None = None + reference_shape: tuple[int, int] | None = None + reference_axis: np.ndarray | None = None + with h5py.File(path, "r") as handle: + stream_group = handle.get("Data/SpectrumStream") + if stream_group is None or not stream_group: + raise ValueError(f"{path} contains no Velox SpectrumStream data") + for stream in stream_group.values(): + settings = _decode_velox_json(stream["AcquisitionSettings"]) + raster = settings["RasterScanDefinition"] + shape = (int(raster["Height"]), int(raster["Width"])) + channels = int(settings["bincount"]) + if reference_shape is None: + reference_shape = shape + maps = { + name: np.zeros(shape[0] * shape[1], dtype=np.uint32) + for name in windows + } + elif shape != reference_shape: + raise ValueError("EDS detector streams have inconsistent scan shapes") + + metadata = _decode_velox_json(stream["Metadata"]) + detector_name = metadata["BinaryResult"]["Detector"] + detector = next( + ( + item + for item in metadata["Detectors"].values() + if item.get("DetectorName") == detector_name + ), + None, + ) + if detector is None: + raise ValueError(f"Missing calibration for EDS detector {detector_name}") + scale = np.float32(float(detector["Dispersion"]) / 1000.0) + offset = np.float32(float(detector["OffsetEnergy"]) / 1000.0) + axis = np.arange(channels, dtype=np.float32) * scale + offset + if reference_axis is None: + reference_axis = axis + elif not np.array_equal(axis, reference_axis): + raise ValueError("EDS detector streams have inconsistent energy axes") + + encoded = np.asarray(stream["Data"][:, 0], dtype=np.uint16) + gates = encoded == np.uint16(65535) + pixel_index = np.cumsum(gates, dtype=np.int32) + pixel_count = shape[0] * shape[1] + assert maps is not None + for name, (low, high) in windows.items(): + selected_channels = (axis >= low) & (axis <= high) + lookup = np.zeros(65536, dtype=bool) + lookup[:channels] = selected_channels + selected_events = (~gates) & lookup[encoded] + counts = np.bincount( + pixel_index[selected_events] % pixel_count, + minlength=pixel_count, + ) + maps[name] += counts[:pixel_count].astype(np.uint32, copy=False) + + assert maps is not None and reference_shape is not None and reference_axis is not None + return ( + { + name: values.reshape(reference_shape).astype(np.float32) + for name, values in maps.items() + }, + reference_axis, + ) + + def _axis_calibration( axes: list[dict], ndim: int, @@ -243,6 +344,7 @@ def read_emd_eds( path: str | Path, *, load_spectrum: bool = False, + energy_windows: dict[str, tuple[float, float]] | None = None, verbose: bool = True, ) -> dict[str, object]: """Read images from a Velox EDS/EELS spectrum-image EMD. @@ -261,6 +363,10 @@ def read_emd_eds( load_spectrum : bool, optional Load the full ``(row, col, energy)`` spectrum. The default ``False`` loads only the HAADF and stored 2-D elemental maps. + energy_windows : dict, optional + Named ``(low_keV, high_keV)`` intervals to count directly from the + compressed Velox SpectrumStream. This avoids expanding the full + spectrum cube. Window endpoints are inclusive. verbose : bool, optional Print a compact summary of the loaded images. @@ -328,8 +434,13 @@ def read_emd_eds( if spectrum_stream is None else _spectrum_dataset(spectrum_stream, path, metadata) ) + window_maps: dict[str, np.ndarray] = {} + if energy_windows is not None: + window_maps, window_axis = _read_emd_energy_windows(path, energy_windows) + else: + window_axis = None if spectrum is None: - energy_axis = None + energy_axis = window_axis else: scale = float(spectrum.sampling[-1]) offset = float(spectrum.origin[-1]) @@ -354,6 +465,7 @@ def read_emd_eds( "cube": None if spectrum is None else spectrum.array, "energy_axis_keV": energy_axis, "element_maps": element_maps, + "window_maps": window_maps, "scan_rotation_deg": scan_rot, "pixel_size_nm": px_nm, "shape": shape, diff --git a/src/quantem/imaging/drift/plot.py b/src/quantem/imaging/drift/plot.py index a4d59123..70447eae 100644 --- a/src/quantem/imaging/drift/plot.py +++ b/src/quantem/imaging/drift/plot.py @@ -878,15 +878,21 @@ def static_comparison(correction, zoom: float, size: int, **kwargs): elif vmin is not None or vmax is not None: norm = {"vmin": vmin, "vmax": vmax} - figure, axes = show_2d( - image_grid, - title=label_grid, - cmap=cmap, - norm=norm, - scalebar={"sampling": panels["pixel_size"], "units": panels["pixel_unit"]}, - axsize=(max(float(size), 200.0) / 100.0,) * 2, - **kwargs, - ) + # Build without registering an inline-backend draw. Otherwise Jupyter + # flushes the open figure after also rendering the returned Figure value. + with plt.ioff(): + figure, axes = show_2d( + image_grid, + title=label_grid, + cmap=cmap, + norm=norm, + scalebar={ + "sampling": panels["pixel_size"], + "units": panels["pixel_unit"], + }, + axsize=(max(float(size), 200.0) / 100.0,) * 2, + **kwargs, + ) axes_array = np.asarray(axes, dtype=object).reshape(len(image_grid), ncols) interpolation = "bilinear" if smooth else "nearest" for axis, image in zip(axes_array.flat, images, strict=True): @@ -912,6 +918,9 @@ def static_comparison(correction, zoom: float, size: int, **kwargs): center_row + visible_height / 2, center_row - visible_height / 2, ) + # Remove the manager while preserving the returned Figure for rich display + # and savefig. A bare ``drift.show(mode="static")`` then renders once. + plt.close(figure) return figure diff --git a/tests/imaging/test_drift.py b/tests/imaging/test_drift.py index 6f74a428..f34ab4d5 100644 --- a/tests/imaging/test_drift.py +++ b/tests/imaging/test_drift.py @@ -2,8 +2,10 @@ import inspect +import matplotlib.pyplot as plt import numpy as np import pytest +from matplotlib.figure import Figure from scipy.ndimage import gaussian_filter, shift from quantem.core.datastructures.dataset2d import Dataset2d @@ -74,6 +76,25 @@ def test_from_images_requires_angles_for_bare_arrays(): DriftCorrection.from_images(image, image.copy(), device="cpu") +def test_static_show_returns_one_closed_matplotlib_figure(): + """A bare static show call does not also queue an inline duplicate.""" + scan_0, scan_90 = _orthogonal_pair() + drift = DriftCorrection.from_images(scan_0, scan_90, device="cpu") + drift.preprocess( + padding_fraction=0.25, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + + figure = drift.show(mode="static", cmap="gray") + + assert isinstance(figure, Figure) + assert len(figure.axes) == 6 + assert figure.number not in plt.get_fignums() + + def test_metadata_driven_affine_workflow_returns_calibrated_dataset(): """The normal image workflow is short, finite, and calibration preserving.""" scan_0, scan_90 = _orthogonal_pair() diff --git a/tests/imaging/test_drift3d.py b/tests/imaging/test_drift3d.py index bdced6e9..186b7467 100644 --- a/tests/imaging/test_drift3d.py +++ b/tests/imaging/test_drift3d.py @@ -292,3 +292,80 @@ def test_read_emd_eds_preserves_native_energy_axis(tmp_path, monkeypatch): np.testing.assert_array_equal(spectrum.origin, [1.0, 2.0, 0.35]) np.testing.assert_array_equal(spectrum.sampling, [0.2, 0.2, 0.01]) assert spectrum.units == ["nm", "nm", "keV"] + + +def test_read_emd_eds_extracts_requested_windows_without_dense_cube( + tmp_path, monkeypatch +): + """Requested EDS windows are counted directly from each sparse stream.""" + path = tmp_path / "spectrum_windows.emd" + _write_metadata_emd( + path, + rotation_degrees=0.0, + stage_position=(0.0, 0.0), + spectrum_image=True, + timestamp=1, + scan_shape=(2, 3), + ) + detector_metadata = { + "BinaryResult": {"Detector": "SuperX-1"}, + "Detectors": { + "Detector-0": { + "DetectorName": "SuperX-1", + "Dispersion": "100", + "OffsetEnergy": "0", + } + }, + } + acquisition_settings = { + "bincount": "8", + "StreamEncoding": "uint16", + "RasterScanDefinition": {"Width": "3", "Height": "2"}, + } + # Six pixels. Values other than 65535 are one X-ray count in that + # energy-channel index; 65535 advances to the next scan pixel. + stream = np.array( + [1, 2, 65535, 2, 65535, 4, 65535, 1, 3, 65535, 2, 2, 65535, 65535], + dtype=np.uint16, + ) + with h5py.File(path, "a") as handle: + group = handle.create_group("Data/SpectrumStream/stream-0") + group.create_dataset( + "AcquisitionSettings", + data=np.array([json.dumps(acquisition_settings).encode()]), + ) + encoded = np.frombuffer(json.dumps(detector_metadata).encode(), dtype=np.uint8) + group.create_dataset("Metadata", data=encoded) + group.create_dataset("Data", data=stream[:, None]) + + streams = [ + { + "data": np.ones((2, 3), dtype=np.float32), + "metadata": {"General": {"title": "HAADF"}}, + "axes": [ + {"index_in_array": 0, "scale": 1.0, "offset": 0.0}, + {"index_in_array": 1, "scale": 1.0, "offset": 0.0}, + ], + } + ] + monkeypatch.setattr("rsciio.emd.file_reader", lambda *args, **kwargs: streams) + + acquisition = read_emd_eds( + path, + energy_windows={"low": (0.1, 0.2), "high": (0.3, 0.4)}, + verbose=False, + ) + + np.testing.assert_array_equal( + acquisition["window_maps"]["low"], + [[2, 1, 0], [1, 2, 0]], + ) + np.testing.assert_array_equal( + acquisition["window_maps"]["high"], + [[0, 0, 1], [1, 0, 0]], + ) + np.testing.assert_allclose( + acquisition["energy_axis_keV"], + np.arange(8, dtype=np.float32) * np.float32(0.1), + ) + assert acquisition["spectrum"] is None