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 3637acdc..2286d509 100644 --- a/src/quantem/imaging/__init__.py +++ b/src/quantem/imaging/__init__.py @@ -1,3 +1,15 @@ -from quantem.imaging.drift import DriftCorrection as DriftCorrection +"""Imaging tools for scientific image analysis.""" + +from quantem.imaging.drift import ( + CorrectionResult as CorrectionResult, + DriftCorrection as DriftCorrection, + StripPass as StripPass, + 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, +) 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.py b/src/quantem/imaging/drift.py deleted file mode 100644 index 415789d8..00000000 --- a/src/quantem/imaging/drift.py +++ /dev/null @@ -1,1832 +0,0 @@ -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.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 diff --git a/src/quantem/imaging/drift/__init__.py b/src/quantem/imaging/drift/__init__.py new file mode 100644 index 00000000..585a04ca --- /dev/null +++ b/src/quantem/imaging/drift/__init__.py @@ -0,0 +1,59 @@ +"""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, + read_emd_eds as read_emd_eds, + read_emd_metadata as read_emd_metadata, +) + +__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 new file mode 100644 index 00000000..b2c00ba1 --- /dev/null +++ b/src/quantem/imaging/drift/apply.py @@ -0,0 +1,944 @@ +"""Apply a solved drift field and return corrected scientific data.""" + +from copy import deepcopy +from typing import Literal + +import numpy as np +import torch +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.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 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 + + +def padding_offset( + canvas_shape: tuple[int, int], + scan_shape: tuple[int, int], + *, + 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 + + +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 + + +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)) + ) + 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, + ) + 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_frame: Literal["auto", "input", "canvas"] = "auto", + strip_padding: bool = False, + smoothing_sigma: float | None = 0.5, + stage: str | None = None, + merge: bool = True, + verbose: bool = True, +): + """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 + Sampling multiplier for the corrected image. + 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 + 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 + ------- + Dataset2d, Dataset3d, or list of Dataset2d + Corrected data with the source calibration and metadata. + + Examples + -------- + >>> corrected = drift.corrected() + >>> scans = drift.corrected(merge=False) + """ + 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 resolved_output_frame != "input" + 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( + "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 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) + 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), + 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(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_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) + 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)) + + output_shape = ( + tuple(int(value) for value in self.imgs[0].shape[:2]) + if resolved_output_frame == "input" + 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: + 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( + corrected_array, + name="drift corrected image", + 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": 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), + } + ) + return image_corr + + +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. + """ + # 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}" + ) + + ndim = len(original_shape) + if ndim < 3: + raise ValueError( + f"ds_4d must be at least 3D, got shape {original_shape}" + ) + + scan_h, scan_w = original_shape[0], original_shape[1] + n_channels = 1 + for d in range(2, ndim): + n_channels *= original_shape[d] + + 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 + ) + # 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) + ) + + # 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) + ) + + # 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) + + # ── Allocate output ── + if use_external_output: + out_flat = output.reshape(scan_h, scan_w, n_channels) + else: + internal_output = torch.empty( + scan_h, scan_w, n_channels, dtype=out_dt, device=target, + ) + + # ── 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) + + # 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 apply_correction( + self, + 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. + + 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 + ---------- + 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 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), + ) + + 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, + ) + + 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: + 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 crop(self, image: NDArray, *, shape: str = "square") -> NDArray: + """Crop an image to the field measured by every corrected scan. + + Parameters + ---------- + 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 + ------- + numpy.ndarray + Cropped data with trailing channel or detector axes unchanged. + + Examples + -------- + >>> cropped = drift.crop(corrected.array, shape="rectangle") + """ + 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. + + 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, + ) + 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 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 new file mode 100644 index 00000000..d0e9e8b6 --- /dev/null +++ b/src/quantem/imaging/drift/core/__init__.py @@ -0,0 +1,6 @@ +"""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 new file mode 100644 index 00000000..94285517 --- /dev/null +++ b/src/quantem/imaging/drift/core/affine.py @@ -0,0 +1,1740 @@ +"""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 + +import quantem.imaging.drift.apply as drift_apply +import quantem.imaging.drift.plot as drift_plot +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 ( + backward_warp_grid_search, + cross_corr_batch, + fixed_overlap_ncc, + translate_align_pair_batch, + warp_and_translate, +) + + +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 = [ + preprocessing.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, + *, + 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, + 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, + show_knot_plot: bool = False, + show_report: bool = False, + verbose: bool = True, + downsample: int | str = "auto", + chunk_size: int | None = None, +): + """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. + + 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 + ---------- + 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, 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. ``"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 "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, 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, 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 + Self, for method chaining. + + Examples + -------- + >>> 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]) + """ + # 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( + "Set both max_drift_rate and num_rates for an explicit grid, " + "or leave both unset for automatic affine alignment." + ) + if not automatic and downsample != "auto": + raise ValueError( + "downsample only applies when max_drift_rate and " + "num_rates are left unset for automatic affine alignment." + ) + 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, + region=region, + smoothing_sigma=region_smoothing_sigma, + max_drift_rate=max_drift_rate, + num_rates=num_rates, + refine=refine, + max_image_shift=max_image_shift, + fixed_scans=fixed_scans, + verbose=verbose, + downsample=downsample, + chunk_size=chunk_size, + ) + 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 = preprocessing.automatic_alignment_normalization( + self.imgs + ) + padding = preprocessing.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, + 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_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() + 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 _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_indices=fixed_set, + progress_desc=f"Affine {label.lower()}" if verbose else None, + ) + _apply_affine_rate(self, candidates[best_idx], fixed_set) + if verbose: + _print_top_candidates(label, candidates, costs) + warped_t = warp_and_translate( + self, + max_image_shift, upsample_factor, fixed_indices=fixed_set + ) + 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, coarse_margin = _search_and_apply(drift_vectors, "Coarse search") + if refine: + drift_fine = drift_vectors / (num_tests - 1) + 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.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"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}%)" + ) + 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 _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), + ) + + 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, + ) + + 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 + 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 + + +@torch.inference_mode() +def automatic_affine_search( + self, + *, + 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, + upsample_factor: int, + chunk_size: int | None, + pyramid_downsample: int | str, +): + """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 = preprocessing.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: + 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), + ) + 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) + ) + + # ------------------------------------------------------------------------- + # 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 + 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, + ) + 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]) + + # 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 + ) + ) + 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": "automatic_pyramid", + "downsample_factor": factor, + "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, + } + if verbose: + print( + "correct_affine: automatic pyramid " + f"{factor}x, {native_evaluations} native candidates, " + f"{elapsed:.2f} s" + ) + 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 grid_search_batch( + self, + drift_vectors, + upsample_factor, + max_image_shift, + chunk_size=None, + fixed_indices=None, + progress_desc=None, + fixed_overlap_check=False, +): + """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. + + 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) + 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. + 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 + ------- + 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 + 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(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 + 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 = 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 + 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) + 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 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. 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) + bytes_per_element = torch.finfo(dtype).bits // 8 + per_candidate_bytes = canvas_shape[0] * canvas_shape[1] * bytes_per_element * 32 + 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: + 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 new file mode 100644 index 00000000..3c39e592 --- /dev/null +++ b/src/quantem/imaging/drift/core/knots.py @@ -0,0 +1,691 @@ +"""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 + + +def interpolator(correction, image_index: int, knots: torch.Tensor | None = None): + """Map one scan's knot positions into its padded correction canvas. + + 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, + ) + + +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 + + +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``. + + 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. + + 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. + + 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)}." + ) + + 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], +) -> 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 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)`` + ``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). + 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)``. + """ + 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 _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, + 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. + """ + 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) + + +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: 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 + 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 new file mode 100644 index 00000000..4dbf2997 --- /dev/null +++ b/src/quantem/imaging/drift/core/nonrigid.py @@ -0,0 +1,777 @@ +"""Per-scanline non-rigid drift optimization.""" +import numpy as np +import torch +import torch.nn.functional as F +from tqdm import tqdm + +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 _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. + """ + 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() + + +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``. + + 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. + """ + 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 + + +@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( + ref_batch, target_batch, knots_batch, + loss_fn, loss_args, + optimizer, steps, grad_mask=None, +): + """Run ``steps`` of Adam on a batched knot tensor against ``loss_fn``.""" + ref_t = ref_batch[:, None] + for _ in range(steps): + optimizer.zero_grad() + 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() + + +def _optimize_knots_lbfgs( + 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 = 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( + 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`` + 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. + """ + # 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 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 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: + 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 new file mode 100644 index 00000000..14a5dae6 --- /dev/null +++ b/src/quantem/imaging/drift/core/strip.py @@ -0,0 +1,839 @@ +"""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 new file mode 100644 index 00000000..2cd24d72 --- /dev/null +++ b/src/quantem/imaging/drift/core/warping.py @@ -0,0 +1,1111 @@ +"""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( + 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) + >>> cost.shape + torch.Size([5]) + """ + _, 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() + image_shifts = _translation_from_cross_correlation( + cross_corr_fft, upsample_factor, max_shift_mask + ) + 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 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, +) -> torch.Tensor: + """Pairwise translation alignment of warped images via cross-correlation. + + 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. + + 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() + 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( + -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 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 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, + 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: 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 + 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: 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. + 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: 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 + 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 new file mode 100644 index 00000000..aaeecb4f --- /dev/null +++ b/src/quantem/imaging/drift/correction.py @@ -0,0 +1,473 @@ +"""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``. +Manual rigid registration remains available through ``align_translation``. +""" + +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.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 +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 +from quantem.core.datastructures.dataset3d import Dataset3d +from quantem.core.io.serialize import AutoSerialize + + +class DriftCorrection(AutoSerialize): + """GPU-accelerated multi-angle scan drift correction. + + 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`. 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). + """ + + def __init__( + self, + *datasets: Dataset2d | Dataset3d | NDArray, + scan_direction_degrees: list[float] | NDArray | float | None = None, + alignment_image: NDArray | None = None, + device: str | None = None, + ): + """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 + + if not datasets and alignment_image is None: + return + preprocessing.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_images( + cls, + *images, + scan_direction_degrees: list[float] | NDArray | tuple[float, ...] | None = None, + device: str | None = None, + ) -> Self: + """Create drift correction from already-loaded 2-D scans. + + 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_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] = preprocessing.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, + *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: + """Build a 0°/90° 4D-STEM *collection* correction (two drifted scans). + + 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 + ---------- + *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 + + Examples + -------- + >>> 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, + ... ) + >>> dc.correct_affine(show_combined=False, verbose=False) + >>> result = dc.corrected_4dstem(merge=True, verbose=True) + """ + result = cls( + *datasets, + scan_direction_degrees=scan_direction_degrees, + device=device, + ) + 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_reference( + cls, + reference_image, + drifted_dataset, + *, + alignment_image: NDArray | None = None, + scan_direction_degrees: list[float] | NDArray | float = 0.0, + device: str | None = None, + ) -> Self: + """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 + + 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) + """ + 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 = 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), + target_sampling=getattr(drifted_dataset, "sampling", None), + ) + 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 = [ + preprocessing.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, + ) + reference_dc.correct_affine( + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + 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), + ) + 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, + ) + # 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 + + drift_field = fourdstem.drift_field + probe_positions = fourdstem.probe_positions + + preprocess = preprocessing.preprocess + align_translation = warping.align_translation + 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 + + 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): + """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, + compression_level=compression_level, + ) + + 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 + # 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 new file mode 100644 index 00000000..0fbd7ea2 --- /dev/null +++ b/src/quantem/imaging/drift/diagnostics.py @@ -0,0 +1,474 @@ +"""Scientific diagnostics for affine and non-rigid drift correction.""" + +import copy +import time + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch +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" + ) + 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([]) + + 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( + "diagnose_nonrigid() requires an affine correction. " + "Run correct_affine() first." + ) + counts = tuple(dict.fromkeys(int(value) for value in num_knots)) + if not counts or min(counts) < 1: + raise ValueError( + f"num_knots must contain positive integers, got {num_knots!r}." + ) + + 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, + } + ) + + common_mask = np.logical_and.reduce( + [np.asarray(candidate["mask"], dtype=bool) for candidate in candidates] + ) + 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, + ) + 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"], + } + ) + + 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, + ) + return figure, pd.DataFrame.from_records(records) diff --git a/src/quantem/imaging/drift/fourdstem.py b/src/quantem/imaging/drift/fourdstem.py new file mode 100644 index 00000000..92a5e5f3 --- /dev/null +++ b/src/quantem/imaging/drift/fourdstem.py @@ -0,0 +1,834 @@ +"""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_correction_to_dataset, + crop_slices, + padding_offset, +) +from quantem.imaging.drift.core import knots as drift_knots + + +@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) + """ + 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]) + 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, 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) + 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) + 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 + + +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 drift_knots.interpolator(self, idx).drift_raw( + self._initial_knots[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.imgs) + knots = ( + self.knots[index] + if corrected + else self._initial_knots[index] + ) + 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.imgs[0].shape[:2], + ) + positions -= torch.tensor( + [pad_row, pad_column], + device=positions.device, + dtype=positions.dtype, + ) + result = to_numpy(positions, dtype=np.float32) + if plot: + self.plot_probe_positions( + image_index=index, + strip_padding=strip_padding, + stride=stride, + ) + 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.imgs) != 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.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.imgs[0].shape}, {self.imgs[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): + image_t = torch.as_tensor( + image, + 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].to(torch.float32) + valid[1].to( + torch.float32 + ) + 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( + torch.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): + 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=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 = torch.rot90(corrected, k=rot_k, dims=(0, 1)) + 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(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.imgs] + 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_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, + *, + 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_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 inputs_on_device: + self._datasets[0] = None + torch.cuda.empty_cache() + 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 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), + ).copy() + + 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): + 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(merge_dtype) + + corrected_4dstem_1[r0:r1].to(merge_dtype) + ) * 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.imgs[0].array) + alignment_vdf_1 = np.asarray(self.imgs[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/src/quantem/imaging/drift/io.py b/src/quantem/imaging/drift/io.py new file mode 100644 index 00000000..f0f7f480 --- /dev/null +++ b/src/quantem/imaging/drift/io.py @@ -0,0 +1,841 @@ +"""Read Velox EMD images, spectrum images, and scan metadata for drift correction.""" + +import json +import math +from pathlib import Path + +import h5py +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, +) -> 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 Velox's repeated-byte metadata encoding for image stacks.""" + 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 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: + 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")) + except json.JSONDecodeError: + continue + return {} + + +def read_emd_metadata(file_path: str | Path) -> dict: + """Read normalized Velox acquisition metadata without loading image data. + + Parameters + ---------- + file_path : str or Path + Velox EMD file. + + Returns + ------- + 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_0.emd") + >>> metadata["scan_rotation_deg"] + 0.0 + """ + 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_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_width = scan_size.get("width") + scan_height = scan_size.get("height", scan_width) + scan_shape = ( + (int(scan_height), int(scan_width)) + if scan_height and scan_width + 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_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, + "fov_m": ( + float(pixel_size_m) * float(scan_width) + if pixel_size_m and scan_width + else None + ), + "acquisition_timestamp": ( + int(timestamp) if timestamp and str(timestamp).isdigit() else None + ), + "acquisition_context": acquisition_context, + "original_metadata": metadata, + } + + +def _image_dataset(stream, path, metadata): + from quantem.core.datastructures.dataset2d import Dataset2d + + 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 = path + image.metadata.update(metadata) + return image + + +def _spectrum_dataset(stream, path, metadata): + """Build a calibrated ``(scan_row, scan_col, energy)`` Dataset3d.""" + from quantem.core.datastructures.dataset3d import Dataset3d + + array = np.asarray(stream["data"]) + axes = [ + axis + for _, axis in sorted( + ( + (int(axis.get("index_in_array", fallback)), axis) + for fallback, axis in enumerate(stream.get("axes", [])[:3]) + ), + key=lambda item: item[0], + ) + ] + energy_axis = next( + ( + index + for index, axis in enumerate(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) + 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), + 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(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, + 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. + + 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, 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. + + Returns + ------- + dict[str, object] + 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") + >>> acquisition["haadf"].shape + (2048, 2048) + """ + from rsciio.emd import file_reader + + 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 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" + ), + 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 + } + 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) + ) + 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 = window_axis + 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)} 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, + "window_maps": window_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: + 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["fov_m"], second["fov_m"]) > 0.02: + return False + if _relative_difference(first["magnification"], second["magnification"]) > 0.02: + return False + 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["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 + 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/plot.py b/src/quantem/imaging/drift/plot.py new file mode 100644 index 00000000..70447eae --- /dev/null +++ b/src/quantem/imaging/drift/plot.py @@ -0,0 +1,1033 @@ +"""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.core import knots as drift_knots +from quantem.imaging.drift.core.warping import ensure_warped_images + + +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: + 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: tuple[float, float] = (7, 4.2), + log_scale: bool = True, + **kwargs, +) -> 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() + """ + 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." + ) + if affine_error.size: + nonrigid_error = np.concatenate((affine_error[-1:], nonrigid_error)) + + 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 + + +# --------------------------------------------------------------------------- +# Building blocks - used internally by the public API functions above +# --------------------------------------------------------------------------- + + +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 _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 _stage_description(stage: str | None) -> str: + """Return a concise, natural-language description of a correction stage.""" + + return { + None: "current correction", + "initial": "raw", + "raw": "raw", + "translation": "after translation alignment", + "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, + *, + 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, + ) + merged_title = kwargs.pop( + "title", + f"Combined: {stage_display}" + (" (RGB)" if rgb else ""), + ) + 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 = 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) + ) + 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_knots( + self, + *, + figsize: tuple[int, int] | None = None, + stage: str | None = None, +) -> 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 = (7 * num_images, 11) + fig, axes = plt.subplots( + 2, num_images, figsize=figsize, + gridspec_kw={"height_ratios": [2, 1]}, + ) + 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 + ) + strip_np = ( + self._knots_after_strip[img_idx].cpu().numpy() + if hasattr(self, "_knots_after_strip") + else None + ) + 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()) + ) + 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} + + # 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): + 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, + ) + 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, + ) + # 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 + + +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, + ) + viewer.frame_idx = 2 + return viewer diff --git a/src/quantem/imaging/drift/preprocess.py b/src/quantem/imaging/drift/preprocess.py new file mode 100644 index 00000000..63773395 --- /dev/null +++ b/src/quantem/imaging/drift/preprocess.py @@ -0,0 +1,741 @@ +"""Preprocess scientific inputs 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.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 a computational downsampling factor.""" + factor = int(value) + 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 integer factor.""" + image = np.asarray(array) + 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, cols = image.shape + if rows % factor or cols % factor: + raise ValueError( + f"downsample={factor} requires image dimensions divisible by " + f"{factor}, got {image.shape}." + ) + downsampled = image.reshape( + rows // factor, + factor, + cols // 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 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 safe average-downsampling factor, capped at eight.""" + height, width = shape + for factor in (8, 4, 2): + if height % factor == 0 and width % factor == 0: + return factor + return 1 + + +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): + 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 " + f"to be divisible by {factor}, got {shape}. Choose a divisor " + "such as 1, 2, 4, or 8." + ) + return factor + + +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 images and update their coordinate metadata.""" + original_records = [] + downsampled_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_units = list(image.units) + original_shape = tuple(int(value) for value in original_array.shape[:2]) + downsampled = average_downsample_2d(original_array, factor) + sampling = scale_coordinate_metadata(original_sampling, factor) + origin_offset = (factor - 1) * original_sampling[:2] / 2.0 + origin = scale_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[:2]), + "sampling_is_downsampled": True, + } + ) + image_dataset = Dataset2d.from_array( + downsampled, + name=f"{image.name} ({factor}x downsample)", + origin=origin, + sampling=sampling, + units=original_units, + signal_units=image.signal_units, + ) + 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": original_units, + } + ) + 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.imgs[0].shape[:2]), + } + + +def preprocess( + self, + *, + 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, + show_knot_plot: bool = False, +): + """Build the shared scanline canvas before drift correction. + + Affine correction calls this automatically. Use it directly only when a + publication requires a fixed padding, normalization, or knot layout. + + Parameters + ---------- + 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 + ------- + object + The same correction object for method chaining. + + Examples + -------- + >>> drift.preprocess(padding_fraction=0.25, show_combined=True) + """ + downsample = validate_downsample(downsample) + if downsample > 1: + if getattr(self, "_datasets", None) is not None: + 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 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." + ) + else: + self.downsample = 1 + self.downsample_method = "none" + self.downsample_metadata = { + "factor": 1, + "method": "none", + "original_images": [ + { + "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 image_index, image in enumerate(self.imgs) + ], + "downsampled_shape": list(self.imgs[0].shape[:2]), + } + + 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 + ) + self.scan_slow = np.stack( + [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.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 + ), + ) + 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 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]) + ] + device = self._device + dtype = self._dtype + 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[index], dtype=dtype, device=device) + for index in range(self.shape[0]) + ] + 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 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 new file mode 100644 index 00000000..e54923ef --- /dev/null +++ b/src/quantem/imaging/drift/report.py @@ -0,0 +1,170 @@ +"""Regional alignment reports for :class:`DriftCorrection`. + +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 Sequence + +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 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 report( + self, + *, + stages: Sequence[str] | None = None, +) -> 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) + + 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/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.py b/tests/imaging/test_drift.py index 3ac72a06..f34ab4d5 100644 --- a/tests/imaging/test_drift.py +++ b/tests/imaging/test_drift.py @@ -1,243 +1,201 @@ -""" -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 matplotlib.pyplot as plt import numpy as np import pytest -from scipy.ndimage import gaussian_filter +from matplotlib.figure import Figure +from scipy.ndimage import gaussian_filter, shift + from quantem.core.datastructures.dataset2d import Dataset2d -from quantem.imaging.drift import DriftCorrection - - -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 - ) +from quantem.imaging import DriftCorrection + - 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 +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"), ) + 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 + 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 + ) + for old_name in ( + "pad_fraction", + "kde_sigma", + "number_knots", + "step", + "num_tests", + "num_iterations", + ): + 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_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, + ) - return im0, im1, base_image + 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_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, +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.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, + drift.correct_affine( + max_drift_rate=0.02, + num_rates=5, + refine=False, + max_image_shift=8, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=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, + corrected = drift.corrected( + upsample_factor=1, + output_frame="input", + verbose=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!", + 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(upsample_factor=1, verbose=False) + canvas = drift.corrected( + 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", + ) -# 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, + returned = drift.align_translation( + max_image_shift=8, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=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, + initial = drift.corrected( + stage="initial", + merge=False, + output_frame="input", + verbose=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, + aligned = drift.corrected( + merge=False, + output_frame="input", + verbose=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, + interior = np.s_[8:-8, 8:-8] + initial_error = np.mean( + np.abs(initial[0].array[interior] - initial[1].array[interior]) ) - 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, + aligned_error = np.mean( + np.abs(aligned[0].array[interior] - aligned[1].array[interior]) ) - 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) + + 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 "" + 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 new file mode 100644 index 00000000..186b7467 --- /dev/null +++ b/tests/imaging/test_drift3d.py @@ -0,0 +1,371 @@ +"""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, + spectrum_image, + alignment_image=alignment_image, + scan_direction_degrees=0.0, + 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_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + corrected = drift.corrected(verbose=False) + + 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].cpu(), + initial_reference_knots.cpu(), + 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"] + + +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 diff --git a/tests/imaging/test_drift4d.py b/tests/imaging/test_drift4d.py new file mode 100644 index 00000000..f927e74d --- /dev/null +++ b/tests/imaging/test_drift4d.py @@ -0,0 +1,443 @@ +"""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( + padding_fraction=0.25, + 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=8, + chunk_size=1, + 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.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(): + """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( + num_knots=1, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=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( + num_knots=1, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=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.imgs[0].sampling, (0.2, 0.3)) + assert drift.imgs[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( + 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) + _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.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(): + """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( + num_knots=1, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=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( + 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) + _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, + ) diff --git a/tests/imaging/test_drift_downsampling.py b/tests/imaging/test_drift_downsampling.py new file mode 100644 index 00000000..e6c4e52b --- /dev/null +++ b/tests/imaging/test_drift_downsampling.py @@ -0,0 +1,157 @@ +"""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.preprocess 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_images( + *datasets, + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + downsample=2, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + + 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_images( + *_calibrated_pair(), + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + downsample=2, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + + corrected = drift.corrected( + upsample_factor=1, + verbose=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_images( + *_calibrated_pair(), + scan_direction_degrees=(0.0, 90.0), + device="cpu", + ).preprocess( + 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_combined=False, + show_scans=False, + show_knots=False, + verbose=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_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) + 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])), + ) + + drift = DriftCorrection.from_reference( + reference, + target, + scan_direction_degrees=0.0, + device="cpu", + ) + drift.correct_affine( + downsample=2, + show_combined=False, + show_scans=False, + show_knots=False, + verbose=False, + ) + + 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 a917ef17..4f0eceae 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,18 @@ 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.apply import fourier_crop_torch +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 @@ -191,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)), @@ -228,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" },