diff --git a/src/quantem/core/datastructures/polar4dstem.py b/src/quantem/core/datastructures/polar4dstem.py index 81ae5d03..903c4114 100644 --- a/src/quantem/core/datastructures/polar4dstem.py +++ b/src/quantem/core/datastructures/polar4dstem.py @@ -1,6 +1,7 @@ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Self import numpy as np +import torch from numpy.typing import NDArray from scipy.ndimage import map_coordinates @@ -11,17 +12,19 @@ class Polar4dstem(Dataset4d): - """4D-STEM dataset in polar coordinates (scan_y, scan_x, phi, r).""" + """4D-STEM dataset in polar coordinates (scan_row, scan_col, phi, r_pix).""" def __init__( self, - array: NDArray | Any, - name: str, - origin: NDArray | tuple | list | float | int, - sampling: NDArray | tuple | list | float | int, - units: list[str] | tuple | list, + array: NDArray | None = None, + tensor: torch.Tensor | None = None, + name: str = "", + origin: NDArray | tuple | list | float | int | None = None, + sampling: NDArray | tuple | list | float | int | None = None, + units: list[str] | tuple | list | None = None, signal_units: str = "arb. units", metadata: dict | None = None, + origin_array: NDArray | None = None, _token: object | None = None, ): if metadata is None: @@ -32,6 +35,7 @@ def __init__( "polar_radial_step", "polar_num_annular_bins", "polar_two_fold_rotation_symmetry", + "polar_ellipticity", "polar_origin_row", "polar_origin_col", "polar_ellipse_params", @@ -41,6 +45,7 @@ def __init__( metadata[k] = None super().__init__( array=array, + tensor=tensor, name=name, origin=origin, sampling=sampling, @@ -49,45 +54,30 @@ def __init__( metadata=metadata, _token=_token, ) + self.origin_array = origin_array @classmethod def from_array( cls, - array: NDArray | Any, + array: NDArray, name: str | None = None, origin: NDArray | tuple | list | float | int | None = None, sampling: NDArray | tuple | list | float | int | None = None, units: list[str] | tuple | list | None = None, signal_units: str = "arb. units", metadata: dict | None = None, - ) -> "Polar4dstem": - """ - Create a Polar4dstem from a 4D array shaped (scan_y, scan_x, phi, r). - - Parameters - ---------- - array : NDArray - Polar data, shape (scan_y, scan_x, n_phi, n_r). - name : str, optional - Dataset name, by default "Polar 4D-STEM dataset". - origin : NDArray | tuple | list | float | int, optional - Origin of each axis in calibrated units, by default zeros. - sampling : NDArray | tuple | list | float | int, optional - Sampling of each axis, by default ones. - units : list[str] | tuple | list, optional - Units of each axis, by default ["pixels", "pixels", "deg", "pixels"]. - signal_units : str, optional - Units of the array values, by default "arb. units". - metadata : dict, optional - Metadata. Missing "polar_*" keys are set to None. - - Returns - ------- - Polar4dstem - """ + ) -> Self: + if isinstance(array, torch.Tensor): + raise TypeError( + f"Polar4dstem.from_array requires a numpy array, got {type(array).__name__}. " + "Use Polar4dstem.from_tensor for torch tensors." + ) array = np.asarray(array) if array.ndim != 4: - raise ValueError("Polar4dstem.from_array expects a 4D array.") + raise ValueError( + f"Found array with shape: {array.shape}. " + "Polar4dstem.from_array expects a 4D array." + ) if origin is None: origin = np.zeros(4, dtype=float) if sampling is None: @@ -107,14 +97,52 @@ def from_array( _token=cls._token, ) + @classmethod + def from_tensor( + cls, + tensor: torch.Tensor, + name: str | None = None, + origin: NDArray | tuple | list | float | int | None = None, + sampling: NDArray | tuple | list | float | int | None = None, + units: list[str] | tuple | list | None = None, + signal_units: str = "arb. units", + metadata: dict | None = None, + ) -> Self: + """Create a Polar4dstem from a torch tensor.""" + if not isinstance(tensor, torch.Tensor): + raise TypeError( + f"Polar4dstem.from_tensor requires torch.Tensor, got {type(tensor).__name__}." + ) + if tensor.ndim != 4: + raise ValueError( + f"Found tensor with shape: {tuple(tensor.shape)}. " + "Polar4dstem.from_tensor expects a 4D tensor." + ) + if origin is None: + origin = np.zeros(4, dtype=float) + if sampling is None: + sampling = np.ones(4, dtype=float) + if units is None: + units = ["pixels", "pixels", "deg", "pixels"] + if metadata is None: + metadata = {} + return cls( + tensor=tensor, + name=name if name is not None else "Polar 4D-STEM dataset", + origin=origin, + sampling=sampling, + units=units, + signal_units=signal_units, + metadata=metadata, + _token=cls._token, + ) + @property def n_phi(self) -> int: - """Number of azimuthal bins (axis 2).""" return int(self.shape[2]) @property def n_r(self) -> int: - """Number of radial bins (axis 3).""" return int(self.shape[3]) diff --git a/src/quantem/core/utils/filter.py b/src/quantem/core/utils/filter.py index 3f9d13b7..a11d5e70 100644 --- a/src/quantem/core/utils/filter.py +++ b/src/quantem/core/utils/filter.py @@ -135,6 +135,17 @@ def gaussian_filter_2d( return img.squeeze_(0).squeeze_(0) # Make 2D again +def gaussian_filter_1d( + arr: torch.Tensor, kernel_1d: torch.Tensor +) -> torch.Tensor: # Replicate-padded torch alternative to ``scipy.ndimage.gaussian_filter1d`` + padding = len(kernel_1d) // 2 # Ensure that signal size does not change + arr = arr.unsqueeze(0).unsqueeze_(0) # Make copy, make 3D for ``conv1d()`` + # Replicate edge values so the output has no zero-padded ringing at boundaries + arr = torch.nn.functional.pad(arr, (padding, padding), mode="replicate") + arr = torch.nn.functional.conv1d(arr, weight=kernel_1d.view(1, 1, -1)) + return arr.squeeze_(0).squeeze_(0) # Make 1D again + + def gaussian_filter_2d_stack(stack: torch.Tensor, kernel_1d: torch.Tensor) -> torch.Tensor: """ Apply 2D Gaussian blur to each slice stack[:, i, :] in a vectorized way. diff --git a/src/quantem/diffraction/__init__.py b/src/quantem/diffraction/__init__.py index f167408e..259181aa 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -1,3 +1,4 @@ +from quantem.diffraction.polar import PairDistributionFunction as PairDistributionFunction from quantem.diffraction.bragg_vectors import BraggVectors as BraggVectors from quantem.diffraction.crystal import Crystal as Crystal from quantem.diffraction.crystal_map import CrystalMap as CrystalMap diff --git a/src/quantem/diffraction/ellipticity.py b/src/quantem/diffraction/ellipticity.py new file mode 100644 index 00000000..a6080c63 --- /dev/null +++ b/src/quantem/diffraction/ellipticity.py @@ -0,0 +1,251 @@ +import math +import warnings + +import numpy as np +import torch +from numpy.typing import NDArray + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem + + +def fit_elliptical_distortion( + data: Dataset4dstem | Dataset2d, + center: tuple[float, float], + fit_radii: tuple[float, float], + p0: torch.Tensor | None = None, + mask: NDArray | torch.Tensor | None = None, + device: str = "cpu", + max_iter: int = 500, + write_metadata: bool = True, +) -> dict: + """Fit the elliptical distortion of a diffraction pattern from its amorphous ring. + + The mean pattern (averaged over scan positions for 4D data) is fit, inside the + annulus ``fit_radii`` around ``center``, with a central-beam Gaussian, a ring with + different widths inside and outside its radius, and a constant background. + + Parameters + ---------- + data : Dataset4dstem or Dataset2d + 4D-STEM data (averaged over scan positions) or a single pattern. + center : (row, col) + Approximate pattern center in pixels; the fit refines it. + fit_radii : (inner, outer) + Radii in pixels of the annulus that is fit. It should contain the ring and + exclude the central beam. + p0 : Tensor or None + Initial model parameters (see :func:`_amorphous_ring_model`). Estimated from + the data if None. + mask : array or None + Boolean mask of pixels to ignore (True = ignore). + device : str + Torch device for the fit. + max_iter : int + Maximum number of iterations. + write_metadata : bool + If True, store the fitted ellipse in ``data.metadata["ellipticity"]``. + + Returns + ------- + dict + ``"ellipse_params"``: (a, b, theta_deg), the semiaxes in pixels (a >= b) and + the tilt of the major axis from the column axis in degrees, in [-90, 90), as + used by :func:`polar_transform`. ``"center"``: refined (row, col). + ``"fit_params"``: the 11 fitted model parameters. ``"cost"``: sum of squared + residuals. + """ + dp = ( + torch.from_numpy(np.ascontiguousarray(data.array)) + if data.array is not None + else data.tensor + ) + if dp.ndim == 4: + dp = dp.mean(dim=(0, 1), dtype=torch.float32) + elif dp.ndim == 2: + dp = dp.float() + else: + raise ValueError( + f"Got data with shape {tuple(data.shape)}. Expected a 2D or 4D-STEM dataset." + ) + dp = dp.to(device) + + # pixel offsets from the nominal center, restricted to the fit annulus + row_c, col_c = map(float, center) + r_in, r_out = map(float, fit_radii) + d_row, d_col = torch.meshgrid( + torch.arange(dp.shape[0], device=device) - row_c, + torch.arange(dp.shape[1], device=device) - col_c, + indexing="ij", + ) + r = torch.sqrt(d_col**2 + d_row**2) + in_fit = (r > r_in) & (r < r_out) & dp.isfinite() + if mask is not None: + in_fit &= ~torch.as_tensor(mask, dtype=torch.bool, device=device) + d_row, d_col, r, values = d_row[in_fit], d_col[in_fit], r[in_fit], dp[in_fit] + if values.numel() == 0: + raise ValueError("No pixels in the fitting annulus. Check center and fit_radii.") + + if p0 is None: + # ring radius = peak of the annulus intensity summed in radial bins + n_bins = max(int(r_out - r_in), 10) + bin_width = (r_out - r_in) / n_bins + bins = ((r - r_in) / bin_width).long().clamp(0, n_bins - 1) + radial_sum = torch.zeros(n_bins, device=device).scatter_add_(0, bins, values) + R = r_in + (radial_sum.argmax().item() + 0.5) * bin_width + near_ring = values[(r - R).abs() < 3.0] + if near_ring.numel() == 0: + raise ValueError( + f"No annulus pixels within 3 px of the estimated ring radius {R:.1f}. " + "Annulus is too narrow or fit_radii is misaligned with the ring." + ) + background = max(values.min().item(), 0.0) + I1 = max(near_ring.mean().item() - background, 0.0) + I0 = max(dp[int(row_c), int(col_c)].item() - background, 0.0) + # start from a circle of radius R centered on the nominal center + p0 = [I0, I1, R * 0.3, R * 0.15, R * 0.2, background, 0.0, 0.0] + p0 += convert_ellipse_params_r(R, R, 0.0) + p0 = torch.as_tensor(p0, dtype=torch.float32, device=device) + + # The fit varies x = (I0, I1, sigma0, sigma1, sigma2, background, row0, col0, R, e1, e2), + # describing the ring ellipse by its mean radius R and ellipticity (e1, e2) (see + # _conic_from_shape), which keeps the fit well conditioned. Bounds: amplitudes and + # background >= 0, R inside fit_radii, |e1| and |e2| <= 0.5. + a0, b0, theta0 = convert_ellipse_params(*p0[8:].tolist()) + e0 = (a0 - b0) / (a0 + b0) + shape0 = [(a0 + b0) / 2, e0 * math.cos(2 * theta0), e0 * math.sin(2 * theta0)] + x = torch.cat([p0[:8], p0.new_tensor(shape0)]) + lower = torch.full_like(x, -math.inf) + upper = torch.full_like(x, math.inf) + lower[[0, 1, 5]] = 0.0 + lower[8], upper[8] = r_in, r_out + lower[9:], upper[9:] = -0.5, 0.5 + + def model_params(x): + # the widths enter the model squared, so their sign does not matter + return torch.cat([x[:2], x[2:5].abs(), x[5:8], _conic_from_shape(*x[8:])]) + + def residuals(x): + return _amorphous_ring_model(model_params(x), d_row, d_col) - values + + x = _least_squares(residuals, x.clamp(lower, upper), lower, upper, max_iter) + if x[8] <= r_in or x[8] >= r_out: + warnings.warn( + f"Fitted ring radius {x[8].item():.1f} px is at the edge of fit_radii {fit_radii}. " + "Choose fit_radii around the ring." + ) + params = model_params(x) + a, b, theta = convert_ellipse_params(*params[8:].tolist()) + ellipse = (a, b, (math.degrees(theta) + 90.0) % 180.0 - 90.0) + if write_metadata: + data.metadata["ellipticity"] = ellipse + return { + "ellipse_params": ellipse, + "center": (row_c + params[6].item(), col_c + params[7].item()), + "fit_params": params, + "cost": (residuals(x) ** 2).sum().item(), + } + + +def _least_squares(residuals, x, lower, upper, max_iter): + """Minimise sum(residuals(x) ** 2) with lower <= x <= upper (Levenberg-Marquardt). + + Each step is a Gauss-Newton step damped by lam; lam is lowered after a step that lowers + the cost and raised until one does. The step is solved in units where every column of + the Jacobian has the same size (the largest seen so far), so that parameters of very + different scales are damped alike and the float32 solve stays accurate. Parameters on + a bound that the gradient pushes outward are not stepped. + """ + res = residuals(x) + cost = res @ res + lam = 1e-3 + col_norm = torch.full_like(x, 1e-12) + for _ in range(max_iter): + # (torch.func.jacfwd would turn float32 into float64 here, which fails on MPS) + J = torch.autograd.functional.jacobian( + residuals, x, vectorize=True, strategy="forward-mode" + ) + col_norm = torch.maximum(col_norm, J.norm(dim=0)) + J = J / col_norm + JtJ, Jtr = J.T @ J, J.T @ res + free = ~(((x <= lower) & (Jtr > 0)) | ((x >= upper) & (Jtr < 0))) + for _ in range(30): + A = JtJ + lam * torch.eye(len(x), device=x.device) + step = torch.zeros_like(x) + step[free] = torch.linalg.solve(A[free][:, free], Jtr[free]) + x_new = (x - step / col_norm).clamp(lower, upper) + res_new = residuals(x_new) + cost_new = res_new @ res_new + if cost_new < cost: + break + lam *= 3.0 + if not cost_new < cost: + return x # no step lowers the cost: converged + converged = cost - cost_new < 1e-6 * cost + x, res, cost, lam = x_new, res_new, cost_new, lam * 0.3 + if converged: + return x + warnings.warn(f"Ellipse fit did not converge in {max_iter} iterations.") + return x + + +def _amorphous_ring_model(params, offset_row, offset_col): + """Pattern model: central-beam Gaussian + amorphous ring + constant background. + + ``params`` are (I0, I1, sigma0, sigma1, sigma2, background, row0, col0, A, B, C): + beam and ring amplitudes, beam width, ring widths inside and outside its radius, + background, center offset from the nominal center in pixels, and the ring + ellipse ``A*x^2 + B*x*y + C*y^2 = 1`` (x, y = column, row offset from the + center). The beam shares the ring's elliptical radius. + """ + I0, I1, sigma0, sigma1, sigma2, background, row0, col0, A, B, C = params + x = offset_col - col0 + y = offset_row - row0 + # elliptical radius in pixels, equal to the mean semiaxis R on the ring + a, b = _ellipse_semiaxes(A, B, C) + R = (a + b) / 2.0 + rho2 = R**2 * (A * x**2 + B * x * y + C * y**2) + dr = torch.sqrt(torch.clamp(rho2, min=1e-12)) - R + sigma = torch.where(dr < 0, sigma1, sigma2) + beam = I0 * torch.exp(-rho2 / (2.0 * sigma0**2)) + ring = I1 * torch.exp(-(dr**2) / (2.0 * sigma**2)) + return beam + ring + background + + +def convert_ellipse_params(A, B, C): + """Semiaxes a >= b and tilt (radians, from the column axis) of the ellipse + ``A*x^2 + B*x*y + C*y^2 = 1``.""" + mean = (A + C) / 2 + half_gap = math.hypot(A - C, B) / 2 + a = 1 / math.sqrt(mean - half_gap) + b = 1 / math.sqrt(mean + half_gap) + theta = math.atan2(-B, C - A) / 2 + return a, b, theta + + +def convert_ellipse_params_r(a, b, theta): + """Coefficients (A, B, C) of the ellipse with semiaxes a, b and tilt theta (inverse + of :func:`convert_ellipse_params`).""" + mean = (1 / a**2 + 1 / b**2) / 2 + half_gap = (1 / b**2 - 1 / a**2) / 2 + A = mean - half_gap * math.cos(2 * theta) + B = -2 * half_gap * math.sin(2 * theta) + C = mean + half_gap * math.cos(2 * theta) + return A, B, C + + +def _ellipse_semiaxes(A, B, C): + """Differentiable semiaxes (a, b) of the ellipse ``A*x^2 + B*x*y + C*y^2 = 1``. + The clamps keep the square roots and their gradients finite for a circle.""" + mean = (A + C) / 2 + half_gap = torch.sqrt(torch.clamp((A - C) ** 2 + B**2, min=1e-24)) / 2 + a = torch.rsqrt(torch.clamp(mean - half_gap, min=1e-24)) + b = torch.rsqrt(torch.clamp(mean + half_gap, min=1e-24)) + return a, b + + +def _conic_from_shape(R, e1, e2): + """Conic coefficients (A, B, C) of the ellipse with mean radius R = (a + b) / 2 and + ellipticity e = (a - b) / (a + b) along (e1, e2) = e * (cos 2 theta, sin 2 theta).""" + e_sq = e1**2 + e2**2 + return torch.stack([1 + e_sq - 2 * e1, -4 * e2, 1 + e_sq + 2 * e1]) / (R * (1 - e_sq)) ** 2 diff --git a/src/quantem/diffraction/polar.py b/src/quantem/diffraction/polar.py new file mode 100644 index 00000000..c600f7a6 --- /dev/null +++ b/src/quantem/diffraction/polar.py @@ -0,0 +1,1423 @@ +from __future__ import annotations + +from collections.abc import Iterable +from typing import Literal, Self + +import matplotlib.pyplot as plt +import numpy as np +import torch +from numpy.typing import NDArray + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.datastructures.polar4dstem import Polar4dstem +from quantem.core.io.serialize import AutoSerialize +from quantem.core.utils.filter import gaussian_filter_1d, gaussian_kernel_1d +from quantem.core.utils.utils import to_numpy +from quantem.diffraction.polar_transform import ( + find_origin_angular_descent, + find_origin_angular_grid, + polar_transform, +) + +# TODO: beamstop mask support (mask diffraction-space pixels before azimuthal averaging) + + +class PairDistributionFunction(AutoSerialize): + """Compute the pair distribution function from 4D-STEM diffuse scattering + + The pair distribution function g(r) gives the probability of finding + pairs of atoms at separation r, and is the standard tool for + characterizing local atomic structure in amorphous materials where + Bragg diffraction is unavailable. This class implements the standard + extraction pipeline from a 4D-STEM scan (or a single averaged + diffraction pattern): + + - polar transform of the diffraction patterns + - azimuthal averaging to obtain I(k) + - parametric background fit B(k) (Gaussian model in k² and k⁴) + - reduced structure factor F(k) = 2π · k · [S(k) − 1] + - windowed sine transform of F(k) to recover the reduced PDF G(r) + - optional low-r correction of S(k) for a known density (Kaplow 1965 / + Yoshimoto–Omote 2022), which forces G(r) onto −4π·ρ₀·r on [0, r_cut] + - normalization to g(r) = 1 + G(r) / (4π · r · ρ₀) + + Diffraction data is held in two complementary forms. ``Dataset4dstem`` + holds the input scan with each DP in Cartesian coordinates, indexed as + ``(scan_row, scan_col, n_row, n_col)``. ``Polar4dstem`` holds the + result of rebinning each DP to polar coordinates ``(phi, r_pix)``. + The polar transform is expensive and irreversible, so its result is + cached as a first-class dataset on ``self.polar`` rather than + recomputed on demand. ``from_data`` runs the polar transform (and + optional origin finding) once. + + Attributes + ---------- + polar : Polar4dstem + Polar-transformed diffraction data wrapped by this instance. + input_data : Dataset4dstem or None + Original input dataset that was polar-transformed to produce + ``self.polar``. A ``Dataset2d`` input to ``from_data`` is wrapped + as a 1×1 ``Dataset4dstem`` before being stored here. + device : str + Torch device used for computation. + Ik : torch.Tensor or None + Azimuthally averaged intensity I(k), set by ``calculate_radial_mean``. + bg : torch.Tensor or None + Fitted background B(k), set by ``fit_bg``. + f : torch.Tensor or None + Empirical approximation of the mean-squared atomic scattering factor + ⟨f(k)⟩². Set by ``fit_bg``. + Sk : torch.Tensor or None + Structure factor S(k) = 1 + [I(k) − B(k)] / f(k). Set by ``calculate_Gr``. + Fk : torch.Tensor or None + Reduced structure factor F(k) = 2π · k · [S(k) − 1]. The 2π factor is + explicitly included (py4dstem omits it). Set by ``calculate_Gr``. + Fk_masked : torch.Tensor or None + F(k) after the k window. Set by ``calculate_Gr``. + Fk_damped : torch.Tensor or None + F(k) after the low-r correction, set by ``damp_origin_oscillations``. + reduced_pdf_damped : torch.Tensor or None + G(r) recomputed from the corrected F(k), set by ``damp_origin_oscillations``. + rho0 : float or None + Atomic number density ρ₀ (atoms/ų), used for the low-rcorrection and for g(r). + + Exposed as read-only numpy-valued properties: + + r : NDArray or None + Real-space radial axis in Å. + reduced_pdf : NDArray or None + Reduced pair distribution function G(r), obtained by windowed sine + transform of F(k): + G(r) = (2/π) ∫ F(k) · sin(2π · k · r) dk. + pdf : NDArray or None + Pair distribution function g(r) = 1 + G(r) / (4π · r · ρ₀). + + Examples + -------- + Construct from a 4D-STEM scan and run the standard pipeline: + + >>> import quantem as em + >>> ds = em.core.io.read_4dstem("scan.h5", file_type="arina") + >>> rdf = em.diffraction.PairDistributionFunction.from_data(ds) + >>> rdf.calculate_Gr(k_min_fit=0.05, k_max_window=2.0, r_max=10.0) + >>> rdf.calculate_gr(density=0.05, set_pdf_positive=True) + + Inspect intermediate results: + + >>> rdf.plot_pdf_results(["background_fits", "reduced_sf", "reduced_pdf", "pdf"]) + + Restrict the radial average to a real-space region of interest: + + >>> mask = np.zeros(ds.array.shape[:2], dtype=bool) + >>> mask[300:, 300:] = True + >>> rdf.calculate_Gr(k_min_fit=0.05, k_max_window=2.0, mask_realspace=mask) + """ + + _token = object() + + def __init__( + self, + polar: Polar4dstem, + input_data: Dataset4dstem | None = None, + device: str = "cpu", + _token: object | None = None, + ): + if _token is not self._token: + raise RuntimeError( + "Direct instantiation of PairDistributionFunction is not allowed. " + "Use PairDistributionFunction.from_data() to instantiate this class." + ) + + super().__init__() + self.polar = polar + self.input_data = input_data + self.device = device + + self._r: torch.Tensor | None = None + self._reduced_pdf: torch.Tensor | None = None + self._pdf: torch.Tensor | None = None + self.Ik: torch.Tensor | None = None + self.Sk: torch.Tensor | None = None + self.Fk: torch.Tensor | None = None + self.bg: torch.Tensor | None = None + self.f: torch.Tensor | None = None + self.Fk_masked: torch.Tensor | None = None + self._Ik_mask: NDArray | None = None + self.Fk_damped: torch.Tensor | None = None + self.reduced_pdf_damped: torch.Tensor | None = None + self.rho0: float | None = None + + # ------------------------------------------------------------------ + # Constructors + # ------------------------------------------------------------------ + @classmethod + def from_data( + cls, + data: Dataset2d | Dataset4dstem, + *, + find_origin: bool = True, + origin_method: str = "grid", + origin_row: float | None = None, + origin_col: float | None = None, + origin_array: NDArray | None = None, + ellipse_params: tuple[float, float, float] | None = None, + num_annular_bins: int = 180, + radial_min: float = 0.0, + radial_max: float | None = None, + radial_step: float = 1.0, + two_fold_rotation_symmetry: bool = False, + device: str | None = None, + ) -> Self: + """Create a PairDistributionFunction from a dataset. + + Parameters + ---------- + data : Dataset4dstem or Dataset2d + - ``Dataset4dstem``: triggers origin finding (optional) and polar + transform. + - ``Dataset2d``: single averaged diffraction pattern (e.g. SAED + or a pre-averaged 4DSTEM result); wrapped as a 1x1 scan + internally. + find_origin : bool + If True, find the diffraction origin at each scan position (method chosen + by ``origin_method``). If False, use ``origin_row`` / ``origin_col`` (or + the image center if those are None). + origin_method : {"grid", "descent"} + Origin finder used when ``find_origin=True``. ``"grid"`` (default) is the + global angular-variance search (:func:`find_origin_angular_grid`); + ``"descent"`` is the COM-anchored local descent + (:func:`find_origin_angular_descent`), better on small detectors. + origin_row, origin_col : float or None + Fixed diffraction-space origin in pixels, used only when + ``find_origin=False`` and ``origin_array`` is None. Defaults to + the center of the diffraction pattern. + origin_array : ndarray or None + Pre-computed per-DP origins of shape ``(scan_row, scan_col, 2)``. + When provided, ``find_origin_angular_grid`` is skipped and these origins + are used directly. Takes precedence over ``find_origin`` and + ``origin_row``/``origin_col``. + ellipse_params : tuple of (float, float, float) or None + Elliptical distortion parameters ``(a, b, theta_deg)`` applied + during origin finding and polar transform. + num_annular_bins : int + Number of angular bins in the polar transform. + radial_min, radial_max : float or None + Radial range of the polar transform, in pixels. + radial_step : float + Radial step size in pixels. + two_fold_rotation_symmetry : bool + If True, sample only ``[0, pi)`` in the angular axis. + device : str or None + Torch device used for computation. If None (default), it is + inferred from the input. + + Returns + ------- + PairDistributionFunction + """ + if device is None: + device = data.device if data.array is None else "cpu" + + # Dataset2d input: wrap as a trivial 4D-STEM (1x1 scan) and fall through + if isinstance(data, Dataset2d): + # Dataset2d is numpy-only on dev for now + arr2d = data.array + if arr2d.ndim != 2: + raise ValueError( + f"Found array with shape: {arr2d.shape}. " + "Dataset2d for PairDistributionFunction must be 2D." + ) + arr4 = arr2d[None, None, ...] # (1, 1, n_row, n_col) + + data = Dataset4dstem.from_array( + array=arr4, + name=f"{data.name}_as4dstem" + if getattr(data, "name", None) + else "rdf_4dstem_from_2d", + origin=np.concatenate( + [np.zeros(2, dtype=float), np.asarray(data.origin, dtype=float)] + ), + sampling=np.concatenate( + [np.ones(2, dtype=float), np.asarray(data.sampling, dtype=float)] + ), + units=["pixels", "pixels"] + list(data.units), + signal_units=data.signal_units, + ) + + # Dataset4dstem input: polar-transform it + if isinstance(data, Dataset4dstem): + scan_row, scan_col, n_row, n_col = data.shape + if origin_array is not None: + origin_array = np.asarray(origin_array, dtype=float) + if origin_array.shape != (scan_row, scan_col, 2): + raise ValueError( + f"origin_array has shape {origin_array.shape}, expected " + f"({scan_row}, {scan_col}, 2)." + ) + elif find_origin: + if origin_method == "descent": + origin_array = find_origin_angular_descent( + data, + ellipse_params=ellipse_params, + radial_min=radial_min, + radial_max=radial_max, + radial_step=radial_step, + device=device, + ) + elif origin_method == "grid": + origin_array = find_origin_angular_grid( + data, + ellipse_params=ellipse_params, + num_annular_bins=num_annular_bins, + radial_min=radial_min, + radial_max=radial_max, + radial_step=radial_step, + two_fold_rotation_symmetry=two_fold_rotation_symmetry, + device=device, + ) + else: + raise ValueError( + f"origin_method must be 'grid' or 'descent', got {origin_method!r}." + ) + else: + if origin_row is None: + origin_row = (n_row - 1) / 2.0 + if origin_col is None: + origin_col = (n_col - 1) / 2.0 + origin_array = np.zeros((scan_row, scan_col, 2), dtype=float) + origin_array[..., 0] = origin_row + origin_array[..., 1] = origin_col + + polar = polar_transform( + data, + origin_array=origin_array, + ellipse_params=ellipse_params, + num_annular_bins=num_annular_bins, + radial_min=radial_min, + radial_max=radial_max, + radial_step=radial_step, + two_fold_rotation_symmetry=two_fold_rotation_symmetry, + device=device, + ) + return cls(polar=polar, input_data=data, device=device, _token=cls._token) + + raise TypeError( + f"Got {type(data).__name__}. PairDistributionFunction.from_data " + "accepts Dataset4dstem or Dataset2d. Wrap numpy arrays with " + "Dataset4dstem.from_array or Dataset2d.from_array first." + ) + + # ------------------------------------------------------------------ + # Convenience accessors + # ------------------------------------------------------------------ + @property + def qq(self) -> NDArray: + """ + Scattering vector coordinate array along the radial dimension of `self.polar`, + in physical units (using Polar4dstem.sampling and origin). + """ + # Polar4dstem dims: (scan_row, scan_col, phi, r_pix) + # radial axis is 3 + # origin[3] is the physical q-value at bin 0 (radial_min * pixel_size), + # sampling[3] is the physical step per bin (radial_step * pixel_size). + n = self.polar.shape[3] + origin_r = float(np.asarray(self.polar.origin)[3]) + sampling_r = float(np.asarray(self.polar.sampling)[3]) + return np.arange(n, dtype=float) * sampling_r + origin_r + + @property + def r(self) -> NDArray | None: + """Real-space radial grid as a numpy array.""" + if self._r is None: + return None + return to_numpy(self._r) + + @property + def reduced_pdf(self) -> NDArray | None: + """Reduced pair distribution function G(r) as a numpy array.""" + if self._reduced_pdf is None: + return None + return to_numpy(self._reduced_pdf) + + @property + def pdf(self) -> NDArray | None: + """Pair distribution function g(r) as a numpy array.""" + if self._pdf is None: + return None + return to_numpy(self._pdf) + + # ------------------------------------------------------------------ + # Analysis + # ------------------------------------------------------------------ + + # TODO: add beamstop mask support (mask diffraction-space pixels before + # azimuthal averaging, e.g. to exclude a beam stop shadow) + + def calculate_radial_mean( + self, + mask_realspace: NDArray | None = None, + returnval: bool = False, + ) -> torch.Tensor | None: + """ + Calculate the radial mean intensity from the Polar4dSTEM dataset. + + The polar array is assumed to have shape (scan_row, scan_col, phi, k). + This method computes, for each scan position, the mean over the azimuthal + axis (phi), then averages across scan positions to produce a single 1D + radial curve. This result is stored in ``self.Ik``. + + If a real-space mask is provided, only the selected scan positions are + used in the scan-position average. The computation streams chunks through torch to keep peak + memory low. + + Parameters + ---------- + mask_realspace : NDArray or None, optional + Boolean mask in real space used to select probe positions. + If ``None``, all probe positions are used. + Must have shape (scan_row, scan_col) where True means "include". + returnval : bool, optional + If True, return the computed 1D radial mean tensor. + + Returns + ------- + radial_mean : torch.Tensor or None + If `returnval=True`, returns the 1D radial mean intensity (Nk,). + Otherwise returns None. + """ + polar_data = ( + self.polar.tensor + if self.polar.array is None + else torch.from_numpy(np.ascontiguousarray(self.polar.array)) + ) # (scan_row, scan_col, phi, k) + scan_row, scan_col, n_phi, n_k = polar_data.shape + intensity_sum = torch.zeros(n_k, device=self.device, dtype=torch.float64) + n_valid = 0 + chunk_row = 16 # number of scan rows to process at a time + for row0 in range(0, scan_row, chunk_row): + row1 = min(row0 + chunk_row, scan_row) + chunk = polar_data[row0:row1].to(self.device) + # mean over phi first -> (chunk, scan_col, k) + radial_mean = chunk.mean(dim=2) + if mask_realspace is not None: + mask_chunk = torch.from_numpy(mask_realspace[row0:row1]).to(self.device) + n_chunk = int(mask_chunk.sum()) + if n_chunk == 0: + continue + # sum unmasked intensities in chunk and count for normalization later + intensity_sum += radial_mean[mask_chunk].sum(dim=0) + n_valid += n_chunk + else: + # sum all intensities in chunk and count for normalization later + intensity_sum += radial_mean.sum(dim=(0, 1)) + n_valid += (row1 - row0) * scan_col + if n_valid == 0: + raise ValueError( + "No valid scan positions selected. The real-space mask is " + "all False or the dataset is empty." + ) + self.Ik = (intensity_sum / n_valid).float() + self._Ik_mask = None if mask_realspace is None else np.array(mask_realspace, dtype=bool) + + if returnval: + return self.Ik + else: + return None + + def fit_bg( + self, + Ik: torch.Tensor, + kmin: float | None = None, + kmax: float | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fit a smooth background B(k) to a radial intensity curve I(k). + + The background model is a constant plus two monotonically decaying + terms (adopted from py4DSTEM): + + B(k) = c + + i0 * exp(-k^2 / (2 s0^2)) + + i1 * exp(-k^4 / (2 s1^4)) + + B(k) is later subtracted from I(k) to isolate the diffuse signal, and + f(k) = B(k) - c is used as the denominator in the structure factor + S(k) = 1 + [I(k) − B(k)] / f(k). + + The five parameters are fit by weighted least squares over k in [kmin, kmax]. + Each residual is weighted by a sin² low-k taper that downweights the central + beam, divided by a linear factor that falls with k so that higher k counts + more (see :meth:`fit_bg_batched`). This is the single-curve convenience over + :meth:`fit_bg_batched`, which holds the torch-native Levenberg-Marquardt solver. + + Parameters + ---------- + Ik + 1D radial intensity tensor (Nk,). Produced by + :meth:`calculate_radial_mean`. + kmin, kmax + Restrict the fit to k in [kmin, kmax]. The returned B(k) is + still evaluated over the full k axis. If None, the full k + range is used. + + Returns + ------- + bg : torch.Tensor + Fitted background curve B(k), shape (Nk,). + f : torch.Tensor + Background minus the constant offset, f(k) = B(k) - c, or functionally + similar to ^2(k). Used later to compute the reduced structure factor F(k). + """ + # A single curve is just a batch of one + Ik_row = torch.as_tensor(Ik)[None, :] + bg_stack, f_stack = self.fit_bg_batched(Ik_row, kmin=kmin, kmax=kmax) + self.bg, self.f = bg_stack[0], f_stack[0] + return self.bg, self.f + + def fit_bg_batched( + self, + Ik_stack: torch.Tensor, + kmin: float | None = None, + kmax: float | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fit B(k) for a stack of N radial curves at once (vectorized LM). + + The fit minimises ``sum_k [w(k) * (B(k) - I(k))]^2`` over k in + [kmin, kmax], with the residual weight + + w(k) = taper(k) / (k_end - 0.9 * k + dk) + + Here ``taper`` rises as sin² from 0 at kmin to 1 at kmin + 0.25 1/A + (points where it is below 1e-4 get zero weight), and ``k_end`` is the + last k of the fit range. The weighting scheme follows the background fit + in py4DSTEM. All N curves are fit together in a single solve rather + than one at a time. + + Parameters + ---------- + Ik_stack + (N, Nk) radial means, one row per curve. + kmin, kmax + Fit range. + + Returns + ------- + bg_stack, f_stack + Each (N, Nk) float32: fitted background and f(k) = B(k) - c per row. + """ + k = torch.as_tensor(self.qq, dtype=torch.float64, device=self.device) + if kmin is None: + kmin = float(k.min()) + if kmax is None: + kmax = float(k.max()) + k2 = k**2 + fit_mask = (k >= kmin) & (k <= kmax) + k2_fit = k2[fit_mask] # (Nf,) + k_fit = k[fit_mask] + dk = k[1] - k[0] + # Residual weights over the fit window (length Nf): a sin^2 taper over the + # first 0.25 1/A above kmin suppresses the central-beam tail, and dividing + # by a line that falls with k favours high k. + taper = torch.sin(0.5 * torch.pi * torch.clamp((k_fit - kmin) / 0.25, 0.0, 1.0)) ** 2 + resid_weights = torch.where(taper > 1e-4, taper, 0.0) / (k_fit[-1] - 0.9 * k_fit + dk) + Ik = torch.as_tensor(Ik_stack, dtype=torch.float64, device=self.device) + Ik_fit = Ik[:, fit_mask].clamp(min=1e-10) # (N, Nf) + n = Ik_fit.shape[0] # number of Ik in batch + + # Build Jacobian + # Initial guesses: constant = min(I), amplitudes = median(I) - min(I), widths = mean(k). + c0 = Ik_fit.amin(dim=1) # (N,) + amp = (Ik_fit.median(dim=1).values - c0).clamp(min=1e-10) # (N,) + sig = torch.full((n,), max(float(k.mean()), 1e-3), dtype=torch.float64, device=self.device) + init = torch.stack([c0.clamp(min=1e-10), amp, sig, amp, sig], dim=1) # (N, 5) + # exp(theta) = parameters ensures that they remain positive + theta = torch.log(init) + + def _rj(theta: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Batched residual (N, Nf) and analytic Jacobian (N, Nf, 5). + + Model B(k): + B(k) = c + i0 * E0 + i1 * E1 + where: + E0 = exp(-k^2 / (2 s0^2)) + E1 = exp(-k^4 / (2 s1^4)) + """ + c, i0, s0, i1, s1 = torch.exp(theta).unbind(dim=1) # each (N,) + E0 = torch.exp( + torch.clamp(-k2_fit[None, :] / (2.0 * s0[:, None] ** 2), min=-100.0, max=0.0) + ) + E1 = torch.exp( + torch.clamp( + -(k2_fit[None, :] ** 2) / (2.0 * s1[:, None] ** 4), min=-100.0, max=0.0 + ) + ) + B = c[:, None] + i0[:, None] * E0 + i1[:, None] * E1 + w = resid_weights[None, :] # residual weights, broadcast over the batch + r = (B - Ik_fit) * w + t0 = i0[:, None] * E0 + t1 = i1[:, None] * E1 + jac = torch.stack( + [ + c[:, None] * w, + t0 * w, + t0 * (k2_fit[None, :] / s0[:, None] ** 2) * w, + t1 * w, + t1 * (2.0 * k2_fit[None, :] ** 2 / s1[:, None] ** 4) * w, + ], + dim=2, + ) + return r, jac + + def _loss(theta: torch.Tensor) -> torch.Tensor: + r, _ = _rj(theta) + return (r * r).sum(dim=1) # (N,) + + # Batched Levenberg-Marquardt, replaces scipy.curve_fit. + # recompute Jacobian every iteration and take one step, accepting + # per-curve where the loss improved. lamba is adapted per-curve + # over the course of fitting (decrease on accept toward Gauss-Newton, + # increase on reject toward gradient descent). Converged curves simply stop being updated. + lam = torch.ones(n, dtype=torch.float64, device=self.device) + loss = _loss(theta) + eye = 1e-12 * torch.eye(5, dtype=torch.float64, device=self.device) + max_log_step = 1.0 + for _ in range(200): # outer loop: recompute the Jacobian at the current theta + r, jac = _rj(theta) + JtJ = torch.einsum("nki,nkj->nij", jac, jac) # (N, 5, 5) + Jtr = torch.einsum("nki,nk->ni", jac, r) # (N, 5) + diag = torch.diagonal(JtJ, dim1=1, dim2=2).clamp(min=1e-12) # (N, 5) + theta_base = theta # all trial steps start from here this iter + prev_loss = loss + accepted = torch.zeros(n, dtype=torch.bool, device=self.device) + # inner loop: per-curve lambda search. increase damping and retry until + # each curve finds a downhill step + for _ in range(30): + A = JtJ + lam[:, None, None] * torch.diag_embed(diag) + eye + delta = torch.linalg.solve(A, Jtr) # (N, 5) + # shrink the step if largest component exceeds max_log_step + step_scale = (max_log_step / delta.abs().amax(dim=1).clamp(min=1e-30)).clamp( + max=1.0 + ) + theta_trial = theta_base - delta * step_scale[:, None] + loss_trial = _loss(theta_trial) + # accept improvements per-curve + improved = torch.isfinite(loss_trial) & (loss_trial < loss) & (~accepted) + theta = torch.where(improved[:, None], theta_trial, theta) + loss = torch.where(improved, loss_trial, loss) + accepted = accepted | improved + # Accepted curves relax lambda (toward Gauss-Newton), curves not accepted tighten it (toward gradient descent). + lam = torch.where( + improved, + (lam * 0.3).clamp(min=1e-12), + torch.where(~accepted, (lam * 3.0).clamp(max=1e12), lam), + ) + if bool(accepted.all()): + break + if not bool(accepted.any()): + break # no curve found a downhill step -> all converged / stuck + # Stop once every curve's relative improvement is negligible. + rel = (prev_loss - loss) / (prev_loss + 1e-30) + if float(rel.max()) < 1e-9: + break + + # Evaluate B(k) over the full k-range (not just k_fit) for every curve + with torch.no_grad(): + c, i0, s0, i1, s1 = torch.exp(theta).unbind(dim=1) + E0 = torch.exp( + torch.clamp(-k2[None, :] / (2.0 * s0[:, None] ** 2), min=-100.0, max=0.0) + ) + E1 = torch.exp( + torch.clamp(-(k2[None, :] ** 2) / (2.0 * s1[:, None] ** 4), min=-100.0, max=0.0) + ) + bg = c[:, None] + i0[:, None] * E0 + i1[:, None] * E1 # (N, Nk) + f = (bg - c[:, None]).clamp(min=1e-10 * bg.amax(dim=1, keepdim=True)) + return bg.to(dtype=torch.float32), f.to(dtype=torch.float32) + + def calculate_Gr( + self, + k_min_fit: float | None = None, + k_max_fit: float | None = None, + k_min_window: float | None = None, + k_max_window: float | None = None, + k_lowpass: float | None = None, + k_highpass: float | None = None, + r_max: float = 20.0, + r_step: float = 0.02, + mask_realspace: NDArray | None = None, + damp_origin_oscillations: bool = False, + density: float | None = None, + r_cut: float = 0.8, + max_iter: int = 6, + returnval: bool = False, + ) -> list[NDArray] | None: + """ + Calculate the reduced pair distribution function G(r) from a 4D-STEM dataset. + + This routine: + * Computes the radial mean intensity I(k) from self.polar (optionally + restricted to a real-space mask). + * Fits a smooth background B(k) and associated f(k) using :meth:`fit_bg`. + * Constructs the reduced structure factor F(k) with optional low/highpass filtering. + * Applies the k window. + * Computes the reduced PDF using a discrete sine transform: + G(r) = sum_k sin(2*pi*k*r) * F_windowed(k) + + If ``damp_origin_oscillations=True``, :meth:`damp_origin_oscillations` is + called with ``density`` and the corrected F(k)/G(r) are stored as + ``self.Fk_damped`` and ``self.reduced_pdf_damped``. The density is cached in + ``self.rho0``. + + Stored attributes: + * self.Ik, self.bg, self.Fk, self.Fk_masked + * self.Sk, self.r, self.reduced_pdf + * self.rho0, self.Fk_damped, self.reduced_pdf_damped (if damping) + + Parameters + ---------- + k_min_fit : float, optional + Minimum k (A^-1) for the background fit. + k_max_fit : float or None, optional + Maximum k (A^-1) for the background fit. + k_min_window : float or None, optional + k (A^-1) at which the low-k taper of the window starts. If None, falls back + to ``k_min_fit``. + k_max_window : float or None, optional + Maximum k (A^-1) of the window. The taper reaches zero there. + If None, falls back to ``k_max_fit``. + k_lowpass : float or None, optional + Low-pass Gaussian filter sigma in k-space. + k_highpass : float or None, optional + High-pass Gaussian filter sigma in k-space. + r_max : float, optional + Maximum r (A) for the real-space grid. + r_step : float, optional + Step size in r (A) for the real-space grid. + mask_realspace : NDArray or None, optional + Boolean real-space mask selecting probe positions. + damp_origin_oscillations : bool, optional + If True, run :meth:`damp_origin_oscillations` with ``density`` and store + the corrected F(k)/G(r). + density : float or None, optional + Number density (atoms/A^3). Required when + ``damp_origin_oscillations=True``. + r_cut : float, optional + Upper end (A) of the low-r correction: G(r) is corrected on [0, r_cut]. + Set it below the nearest-neighbour distance. Forwarded to + :meth:`damp_origin_oscillations`. + max_iter : int, optional + Maximum number of iterations of the low-r correction. Forwarded to + :meth:`damp_origin_oscillations`. + returnval : bool, optional + If True, return ``[r, G(r)]`` as numpy arrays. + + Returns + ------- + list[np.ndarray] or None + """ + # clear results from any previous run so stale state doesn't leak + self._pdf = None + self.Fk_damped = None + self.reduced_pdf_damped = None + self.rho0 = None + # this is missing a 2pi term that we add back during the pdf calc later + k_np = np.asarray(self.qq) + k = torch.from_numpy(k_np.astype(np.float32)).to(device=self.device) + dk = k[1] - k[0] + # small epsilon to avoid division by very small k values + k_safe = torch.clamp(k, min=1e-10) + self.kmax_fit = k_max_fit if k_max_fit is not None else float(k.max()) + self.kmin_fit = k_min_fit if k_min_fit is not None else float(k.min()) + # window range defaults to bg-fit range when not specified + self.kmin_window = k_min_window if k_min_window is not None else self.kmin_fit + self.kmax_window = k_max_window if k_max_window is not None else self.kmax_fit + + # Validate the real-space mask, if provided, before using it downstream + mask_bool = None + if mask_realspace is not None: + scan_row, scan_col = self.polar.shape[:2] + mask_realspace = np.asarray(mask_realspace) + if mask_realspace.dtype == bool and mask_realspace.shape == (scan_row, scan_col): + mask_bool = mask_realspace + else: + raise ValueError( + f"Got shape {mask_realspace.shape}. " + "mask_realspace must be boolean array of shape " + f"({scan_row}, {scan_col})." + ) + # reuse the cached radial mean only if it was computed with the same + # real-space mask, otherwise recompute + cached_mask = getattr(self, "_Ik_mask", None) + same_mask = (mask_bool is None and cached_mask is None) or ( + mask_bool is not None + and cached_mask is not None + and np.array_equal(mask_bool, cached_mask) + ) + if self.Ik is not None and same_mask: + Ik = self.Ik + else: + Ik = self.calculate_radial_mean(mask_realspace=mask_bool, returnval=True) + # always refit the background fit bc it is cheap + bg, f = self.fit_bg(Ik, self.kmin_fit, self.kmax_fit) + # prevent division by near-zero values which cause NaNs at high k + f_safe = torch.clamp(f, min=1e-10 * f.max()) + + # below is the standard definition of F(k) used in PDF analysis, except for missing 2pi factor + Fk = (Ik - bg) * k_safe / f_safe + # apply optional frequency filtering for noise reduction + Fk = self._frequency_filtering(Fk, k_lowpass, k_highpass, dk) + wk = self._k_window(k, self.kmin_window, self.kmax_window) + # Compute Sk from Fk BEFORE applying the 2pi scaling, + # so that the low-r correction is on the same scale + self.Sk = torch.ones_like(k) + mask = k > 0 + self.Sk = torch.where(mask, 1.0 + (Fk / k_safe), self.Sk) + # apply that missing 2pi factor + Fk = Fk * 2 * torch.pi + Fk_win = Fk * wk + + # build r in float64 and convert + r = torch.arange(0.0, r_max, r_step, dtype=torch.float64).to( + device=self.device, dtype=torch.float32 + ) + ka, ra = torch.meshgrid(k, r, indexing="ij") + # compute reduced PDF using discrete sine transform + reduced_pdf = ( + (2 / torch.pi) + * dk + * 2 + * torch.pi + * torch.sum( + torch.sin(2 * torch.pi * ra * ka) * Fk_win[:, None], + dim=0, + ) + ) + reduced_pdf[0] = 0 # physically must be at 0 when r = 0 + + self.Ik = Ik + self.bg = bg + self.Fk = Fk + self.Fk_masked = Fk_win + self._r = r + self._reduced_pdf = reduced_pdf + + # optionally correct the unphysical oscillations near the origin for a known density + if damp_origin_oscillations: + if density is None: + raise ValueError("damp_origin_oscillations=True requires `density` (atoms/A^3). ") + self._apply_low_r_correction(float(density), r_cut, max_iter) + + if returnval: + Gr = ( + self.reduced_pdf_damped + if self.reduced_pdf_damped is not None + else self._reduced_pdf + ) + return [to_numpy(self._r), to_numpy(Gr)] + return None + + def calculate_gr( + self, + density: float | None = None, + r_cut: float = 0.8, + max_iter: int = 6, + set_pdf_positive: bool = False, + returnval: bool = False, + ) -> list[NDArray] | None: + """ + Calculate the pair distribution function g(r) from G(r). + + Requires :meth:`calculate_Gr` to have been run first. The density rho0 is + the ``density`` argument, or, if None, the value cached by + ``calculate_Gr(damp_origin_oscillations=True, density=...)``. + + The G(r) used is ``self.reduced_pdf_damped`` if it exists (i.e. the user + chose damping in :meth:`calculate_Gr`), otherwise ``self.reduced_pdf``. If a + ``density`` different from the one used for the damping is given, the damping + is redone with the new ``density``. + + Parameters + ---------- + density : float or None, optional + Number density (atoms/A^3). If None, uses cached if available and raises if there is none. + r_cut : float, optional + Upper end (A) of the low-r correction interval [0, r_cut], used only if + the damping is redone with a new density. + max_iter : int, optional + Maximum number of iterations of the low-r correction, used only if the + damping is redone with a new density. + set_pdf_positive : bool, optional + If True, clamp negative g(r) values to 0. + returnval : bool, optional + If True, return ``[r, g(r)]`` as numpy arrays. + + Returns + ------- + list[np.ndarray] or None + """ + if self._reduced_pdf is None or self._r is None: + raise RuntimeError( + "Reduced PDF not computed." + "Run PairDistributionFunction.calculate_Gr() before calculate_gr()." + ) + + # Determine density + if density is None: + if self.rho0 is None: + raise ValueError( + "calculate_gr needs the number density (atoms/A^3): pass `density`, or run " + "calculate_Gr(damp_origin_oscillations=True, density=...)." + ) + rho0 = self.rho0 + else: + rho0 = density + # if G(r) was damped with a different density, redo the correction + if self.reduced_pdf_damped is not None and self.rho0 != rho0: + self._apply_low_r_correction(rho0, r_cut, max_iter) + + # Use damped G(r) if the user opted into damping, otherwise undamped + Gr = self.reduced_pdf_damped if self.reduced_pdf_damped is not None else self._reduced_pdf + Gr = Gr.clone() + + r = self._r + mask = r > 0 + pdf = torch.ones_like(Gr) + # the formula for g(r) from G(r) is: g(r) = 1 + G(r) / (4 * pi * r * rho0) + pdf = torch.where(mask, 1 + Gr / (4 * torch.pi * r * rho0), torch.zeros_like(pdf)) + if set_pdf_positive: # negative values are unphysical + pdf = torch.maximum(pdf, torch.zeros_like(pdf)) + + self._pdf = pdf + if returnval: + return [to_numpy(self._r), to_numpy(self._pdf)] + return None + + def _apply_low_r_correction(self, density: float, r_cut: float, max_iter: int) -> None: + """Run :meth:`damp_origin_oscillations` and store its results on the object.""" + Fk_damped, G_damped = self.damp_origin_oscillations( + density=density, r_cut=r_cut, max_iter=max_iter + ) + self.rho0 = density + self.Fk_damped = Fk_damped + self.reduced_pdf_damped = G_damped + + def damp_origin_oscillations( + self, + density: float, + r_cut: float = 0.8, + max_iter: int = 6, + tolerance: float = 1e-3, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Correct S(k) so that G(r) follows -4*pi*rho0*r on [0, r_cut]. + + Iterative low-r correction of Kaplow, Strong & Averbach (1965) in the form + given by Yoshimoto & Omote (2022). The deviation of G(r) from the ideal line + -4*pi*rho0*r on [0, r_cut] is transformed back to k-space and subtracted from + S(k), and the transform is repeated until the correction stops changing. The + density must be supplied. + + This method requires that :meth:`calculate_Gr` has already been run, + because it depends on ``self.reduced_pdf``, ``self.r`` and the + k-window bounds). + + Parameters + ---------- + density : float + Number density rho0 (atoms/A^3). + r_cut : float, optional + Upper end (A) of the correction interval [0, r_cut]. Set it below the + nearest-neighbour distance, where g(r) is zero. + max_iter : int, optional + Maximum number of iterations. Keep this small: with more iterations, the + correction can absorb background errors into S(k) rather than just the + origin oscillations. + tolerance : float, optional + Stop when the largest change of the S(k) correction between two + iterations is below ``tolerance`` (in S(k) units). + + Returns + ------- + Fk_win_damped : torch.Tensor + Windowed corrected reduced structure function used for the transform. + G_cor : torch.Tensor + Reduced PDF G(r) with the oscillations near the origin corrected. + """ + if self.Sk is None or self._reduced_pdf is None or self._r is None: + raise RuntimeError( + "This method depends on Sk, reduced_pdf, and r from calculate_Gr. " + "Run PairDistributionFunction.calculate_Gr() before damp_origin_oscillations()." + ) + rho0 = float(density) + + k = torch.from_numpy(np.asarray(self.qq).astype(np.float32)).to(device=self.device) + dk = k[1] - k[0] + k_fit_mask = (k >= self.kmin_fit) & (k <= self.kmax_window) + k_fit = k[k_fit_mask] + ka, ra = torch.meshgrid(k, self._r, indexing="ij") + + # Restrict r to [0, r_cut] for the correction + r_mask = (self._r >= 0.0) & (self._r <= r_cut) + r_short = self._r[r_mask] + k_fit_scaled = k_fit * 2 * torch.pi + k2d_fit, r2d_fit = torch.meshgrid(k_fit_scaled, r_short, indexing="ij") + + Sk_cor = self.Sk.clone() + # the window does not change during the iteration + wk = self._k_window(k, self.kmin_window, self.kmax_window) + # windowed G(r) for the iteration + Fk_win = k * (Sk_cor - 1.0) * wk * 2 * torch.pi + G_iter = ( + (2.0 / torch.pi) + * dk + * 2 + * torch.pi + * torch.sum(torch.sin(2 * torch.pi * ka * ra) * Fk_win[:, None], dim=0) + ) + G_iter[0] = 0.0 + beta_prev = None + for _ in range(max_iter): + # alpha and beta are the ideal and actual contributions to G(r) in the short-r range + # from the current S(k) and G(r) + alpha, beta = self._compute_alpha_beta(k2d_fit, r2d_fit, G_iter[r_mask], r_short) + # converge on the S(k) correction magnitude + if beta_prev is not None: + delta = float(torch.max(torch.abs(beta - beta_prev))) + if delta < tolerance: + break + beta_prev = beta.clone() + # Update S_cor(k) and G_cor + Sk_cor[k_fit_mask] = Sk_cor[k_fit_mask] - beta + rho0 * alpha + Fk_win = k * (Sk_cor - 1.0) * wk * 2 * torch.pi + G_iter = ( + (2.0 / torch.pi) + * dk + * 2 + * torch.pi + * torch.sum(torch.sin(2 * torch.pi * ka * ra) * Fk_win[:, None], dim=0) + ) + G_iter[0] = 0.0 + return Fk_win, G_iter + + # ------------------------------------------------------------------ + # Plotting functions + # ------------------------------------------------------------------ + + PlotName = Literal[ + "radial_mean", + "background_fits", + "reduced_sf", + "reduced_pdf", + "pdf", + "oscillation_damping", + ] + + def plot_pdf_results( + self, + which: Iterable[PlotName] = ("reduced_pdf",), + *, + qmin: float | None = None, + qmax: float | None = None, + rmin: float | None = None, + rmax: float | None = None, + figsize: tuple[float, float] = (6, 4), + returnfigs: bool = False, + ): + """ + Convenience plotting dispatcher. + + Examples + -------- + pdfc.calculate_Gr(...) + pdfc.plot(["radial_mean", "background", "reduced_pdf"]) + """ + mapping = { + "radial_mean": self.plot_radial_mean, + "background_fits": self.plot_background_fits, + "reduced_sf": self.plot_reduced_sf, + "reduced_pdf": self.plot_reduced_pdf, + "pdf": self.plot_pdf, + "oscillation_damping": self.plot_oscillation_damping, + } + + figs = [] + for name in which: + if name not in mapping: + raise ValueError(f"Unknown plot '{name}'. Options: {tuple(mapping)}") + fig = mapping[name]( + qmin=qmin, qmax=qmax, rmin=rmin, rmax=rmax, figsize=figsize, returnfig=returnfigs + ) + if returnfigs: + figs.append(fig) + + return figs if returnfigs else None + + def plot_radial_mean( + self, + qmin: float | None = None, + qmax: float | None = None, + rmin: float | None = None, # accepted for dispatcher compatibility, unused + rmax: float | None = None, # accepted for dispatcher compatibility, unused + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + """ + Plotting radial mean intensity vs scattering vector. + """ + + if self.Ik is None: + raise RuntimeError( + "Radial mean intensity has not been calculated yet." + "Run PairDistributionFunction.calculate_Gr() or PairDistributionFunction.calculate_radial_mean() before plotting." + ) + + x = np.asarray(self.qq) + y = to_numpy(self.Ik) + x, y = self._apply_xrange(x, y, qmin, qmax) + + fig, ax = plt.subplots(figsize=figsize) + ax.plot(x, y, label="Radial Mean Intensity I(k)") + ax.set_xlabel("Scattering Vector q (1/Å)") + ax.set_ylabel("Intensity (a.u.)") + ax.set_title("Radial Mean Intensity vs Scattering Vector") + ax.legend() + ax.set_yscale("log") + plt.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + def plot_background_fits( + self, + qmin: float | None = None, + qmax: float | None = None, + rmin: float | None = None, # accepted for dispatcher compatibility, unused + rmax: float | None = None, # accepted for dispatcher compatibility, unused + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + """ + Plotting background fit vs radial mean intensity. + """ + if self.Ik is None or self.bg is None: + raise RuntimeError( + "Radial mean intensity or background has not been calculated yet." + "Run PairDistributionFunction.calculate_Gr() or both calculate_radial_mean() and calculate_background() before plotting." + ) + + x = np.asarray(self.qq) + y1 = to_numpy(self.Ik) + x, y1 = self._apply_xrange(x, y1, qmin, qmax) + x = np.asarray(self.qq) + y2 = to_numpy(self.bg) + x, y2 = self._apply_xrange(x, y2, qmin, qmax) + + fig, ax = plt.subplots(figsize=figsize) + ax.plot(x, y1, label="Radial Mean Intensity I(k)") + ax.plot(x, y2, label="Background B(k)", linestyle="--") + ax.set_xlabel("Scattering Vector q (1/Å)") + ax.set_ylabel("Intensity (a.u.)") + ax.set_title("Radial Mean Intensity and Background Fit") + ax.legend() + ax.set_yscale("log") + plt.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + def plot_reduced_sf( + self, + qmin: float | None = None, + qmax: float | None = None, + rmin: float | None = None, # accepted for dispatcher compatibility, unused + rmax: float | None = None, # accepted for dispatcher compatibility, unused + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + """ + Plotting reduced structure factor F(k). + """ + if self.Fk_masked is None: + raise RuntimeError( + "Reduced structure factor F(k) has not been calculated yet." + "Run PairDistributionFunction.calculate_Gr() before plotting." + ) + + Fk = getattr(self, "Fk_damped", None) + if Fk is None: + Fk = self.Fk_masked + + x = np.asarray(self.qq) + y = to_numpy(Fk) + x, y = self._apply_xrange(x, y, qmin, qmax) + + fig, ax = plt.subplots(figsize=figsize) + ax.plot(x, y, label="Reduced Structure Factor F(k)") + ax.set_xlabel("Scattering Vector q (1/Å)") + ax.set_ylabel("Reduced Structure Factor F(k)") + plt.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + def plot_reduced_pdf( + self, + qmin: float | None = None, # accepted for dispatcher compatibility, unused + qmax: float | None = None, # accepted for dispatcher compatibility, unused + rmin: float | None = None, + rmax: float | None = None, + padding_frac: float = 0.1, + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + """ + Plotting reduced PDF g(r). + """ + if self._reduced_pdf is None: + raise RuntimeError( + "Reduced PDF has not been calculated yet." + "Run PairDistributionFunction.calculate_Gr() before plotting." + ) + Gr = self.reduced_pdf_damped if self.reduced_pdf_damped is not None else self._reduced_pdf + + x = to_numpy(self._r) + y = to_numpy(Gr) + x, y = self._apply_xrange(x, y, rmin, rmax) + + # Filter out NaN and Inf values to avoid plot errors + valid_mask = np.isfinite(y) + if np.any(valid_mask): + y_valid = y[valid_mask] + y_max = np.max(y_valid) + y_min = np.min(y_valid) + else: + # Fallback if all values are invalid + y_max = 1.0 + y_min = -1.0 + yrange = y_max - y_min + pad = padding_frac * yrange + + fig, ax = plt.subplots(figsize=figsize) + ax.plot(x, y, label="Reduced Pair Distribution Function G(r)") + ax.set_xlabel("Radial Distance r (Å)") + ax.set_ylabel("Reduced Pair Distribution Function G(r)") + ax.set_ylim(y_min - pad, y_max + pad) + plt.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + def plot_pdf( + self, + qmin: float | None = None, # accepted for dispatcher compatibility, unused + qmax: float | None = None, # accepted for dispatcher compatibility, unused + rmin: float | None = None, + rmax: float | None = None, + padding_frac: float = 0.1, + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + """ + Plotting pair distribution function g(r). + """ + if self._reduced_pdf is None or self._pdf is None: + raise RuntimeError( + "PDF has not been calculated yet." + "Run PairDistributionFunction.calculate_gr() before plotting." + ) + + x = to_numpy(self._r) + y = to_numpy(self._pdf) + x, y = self._apply_xrange(x, y, rmin, rmax) + + # Filter out NaN and Inf values to avoid plot errors + valid_mask = np.isfinite(y) + if np.any(valid_mask): + y_valid = y[valid_mask] + y_max = np.max(y_valid) + y_min = np.min(y_valid) + else: + # Fallback if all values are invalid + y_max = 1.0 + y_min = -1.0 + yrange = y_max - y_min + pad = padding_frac * yrange + + fig, ax = plt.subplots(figsize=figsize) + ax.plot(x, y, label="Pair Distribution Function g(r)") + ax.set_xlabel("Radial Distance r (Å)") + ax.set_ylabel("Pair Distribution Function g(r)") + ax.set_ylim(y_min - pad, y_max + pad) + plt.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + def plot_oscillation_damping( + self, + qmin: float | None = None, # accepted for dispatcher compatibility, unused + qmax: float | None = None, # accepted for dispatcher compatibility, unused + rmin: float | None = None, + rmax: float | None = None, + padding_frac: float = 0.1, + figsize: tuple[float, float] = (8, 4), + returnfig: bool = False, + ): + if self.Fk_masked is None or self.Fk_damped is None or self.reduced_pdf_damped is None: + raise RuntimeError( + "Oscillation damping data not available. " + "Run calculate_Gr(damp_origin_oscillations=True) first." + ) + + k = np.asarray(self.qq) + + # Convert torch tensors to numpy for plotting + Fk_masked = to_numpy(self.Fk_masked) + Fk_damped = to_numpy(self.Fk_damped) + r = to_numpy(self._r) + reduced_pdf = to_numpy(self._reduced_pdf) + reduced_pdf_damped = to_numpy(self.reduced_pdf_damped) + + fig, axes = plt.subplots(2, 2, figsize=figsize) + + # F(k) + axS_top = axes[0, 0] + axS_res = axes[1, 0] + axS_top.plot(k, Fk_masked, label="F_obs(k)", color="gray") + axS_top.plot(k, Fk_damped, label="F_cor(k)", color="red") + axS_top.set_xlabel("k (A$^{-1}$)") + axS_top.set_ylabel("F(k)") + axS_top.legend() + + axS_res.plot(k, Fk_damped - Fk_masked, color="blue") + axS_res.set_xlabel("k (A$^{-1}$)") + axS_res.set_ylabel("F_cor - F_obs") + + # G(r) + axG_top = axes[0, 1] + axG_res = axes[1, 1] + axG_top.plot(r, reduced_pdf, label="G_obs(r)", color="gray") + axG_top.plot(r, reduced_pdf_damped, label="G_cor(r)", color="red") + axG_top.set_xlabel("r (A)") + axG_top.set_ylabel("G(r)") + axG_top.legend() + + axG_res.plot(r, reduced_pdf_damped - reduced_pdf, color="blue") + axG_res.set_xlabel("r (A)") + axG_res.set_ylabel("G_cor - G_obs") + + fig.tight_layout() + + if returnfig: + return fig + else: + plt.show() + + # ------------------------------------------------------------------ + # Helper functions + # ------------------------------------------------------------------ + + def _frequency_filtering( + self, + Fk: torch.Tensor, + k_lowpass: float | None, + k_highpass: float | None, + dk: torch.Tensor, + ) -> torch.Tensor: + """Band pass filtering using torch""" + if ( + k_lowpass is not None + and k_lowpass > 0.0 + and k_highpass is not None + and k_highpass > 0.0 + ): + if k_highpass > k_lowpass: + raise ValueError( + "k_highpass is greater than k_lowpass." + "Gaussian band-pass filtering requires k_highpass < k_lowpass." + ) + low_kernel = gaussian_kernel_1d(k_lowpass / dk.item()).to(self.device) + high_kernel = gaussian_kernel_1d(k_highpass / dk.item()).to(self.device) + Fk_low = gaussian_filter_1d(Fk, low_kernel) + Fk_high = gaussian_filter_1d(Fk, high_kernel) + Fk = Fk_high - Fk_low + elif k_lowpass is not None and k_lowpass > 0.0: + low_kernel = gaussian_kernel_1d(k_lowpass / dk.item()).to(self.device) + Fk = gaussian_filter_1d(Fk, low_kernel) + elif k_highpass is not None and k_highpass > 0.0: + high_kernel = gaussian_kernel_1d(k_highpass / dk.item()).to(self.device) + Fk_high = gaussian_filter_1d(Fk, high_kernel) + Fk = Fk - Fk_high + return Fk + + def _k_window(self, k: torch.Tensor, kmin: float, kmax: float) -> torch.Tensor: + """ + Combined low-q taper and high-q Lorch window. + + - zero outside [kmin, kmax] + - smoothly rises from 0->1 near kmin using a sin^2 ramp over 10% of the band + - applies a Lorch-style sinc factor over the full in-band region: + sin(pi * k/kmax) / (pi * k/kmax) + """ + edge_width_low = 0.1 * (kmax - kmin) + low = (k >= kmin) & (k < kmin + edge_width_low) + t = (k - kmin) / edge_width_low + wk = torch.ones_like(k) + wk = torch.where(low, torch.sin(0.5 * torch.pi * t) ** 2, wk) + wk = torch.where((k < kmin) | (k > kmax), torch.zeros_like(wk), wk) + x = k / kmax + inband = (k >= kmin) & (k <= kmax) + sinc_val = torch.where( + x == 0, torch.ones_like(x), torch.sin(torch.pi * x) / (torch.pi * x) + ) + return wk * torch.where(inband, sinc_val, torch.zeros_like(k)) + + def _compute_alpha_beta( + self, + Q2d: torch.Tensor, + r2d: torch.Tensor, + G_beta: torch.Tensor, + r_1d: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute Yoshimoto-Omote alpha(Q) and beta(Q) integrals used for the low-r correction. + """ + Qsafe = torch.where( + Q2d == 0.0, + torch.tensor(1e-12, device=self.device, dtype=torch.float32), + Q2d, + ) + alpha_int = -4 * torch.pi * r2d * torch.sin(Qsafe * r2d) / Qsafe + beta_int = G_beta.unsqueeze(0) * torch.sin(Qsafe * r2d) / Qsafe + alpha = torch.trapezoid(alpha_int, x=r_1d, dim=1) + beta = torch.trapezoid(beta_int, x=r_1d, dim=1) + return alpha, beta + + def _apply_xrange( + self, + x: NDArray, + y: NDArray, + xmin: float | None, + xmax: float | None, + ) -> tuple[NDArray, NDArray]: + if xmin is None and xmax is None: + return x, y + xmin_eff = x.min() if xmin is None else xmin + xmax_eff = x.max() if xmax is None else xmax + if xmax_eff <= xmin_eff: + raise ValueError(f"xmax must be > xmin (got xmin={xmin_eff}, xmax={xmax_eff}).") + m = (x >= xmin_eff) & (x <= xmax_eff) + # avoid empty plots + if not np.any(m): + raise ValueError("Requested plot range contains no data.") + return x[m], y[m] diff --git a/src/quantem/diffraction/polar_transform.py b/src/quantem/diffraction/polar_transform.py new file mode 100644 index 00000000..bac53001 --- /dev/null +++ b/src/quantem/diffraction/polar_transform.py @@ -0,0 +1,1051 @@ +import warnings + +import numpy as np +import torch +import torch.nn.functional as F +from numpy.typing import NDArray +from tqdm import tqdm + +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.datastructures.polar4dstem import Polar4dstem +from quantem.core.utils.utils import to_numpy + +# Percentile of the per-origin in-frame radius budget used to pick radial_max +# when the caller does not supply one. A plain minimum lets a single failed +# origin fit truncate the q-axis of an entire dataset, so a small fraction of +# the worst origins is allowed to fall outside (and zero-pad) instead +RADIAL_MAX_PERCENTILE = 0.1 + +# Standard DPs use (row, col) convention. Polar coordinates use (phi, r_pix), +# grid_sample's grid tensor requires them to be ordered (col, row) +# but is noted where the call occures + + +def find_origin_angular_grid( + data: Dataset4dstem, + *, + ellipse_params: tuple[float, float, float] | None = None, + num_annular_bins: int = 180, + radial_min: float = 0.0, + radial_max: float | None = None, + radial_step: float = 2.0, + two_fold_rotation_symmetry: bool = False, + device: str = "cpu", + batch_size: int = 16, + local_margin: int = 40, + kpow: float = 0.0, +) -> NDArray: + """ + Automatic diffraction center finding by minimizing angular intensity + variation in the polar transform. A correctly centered diffraction + pattern has uniform intensity along each ring, so the center that + minimizes the angular standard deviation is the true beam center. + + Uses a coarse-to-fine search on the mean diffraction pattern to find + a global center, then refines per scan position to account for descan + across the scan. + + Parameters + ---------- + data : Dataset4dstem + A 4D-STEM dataset (or 2D wrapped as 4D) + ellipse_params : tuple or None + Ellipse parameters (a, b, theta_deg) for distortion correction + num_annular_bins : int + Number of angular bins for the final polar transform + radial_min : float + Minimum radius in pixels + radial_max : float or None + Maximum radius in pixels + radial_step : float + Radial step size in pixels for the search polar grid + two_fold_rotation_symmetry : bool + If True, use only 0 to pi range for angles + device : str + Torch device for computation + batch_size : int + Number of scan positions evaluated per coarse-stage kernel call. + Larger values reduce per-iteration overhead but use more memory. + local_margin : int + Half-width (in pixels) of the search window used to refine each + scan position's origin. After the global center is found on the + mean DP, each DP's origin is searched within a + ``(2*local_margin+1)`` square window centered on the global + origin. Set this large enough to cover the worst-case descan + drift across the scan. + kpow : float + Up-weight each ring's contribution to the score by ``k**kpow`` + + Returns + ------- + origin_array : np.ndarray + Array of shape (scan_row, scan_col, 2) containing (row, col) origin + estimates in pixels. + """ + if data.ndim == 2: + n_row, n_col = data.shape + scan_row, scan_col = 1, 1 + elif data.ndim == 4: + scan_row, scan_col, n_row, n_col = data.shape + else: + raise ValueError( + f" Got array with shape {data.shape}." + "To use find_origin_angular_grid, pass a 2D or 4DSTEM dataset." + ) + + # Move the full dataset to the chosen device once + dps = ( + torch.from_numpy(np.ascontiguousarray(data.array)) + if data.array is not None + else data.tensor + ) + array_t = dps.to(device=device, dtype=torch.float32) + if array_t.ndim == 2: + array_t = array_t[None, None] + # first get COM of mean DP because it gives a robust rough center + mean_dp_t = array_t.mean(dim=(0, 1)) + total_intensity = mean_dp_t.sum() + row_grid_t = torch.arange(n_row, dtype=torch.float32, device=device)[:, None] + col_grid_t = torch.arange(n_col, dtype=torch.float32, device=device)[None, :] + com_row = int(round(float(((row_grid_t * mean_dp_t).sum() / total_intensity).item()))) + com_col = int(round(float(((col_grid_t * mean_dp_t).sum() / total_intensity).item()))) + # Radial max of the search polar grid, so dp_mean search candidates + # (at ±global_margin from COM) stay within image bounds + # in-image. Single pos candidates further from COM might be out of bounds + # and are masked with [safe_low, safe_high_*] if so + # (zero-padded samples would otherwise produce a falsely low score) + # global_margin is the half-width of the candidate-center search around the COM. + com_edge_budget = min(com_row, com_col, (n_row - 1) - com_row, (n_col - 1) - com_col) + global_margin = int(min(40, max(2, com_edge_budget // 2))) + safe_radial_max = float( + min( + com_row - global_margin, + (n_row - 1) - (com_row + global_margin), + com_col - global_margin, + (n_col - 1) - (com_col + global_margin), + ) + ) + if radial_max is not None: + safe_radial_max = min(safe_radial_max, float(radial_max)) + if safe_radial_max <= radial_min: + safe_radial_max = radial_min + radial_step + safe_low = int(np.ceil(safe_radial_max)) + safe_high_row = n_row - 1 - safe_low + safe_high_col = n_col - 1 - safe_low + # Internal search-grid resolution. Balancing speed against robustness + search_n_phi = 18 + local_coarse_step = 5 + offset_row, offset_col, _, radial_bins = _build_polar_sampling_offsets( + ellipse_params, + search_n_phi, + radial_min, + safe_radial_max, + radial_step, + two_fold_rotation_symmetry, + device, + ) + n_r = radial_bins.numel() + min_r_idx = 0 + max_r_idx = int(np.ceil(0.9 * n_r)) + # k-weighting, none when kpow == 0 + ring_weights = ( + None + if kpow == 0.0 + else (radial_bins[min_r_idx:max_r_idx].to(device) ** kpow).to(torch.float32) + ) + # Normalize offsets to [-1, 1] because grid_sample expects normalized coordinates + col_norm_scale = 2.0 / (n_col - 1) + row_norm_scale = 2.0 / (n_row - 1) + base_col_norm = offset_col * col_norm_scale + base_row_norm = offset_row * row_norm_scale + + # Mean-DP global center search: coarse → fine, masking candidates + # whose polar grid would extend OOB at each step. + mean_dp_batch = mean_dp_t[None, None] + # Coarse: step=4 over ±global_margin around the COM + rows, cols, grids = _build_candidate_grids( + base_col_norm, + base_row_norm, + com_row, + com_col, + global_margin, + n_row, + n_col, + col_norm_scale, + row_norm_scale, + device, + step=2, + ) + scores = _angular_std_scores(mean_dp_batch, grids, min_r_idx, max_r_idx, ring_weights) + valid = ( + (rows >= safe_low) & (rows <= safe_high_row) & (cols >= safe_low) & (cols <= safe_high_col) + ) + best = int(scores.masked_fill(~valid, float("inf")).argmin().item()) + coarse_row, coarse_col = int(rows[best].item()), int(cols[best].item()) + # Fine: step=1 over ±6 around the coarse winner + rows, cols, grids = _build_candidate_grids( + base_col_norm, + base_row_norm, + coarse_row, + coarse_col, + 10, + n_row, + n_col, + col_norm_scale, + row_norm_scale, + device, + step=1, + ) + scores = _angular_std_scores(mean_dp_batch, grids, min_r_idx, max_r_idx, ring_weights) + valid = ( + (rows >= safe_low) & (rows <= safe_high_row) & (cols >= safe_low) & (cols <= safe_high_col) + ) + best = int(scores.masked_fill(~valid, float("inf")).argmin().item()) + global_row, global_col = int(rows[best].item()), int(cols[best].item()) + + # Per-scan-position refinement (coarse → medium → fine) for descan + # medium and fine search per-DP around the previous winner + coarse_rows, coarse_cols, coarse_grids = _build_candidate_grids( + base_col_norm, + base_row_norm, + global_row, + global_col, + local_margin, + n_row, + n_col, + col_norm_scale, + row_norm_scale, + device, + step=local_coarse_step, + ) + coarse_valid = ( + (coarse_rows >= safe_low) + & (coarse_rows <= safe_high_row) + & (coarse_cols >= safe_low) + & (coarse_cols <= safe_high_col) + ) + n_coarse = coarse_grids.shape[0] + # Per-DP relative offsets used by the medium and fine stages + med_search_range = torch.arange( + -local_coarse_step, local_coarse_step + 1, 1, dtype=torch.long, device=device + ) + med_drow, med_dcol = ( + m.reshape(-1) for m in torch.meshgrid(med_search_range, med_search_range, indexing="ij") + ) + fine_search_range = torch.arange(-2, 3, dtype=torch.long, device=device) # [-2,-1,0,1,2] px + fine_drow, fine_dcol = ( + m.reshape(-1) for m in torch.meshgrid(fine_search_range, fine_search_range, indexing="ij") + ) + flat_dps_t = array_t.reshape(-1, n_row, n_col) + n_pos = flat_dps_t.shape[0] + origin_flat_t = torch.zeros(n_pos, 2, dtype=torch.float32, device=device) + + def refine(dp_batch, current_row, current_col, drow, dcol): + """Score candidates per DP and return the best (row, col) per DP, plus + the raw score grid and validity mask (used for sub-pixel refinement). + Invalid (out of bounds) candidates are masked for the argmin only.""" + n_cands = drow.numel() + cand_rows = (current_row[:, None] + drow[None, :]).clamp(0, n_row - 1) + cand_cols = (current_col[:, None] + dcol[None, :]).clamp(0, n_col - 1) + g_col = ( + base_col_norm + (cand_cols.reshape(-1).float() * col_norm_scale - 1.0)[:, None, None] + ) + g_row = ( + base_row_norm + (cand_rows.reshape(-1).float() * row_norm_scale - 1.0)[:, None, None] + ) + grids = torch.stack([g_col, g_row], dim=-1) + # Stack each DP's candidate grids along the grid's row dim, so each DP + # is sampled from its own image instead of from n_cands copies of it + polars = F.grid_sample( + dp_batch, + grids.view(dp_batch.shape[0], n_cands * search_n_phi, n_r, 2), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + region = polars.view(dp_batch.shape[0], n_cands, *base_col_norm.shape)[ + ..., min_r_idx:max_r_idx + ] + std_r = _angular_std(region, dim=2) + mean_r = region.mean(dim=2) + if ring_weights is not None: + std_r = std_r * ring_weights + mean_r = mean_r * ring_weights + scores = std_r.sum(dim=2) / (mean_r.sum(dim=2) + 1e-6) + valid = ( + (cand_rows >= safe_low) + & (cand_rows <= safe_high_row) + & (cand_cols >= safe_low) + & (cand_cols <= safe_high_col) + ) + best = scores.masked_fill(~valid, float("inf")).argmin(dim=1) + best_row = cand_rows.gather(1, best[:, None]).squeeze(1) + best_col = cand_cols.gather(1, best[:, None]).squeeze(1) + return best_row, best_col, scores, valid + + n_not_converged = 0 # positions whose fine search did not bracket the minimum + pbar = tqdm(total=n_pos, desc="Finding origin for each scan position") + for start in range(0, n_pos, batch_size): + end = min(start + batch_size, n_pos) + n_dp = end - start # number of diffraction patterns in this batch chunk + dp_b = flat_dps_t[start:end].unsqueeze(1) # (B, 1, H, W) + # Coarse (shared grids): broadcast B DPs across n_coarse candidate + # grids in one grid_sample call by stacking DPs in the channel dim + # and stride-0 expanding along the candidate dim + polars_coarse = F.grid_sample( + dp_b.transpose(0, 1).expand(n_coarse, n_dp, n_row, n_col), + coarse_grids, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + region_coarse = polars_coarse[:, :, :, min_r_idx:max_r_idx] + std_r_coarse = _angular_std(region_coarse, dim=2) + mean_r_coarse = region_coarse.mean(dim=2) + if ring_weights is not None: + std_r_coarse = std_r_coarse * ring_weights + mean_r_coarse = mean_r_coarse * ring_weights + scores_coarse = std_r_coarse.sum(dim=2) / (mean_r_coarse.sum(dim=2) + 1e-6) + scores_coarse = scores_coarse.masked_fill(~coarse_valid[:, None], float("inf")) + best_coarse = scores_coarse.argmin(dim=0) # best candidate per DP + current_row, current_col = coarse_rows[best_coarse], coarse_cols[best_coarse] + # Medium: per-DP search around the coarse winner + current_row, current_col, _, _ = refine( + dp_b, current_row, current_col, med_drow, med_dcol + ) + # Fine: per-DP search around the medium winner + best_row, best_col, fine_scores, fine_valid = refine( + dp_b, current_row, current_col, fine_drow, fine_dcol + ) + # Sub-pixel refine: fit a 2D quadratic to the 3x3 of fine scores around + # the integer winner and add the vertex offset. The flat (B, n) grid + # reshapes to (B, side, side) with [:, i, j] = (fine_search_range[i], fine_search_range[j]). + side = fine_search_range.numel() + scores_grid = fine_scores.view(n_dp, side, side) + valid_grid = fine_valid.view(n_dp, side, side) + flat_best = ( + scores_grid.masked_fill(~valid_grid, float("inf")).view(n_dp, -1).argmin(dim=1) + ) + # unravel the flat argmin (best candidate pixel) into 2D (row, col) indices in the side x side grid + i_star, j_star = flat_best // side, flat_best % side + # Gather the 3x3 patch around (i_star, j_star) + batch_idx = torch.arange(n_dp, device=device) + ii = torch.stack([i_star - 1, i_star, i_star + 1], dim=1).clamp(0, side - 1) + jj = torch.stack([j_star - 1, j_star, j_star + 1], dim=1).clamp(0, side - 1) + patch = scores_grid[batch_idx[:, None, None], ii[:, :, None], jj[:, None, :]] + # A winner "on the border" of the fine grid (its 3x3 would fall off the + # edge) means coarse+medium did not bracket the minimum -- genuine + # non-convergence, counted and warned about below. Keep the integer + # winner there; apply the sub-pixel offset only to interior winners. + on_border = (i_star < 1) | (i_star > side - 2) | (j_star < 1) | (j_star > side - 2) + n_not_converged += int(on_border.sum()) + offset = _quadratic_subpixel_offset(patch).to(torch.float32) # (B, 2) + offset = torch.where(on_border[:, None], torch.zeros_like(offset), offset) + origin_flat_t[start:end, 0] = best_row.to(torch.float32) + offset[:, 0] + origin_flat_t[start:end, 1] = best_col.to(torch.float32) + offset[:, 1] + pbar.update(n_dp) + pbar.close() + if n_not_converged: + warnings.warn( + f"find_origin_angular_grid: {n_not_converged} of {n_pos} scan positions did not " + "converge to a sub-pixel origin (descan may exceed local_margin). The integer-pixel origin was " + "used at those positions.", + stacklevel=2, + ) + return origin_flat_t.cpu().numpy().reshape(scan_row, scan_col, 2) + + +def _ellipse_stretch_factors( + ellipse_params: tuple[float, float, float] | None, +) -> tuple[float, float]: + """Largest factor by which polar sampling at radius ``r_pix`` can exceed + ``r_pix`` along each image axis. + + :func:`_polar_to_cartesian_offsets` stretches the ellipse-frame major axis + by ``a / b``, so a sample nominally at ``r_pix`` actually lands as far as + ``r_pix * stretch`` from the origin. Dividing an origin's distance-to-edge + budget by these gives the largest radius that still samples inside the + frame. Both factors are 1 when no ellipse correction is applied. + """ + if ellipse_params is None: + return 1.0, 1.0 + if len(ellipse_params) != 3: + raise ValueError("ellipse_params must be (a, b, theta_deg).") + a, b, theta_deg = ellipse_params + ab = float(a) / float(b) + theta = np.deg2rad(float(theta_deg)) + # offset_col = (a/b) r cos(alpha) cos(theta) - r sin(alpha) sin(theta), so + # max over alpha is r * hypot((a/b) cos(theta), sin(theta)); likewise for row. + stretch_row = float(np.hypot(ab * np.sin(theta), np.cos(theta))) + stretch_col = float(np.hypot(ab * np.cos(theta), np.sin(theta))) + return stretch_row, stretch_col + + +def polar_transform( + data: Dataset4dstem, + origin_array: NDArray | torch.Tensor | None = None, + ellipse_params: tuple[float, float, float] | None = None, + num_annular_bins: int = 180, + radial_min: float = 0.0, + radial_max: float | None = None, + radial_step: float = 1.0, + two_fold_rotation_symmetry: bool = False, + name: str | None = None, + signal_units: str | None = None, + scan_pos: tuple[int, int] | None = None, + device: str = "cpu", + batch_size: int = 128, +) -> Polar4dstem | torch.Tensor: + if data.ndim != 4: + raise ValueError( + f"Found array with shape: {data.shape}. " + "polar_transform requires a 4D-STEM dataset (ndim=4)." + ) + scan_row, scan_col, n_row, n_col = data.shape + + # Standardize origin_array input + if isinstance(origin_array, torch.Tensor): + origin_array = to_numpy(origin_array) + origin_array = np.asarray(origin_array) if origin_array is not None else None + if origin_array is None: + center = np.array([(n_row - 1) / 2.0, (n_col - 1) / 2.0], dtype=float) + origins = np.broadcast_to(center, (scan_row, scan_col, 2)).copy() + elif origin_array.shape == (2,): + origins = np.empty((scan_row, scan_col, 2), dtype=float) + origins[...] = origin_array + elif origin_array.shape == (scan_row, scan_col, 2): + origins = origin_array + else: + raise ValueError( + f" Got {origin_array.shape}. " + "origin_array must have shape None, (2,) or (scan_row, scan_col, 2)." + ) + + # If scan_pos is provided, compute polar transform only for that position + if scan_pos is not None: + i_row, i_col = scan_pos + dps = ( + torch.from_numpy(np.ascontiguousarray(data.array)) + if data.array is not None + else data.tensor + ) + dp = dps[i_row, i_col].to(device=device, dtype=torch.float32) + r0 = float(origins[i_row, i_col, 0]) + c0 = float(origins[i_row, i_col, 1]) + # Clamp radial range to image bounds for this origin, accounting for the + # ellipse correction stretching the sampling beyond r_pix + if radial_max is None: + stretch_row, stretch_col = _ellipse_stretch_factors(ellipse_params) + radial_max_eff = float( + min( + r0 / stretch_row, + ((n_row - 1) - r0) / stretch_row, + c0 / stretch_col, + ((n_col - 1) - c0) / stretch_col, + ) + ) + else: + radial_max_eff = float(radial_max) + if radial_max_eff <= radial_min: + radial_max_eff = radial_min + radial_step + # Build offsets, translate to this origin, normalize for grid_sample + offset_row, offset_col, _, _ = _build_polar_sampling_offsets( + ellipse_params, + num_annular_bins, + radial_min, + radial_max_eff, + radial_step, + two_fold_rotation_symmetry, + device, + ) + col_norm = 2.0 * (offset_col + c0) / (n_col - 1) - 1.0 + row_norm = 2.0 * (offset_row + r0) / (n_row - 1) - 1.0 + # grid_sample requires (col, row) ordering in the last dim + grid = torch.stack([col_norm, row_norm], dim=-1).unsqueeze(0) # (1, n_phi, n_r, 2) + polar2d = F.grid_sample( + dp[None, None], + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + return polar2d.squeeze(0).squeeze(0) # (n_phi, n_r) + + # Every scan position must map to the same-size polar grid (required for a + # uniform 4D output), so one radius has to serve all origins. Using from a + # low percentile rather than the minimum prevents a of failed origin fit + # from truncating the whole dataset's q-axis. + if radial_max is None: + stretch_row, stretch_col = _ellipse_stretch_factors(ellipse_params) + r_row_pos = origins[:, :, 0] / stretch_row + r_row_neg = ((n_row - 1) - origins[:, :, 0]) / stretch_row + r_col_pos = origins[:, :, 1] / stretch_col + r_col_neg = ((n_col - 1) - origins[:, :, 1]) / stretch_col + radial_max_eff_array = np.minimum.reduce([r_row_pos, r_row_neg, r_col_pos, r_col_neg]) + radial_max = float( + max( + np.percentile(radial_max_eff_array, RADIAL_MAX_PERCENTILE), + radial_min + radial_step, + ) + ) + # Positions in the percentile tail sample past the detector edge and are + # zero-filled there. The tail is non-empty by construction, and most of + # it only grazes the edge by a fraction of a pixel, which costs nothing. + # Warn only where a position loses a whole radial bin or more: that is a + # bad origin fit, not an artifact of choosing a percentile. + shortfall = radial_max - radial_max_eff_array + n_bad = int((shortfall > radial_step).sum()) + if n_bad: + worst = float(radial_max_eff_array.min()) + warnings.warn( + f"polar_transform: {n_bad} of {radial_max_eff_array.size} scan " + "positions have origins too close to the detector edge for " + f"radial_max={radial_max:.2f} px and lose more than one radial " + f"bin to zero-padding (the worst supports only {worst:.2f} px). " + "This is likely due to a failed origin fit.", + stacklevel=2, + ) + + # Build origin-independent polar offsets ONCE. Only the per-origin shift + # changes from one scan position to the next, so we can reuse these. + offset_row, offset_col, phi_bins, radial_bins = _build_polar_sampling_offsets( + ellipse_params, + num_annular_bins, + radial_min, + float(radial_max), + radial_step, + two_fold_rotation_symmetry, + device, + ) + n_phi = phi_bins.numel() + n_r = radial_bins.numel() + radial_max_eff = float(radial_max) + + # Pre-normalize offsets into grid_sample's [-1, 1] coordinate convention + col_norm_scale = 2.0 / (n_col - 1) + row_norm_scale = 2.0 / (n_row - 1) + base_col_norm = offset_col * col_norm_scale # (n_phi, n_r) + base_row_norm = offset_row * row_norm_scale # (n_phi, n_r) + + # Flatten scan dims so we can iterate in flat batches + n_pos = scan_row * scan_col + dps = ( + torch.from_numpy(np.ascontiguousarray(data.array)) + if data.array is not None + else data.tensor + ) + dp_view = dps.reshape(n_pos, n_row, n_col) + origins_t = torch.from_numpy( + np.ascontiguousarray(origins.reshape(n_pos, 2), dtype=np.float32) + ).to(device) + + out = torch.empty((n_pos, n_phi, n_r), dtype=torch.float32, device=device) + for start in tqdm(range(0, n_pos, batch_size), desc="Polar transform"): + end = min(start + batch_size, n_pos) + # Translate the precomputed offsets to each origin in this batch + row_origins = origins_t[start:end, 0] + col_origins = origins_t[start:end, 1] + grid_col = base_col_norm.unsqueeze(0) + (col_origins * col_norm_scale - 1.0)[:, None, None] + grid_row = base_row_norm.unsqueeze(0) + (row_origins * row_norm_scale - 1.0)[:, None, None] + # grid_sample requires (col, row) ordering in the last dim + grids = torch.stack([grid_col, grid_row], dim=-1) + + dp_batch = dp_view[start:end].to(device=device, dtype=torch.float32) + polars = F.grid_sample( + dp_batch.unsqueeze(1), + grids, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + out[start:end] = polars.squeeze(1) + out = out.reshape(scan_row, scan_col, n_phi, n_r) + + # Get polar axes in physical units matching the input dataset's calibration + phi_range = np.pi if two_fold_rotation_symmetry else 2.0 * np.pi + phi_step_deg = (phi_range / float(n_phi)) * (180.0 / np.pi) + sampling = np.zeros(4, dtype=float) + origin = np.zeros(4, dtype=float) + sampling[0:2] = np.asarray(data.sampling)[0:2] + sampling[2] = phi_step_deg + sampling[3] = float(np.asarray(data.sampling)[-1]) * radial_step + origin[0:2] = np.asarray(data.origin)[0:2] + origin[2] = 0.0 + origin[3] = radial_min * float(np.asarray(data.sampling)[-1]) + units = [ + data.units[0], + data.units[1], + "deg", + data.units[-1], + ] + metadata = dict(data.metadata) + metadata.update( + { + "polar_radial_min": float(radial_min), + "polar_radial_max": float(radial_max_eff), + "polar_radial_step": float(radial_step), + "polar_num_annular_bins": int(n_phi), + "polar_two_fold_rotation_symmetry": bool(two_fold_rotation_symmetry), + "polar_ellipticity": tuple(ellipse_params) if ellipse_params is not None else None, + } + ) + return Polar4dstem( + tensor=out, + name=name if name is not None else f"{data.name}_polar", + origin=origin, + sampling=sampling, + units=units, + signal_units=signal_units if signal_units is not None else data.signal_units, + metadata=metadata, + origin_array=origins, + _token=Polar4dstem._token, + ) + + +def _polar_to_cartesian_offsets( + phi: torch.Tensor, + r_pix: torch.Tensor, + ellipse_params: tuple[float, float, float] | None, + device: str = "cpu", +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert polar (phi, r_pix) grids to Cartesian (row, col) pixel offsets + from the origin, optionally correcting for elliptical distortion. + + Returns ``(offset_row, offset_col)`` where + ``col_offset = r_pix * cos(phi)`` and ``row_offset = r_pix * sin(phi)``. + """ + if ellipse_params is None: + offset_col = r_pix * torch.cos(phi) + offset_row = r_pix * torch.sin(phi) + else: + if len(ellipse_params) != 3: + raise ValueError("ellipse_params must be (a, b, theta_deg).") + a, b, theta_deg = ellipse_params + theta = torch.deg2rad(torch.tensor(theta_deg, dtype=torch.float32, device=device)) + # Rotate into the ellipse frame, scale by a/b to undo the distortion, + # then rotate back so sampling follows the true circular rings + alpha = phi - theta + u = (a / b) * r_pix * torch.cos(alpha) + v_prime = r_pix * torch.sin(alpha) + cos_t = torch.cos(theta) + sin_t = torch.sin(theta) + offset_col = u * cos_t - v_prime * sin_t + offset_row = u * sin_t + v_prime * cos_t + return offset_row, offset_col + + +def _build_polar_sampling_offsets( + ellipse_params: tuple[float, float, float] | None, + num_annular_bins: int, + radial_min: float, + radial_max_eff: float, + radial_step: float, + two_fold_rotation_symmetry: bool, + device: str = "cpu", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build origin-independent Cartesian (row, col) offsets for a polar + sampling grid. + + Returns ``(offset_row, offset_col, phi_bins, radial_bins)`` where + ``offset_row`` and ``offset_col`` have shape ``(n_phi, n_r)`` and + represent pixel displacements from an arbitrary origin. + """ + if radial_step <= 0: + raise ValueError(f"Got radial_step = {radial_step}. radial_step must be > 0.") + if num_annular_bins < 1: + raise ValueError("num_annular_bins must be >= 1.") + + radial_bins = torch.arange( + radial_min, radial_max_eff, radial_step, dtype=torch.float32, device=device + ) + if radial_bins.numel() == 0: + radial_bins = torch.tensor([radial_min], dtype=torch.float32, device=device) + phi_range = torch.pi if two_fold_rotation_symmetry else 2.0 * torch.pi + # Drop the last endpoint because 0 and 2pi (or pi) are the same angle + phi_bins = torch.linspace( + 0.0, phi_range, num_annular_bins + 1, dtype=torch.float32, device=device + )[:-1] + phi_grid, r_pix_grid = torch.meshgrid(phi_bins, radial_bins, indexing="ij") + # Compute offsets relative to origin (0,0) so they can be reused + # for any candidate origin by simple translation + offset_row, offset_col = _polar_to_cartesian_offsets( + phi_grid, r_pix_grid, ellipse_params, device + ) + return offset_row, offset_col, phi_bins, radial_bins + + +def _build_candidate_grids( + base_col_norm: torch.Tensor, + base_row_norm: torch.Tensor, + center_row: int, + center_col: int, + margin: int, + n_row: int, + n_col: int, + col_norm_scale: float, + row_norm_scale: float, + device: str = "cpu", + step: int = 1, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build a batch of normalized sampling grids, one per candidate origin + pixel in a search window around (center_row, center_col). Candidates are + produced in a single batched tensor so that they can be evaluated + simultaneously by ``_angular_std_scores``. + + Parameters + ---------- + base_col_norm, base_row_norm : torch.Tensor of shape (n_phi, n_r) + Polar sampling offsets, already expressed in ``grid_sample``'s + normalized [-1, 1] coordinates, relative to origin (0, 0) + center_row, center_col : int + Center of the candidate search window + margin : int + Half-width of the search window in pixels + n_row, n_col : int + Diffraction-pattern image dimensions + col_norm_scale, row_norm_scale : float + Conversion factor from an offset in pixel units to the equivalent + offset in ``grid_sample``'s normalized coordinates + + Returns + ------- + row_flat, col_flat : torch.Tensor of shape (N,) + Candidate origin positions + grids : torch.Tensor of shape (N, n_phi, n_r, 2) + Stacked sampling grids ready for ``F.grid_sample`` (ordered ``(col, row)`` ) + """ + # Enumerate all pixel positions in the search window, clamped to image bounds + rows = torch.arange( + max(0, center_row - margin), + min(n_row, center_row + margin + 1), + step, + dtype=torch.long, + device=device, + ) + cols = torch.arange( + max(0, center_col - margin), + min(n_col, center_col + margin + 1), + step, + dtype=torch.long, + device=device, + ) + row_grid, col_grid = torch.meshgrid(rows, cols, indexing="ij") + row_flat, col_flat = row_grid.reshape(-1), col_grid.reshape(-1) + # Shift the pre-computed polar offsets to each candidate origin, + # converting to grid_sample's [-1, 1] normalized coordinates + grid_col = ( + base_col_norm.unsqueeze(0) + (col_flat.float() * col_norm_scale - 1.0)[:, None, None] + ) + grid_row = ( + base_row_norm.unsqueeze(0) + (row_flat.float() * row_norm_scale - 1.0)[:, None, None] + ) + # grid_sample requires (col, row) ordering in the last dim + grids = torch.stack([grid_col, grid_row], dim=-1) # (N, n_phi, n_r, 2) + return row_flat, col_flat, grids + + +def _angular_std_scores( + dp_batch: torch.Tensor, + grids: torch.Tensor, + min_r_idx: int, + max_r_idx: int, + ring_weights: torch.Tensor | None = None, +) -> torch.Tensor: + """Score candidate origins by angular std over a mid-radius band. + Lower scores indicate better centering. Optional ``ring_weights`` + applies a per-ring weighting before summation.""" + n = grids.shape[0] + # Sample the diffraction pattern at each candidate's polar grid positions + polars = F.grid_sample( + dp_batch.expand(n, -1, -1, -1), + grids, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + # A correctly centered pattern has uniform intensity along each ring, + # so the angular std is minimized at the true center. Normalize by + # mean intensity so candidates whose polar grid samples only background + # do not win on absolute std alone (relevant when search windows are + # wide enough to include centers far from the direct beam). + region = polars.squeeze(1)[:, :, min_r_idx:max_r_idx] + std_r = _angular_std(region, dim=1) + mean_r = region.mean(dim=1) + if ring_weights is not None: + std_r = std_r * ring_weights + mean_r = mean_r * ring_weights + return std_r.sum(dim=1) / (mean_r.sum(dim=1) + 1e-6) + + +def _angular_std(region: torch.Tensor, dim: int) -> torch.Tensor: + """``region.std(dim=dim)``, with a faster route on CPU.""" + if region.device.type != "cpu": + return region.std(dim=dim) + region = region.double() + centred = region - region.mean(dim=dim, keepdim=True) + return (centred.square().sum(dim=dim) / (region.shape[dim] - 1)).sqrt().float() + + +def _quadratic_subpixel_offset(patch: torch.Tensor) -> torch.Tensor: + """Sub-pixel location of the minimum of a 3x3 score patch. + + Fits a 2D quadratic surface + s(u, v) = a + b u + c v + d u^2 + e v^2 + f u v + (u = row offset, v = col offset, both over {-1, 0, 1}) to each patch by + least squares, then returns the vertex (the point where both partial + derivatives vanish): + [2d f ] [dr] [-b] + [f 2e] [dc] = [-c] + The fit depends only on the fixed 3x3 offset positions, so the + least-squares matrix (``fit_matrix``) is identical for every patch. + + Parameters + ---------- + patch + (N, 3, 3) score patches. ``patch[:, i, j]`` is the score at row + offset ``i - 1`` and col offset ``j - 1``. + + Returns + ------- + offset + (N, 2) sub-pixel ``(drow, dcol)`` offsets in pixels, each clamped to + [-1, 1]. Rows whose fit is not a clean minimum (Hessian not positive + definite, or any non-finite score) get ``(0, 0)``. + """ + device = patch.device + scores_flat = patch.reshape(patch.shape[0], 9).to(torch.float64) # (N, 9) + # The 9 fixed offset positions of the 3x3 patch, flattened row-major to + # match scores_flat. + grid = torch.tensor([-1.0, 0.0, 1.0], dtype=torch.float64, device=device) + uu, vv = torch.meshgrid(grid, grid, indexing="ij") + u, v = uu.reshape(9), vv.reshape(9) + # Design matrix: each row evaluates the quadratic basis + # [1, u, v, u^2, v^2, u v] at one offset position. + basis_matrix = torch.stack([torch.ones_like(u), u, v, u * u, v * v, u * v], dim=1) # (9, 6) + # Least-squares pseudo-inverse, same for every patch: coeffs = fit_matrix @ scores. + fit_matrix = torch.linalg.pinv(basis_matrix) # (6, 9) + # fit coeffs + a, b, c, d, e, f = (scores_flat @ fit_matrix.T).unbind(dim=1) # each (N,) + # find minimum + det = 4.0 * d * e - f * f # determinant of the 2x2 curvature (Hessian) matrix + # Accept only clean minima (Hessian positive definite) with finite scores. + valid = (2.0 * d > 0) & (det > 1e-12) & torch.isfinite(scores_flat).all(dim=1) + det_safe = torch.where(valid, det, torch.ones_like(det)) + drow = torch.where(valid, (f * c - 2.0 * e * b) / det_safe, torch.zeros_like(det)) + dcol = torch.where(valid, (f * b - 2.0 * d * c) / det_safe, torch.zeros_like(det)) + return torch.stack([drow.clamp(-1.0, 1.0), dcol.clamp(-1.0, 1.0)], dim=1) + + + +def find_origin_angular_descent( + data: Dataset4dstem, + *, + ellipse_params: tuple[float, float, float] | None = None, + radial_min: float = 0.0, + radial_max: float | None = None, + n_phi: int = 120, + radial_step: float = 1.0, + kpow: float = 0.0, + device: str = "cpu", +) -> NDArray: + """Margin-free diffraction-center finding by COM-anchored local descent. + + Anchors every pattern at the intensity centroid and walks downhill, then + fits a 2D quadratic for the sub-pixel center. Robust on small detectors. + + Parameters + ---------- + data : Dataset4dstem + A 4D-STEM dataset (or a 2D dataset, treated as a 1x1 scan). + ellipse_params : tuple or None + ``(a, b, theta_deg)`` to sample along elliptical rings. Origin finding stays + accurate without it (the center is located by the pattern's inversion + symmetry, independent of ring shape); supplying it deepens the score minimum + and is a mild refinement on strongly elliptical data. ``None`` = circular. + radial_min, radial_max : float + Inner / outer radius (px) of the scoring annulus. Set ``radial_min`` above + the central beam (or any saturated core). ``radial_max=None`` defaults to + ``min(n_row, n_col) // 2 - 2``. + n_phi : int + Number of angular samples per ring. + radial_step : float + Radial sampling step (px). + kpow : float + Up-weight each ring's contribution by ``k**kpow`` (k == radius). ``0`` is the + plain ratio-of-sums objective; ``1``/``2`` emphasize high-k. + device : str + Torch device. + + Returns + ------- + origin_array : np.ndarray + Shape ``(scan_row, scan_col, 2)`` of ``(row, col)`` origins in pixels. + + See Also + -------- + find_origin_angular_grid : global angular-variance grid search; more robust to a + poor initial guess. + """ + if data.ndim == 2: + scan_row, scan_col, n_row, n_col = 1, 1, *data.shape + elif data.ndim == 4: + scan_row, scan_col, n_row, n_col = data.shape + else: + raise ValueError( + f"Got array with shape {data.shape}. " + "To use find_origin_angular_descent, pass a 2D or 4D-STEM dataset." + ) + if radial_max is None: + radial_max = float(min(n_row, n_col) // 2 - 2) + n_radial = max(4, int(round((radial_max - radial_min) / radial_step)) + 1) + data_tensor = ( + torch.from_numpy(np.ascontiguousarray(data.array)) + if data.array is not None + else data.tensor + ) + array_tensor = data_tensor.to(device=device, dtype=torch.float32) + if array_tensor.ndim == 2: + array_tensor = array_tensor[None, None] + patterns = array_tensor.reshape(-1, n_row, n_col) + n_patterns = patterns.shape[0] + offset_row, offset_col, ring_weights = _local_sampling( + radial_min, radial_max, n_phi, n_radial, kpow, ellipse_params, device + ) + # Global anchor: descend on the high-SNR mean DP (a batch of one), seeded at its COM. + mean_pattern = array_tensor.mean(dim=(0, 1)) + global_origin = _descend_batched( + mean_pattern[None], torch.round(_com_anchor(mean_pattern))[None], + offset_row, offset_col, ring_weights, n_phi, n_radial, device, + )[0] + # Per-position: start every DP at the integer global anchor and descend (tracks descan). + start_centers = torch.round(global_origin)[None].expand(n_patterns, 2).clone() + origins = _descend_batched( + patterns, start_centers, offset_row, offset_col, ring_weights, n_phi, n_radial, device, + ) + return origins.reshape(scan_row, scan_col, 2).cpu().numpy() + + +# --- helpers for find_origin_angular_descent ------------------------------------- +# The 8 neighbor directions (axial + diagonal) for the pattern search. +_NEIGHBOR_STEPS_8 = [[1.0, 0], [-1, 0], [0, 1.0], [0, -1], [1, 1.0], [1, -1], [-1, 1], [-1, -1]] +# The 9 offsets of a 3x3 patch, row-major: index k = i*3 + j -> (row i-1, col j-1); +# the patch center is k = 4. +_PATCH_OFFSETS_3X3 = [[-1, -1], [-1, 0], [-1, 1], [0, -1], [0, 0], [0, 1], [1, -1], [1, 0], [1, 1]] + + +def _com_anchor(pattern: torch.Tensor) -> torch.Tensor: + """Rough center guess from the intensity-weighted centroid of the DP.""" + n_row, n_col = pattern.shape + clipped = pattern.clamp(min=0) + total = clipped.sum() + 1e-9 + rows = torch.arange(n_row, device=pattern.device, dtype=torch.float32) + cols = torch.arange(n_col, device=pattern.device, dtype=torch.float32) + center_row = (rows[:, None] * clipped).sum() / total + center_col = (cols[None, :] * clipped).sum() / total + return torch.tensor([center_row.item(), center_col.item()], device=pattern.device) + + +def _local_sampling(radial_min, radial_max, n_phi, n_radial, kpow, ellipse_params, device): + """Polar-sampling template (row/col offset per (phi, radius)) + per-ring k-weights. + + With ``ellipse_params=(a, b, theta_deg)`` the template follows the elliptical rings. ``None`` samples on + circles. Uses the same ``_polar_to_cartesian_offsets`` transform as the grid finder.""" + phi = torch.linspace(0, 2 * np.pi, n_phi + 1, device=device)[:-1] + radii = torch.linspace(radial_min, radial_max, n_radial, device=device) + phi_grid, radius_grid = torch.meshgrid(phi, radii, indexing="ij") # (n_phi, n_radial) + offset_row, offset_col = _polar_to_cartesian_offsets( + phi_grid, radius_grid, ellipse_params, device + ) + ring_weights = radii**kpow # (n_radial,) + return offset_row, offset_col, ring_weights + + +def _local_polar_score(polar_values, valid_mask, n_phi, ring_weights, min_valid_frac): + """Angular-uniformity score from sampled polar values.""" + n_valid = valid_mask.sum(dim=-2).clamp(min=1) # (..., n_radial) + ring_mean = (polar_values * valid_mask).sum(dim=-2) / n_valid + ring_var = (((polar_values - ring_mean.unsqueeze(-2)) ** 2) * valid_mask).sum(dim=-2) / n_valid + ring_std = ring_var.sqrt() + ring_usable = valid_mask.sum(dim=-2) >= (min_valid_frac * n_phi) + weights = ring_weights * ring_usable + score = (weights * ring_std).sum(dim=-1) / ((weights * ring_mean.abs()).sum(dim=-1) + 1e-6) + return score, ring_usable + + +def _local_score_pairs(patterns, pattern_index, centers, offset_row, offset_col, + ring_weights, n_phi, device, min_valid_frac=0.5, chunk=4096): + """Score K (pattern, center) pairs: patterns[pattern_index[k]] scored at centers[k]. + + patterns:(M,H,W), pattern_index:(K,), centers:(K,2) -> (K,) angular-uniformity scores.""" + _, n_row, n_col = patterns.shape + ones_image = torch.ones(1, 1, n_row, n_col, device=device) + scores = torch.empty(centers.shape[0], device=device) + for start in range(0, centers.shape[0], chunk): + index = pattern_index[start:start + chunk] + n_chunk = index.shape[0] + center_row = centers[start:start + chunk, 0][:, None, None] + center_col = centers[start:start + chunk, 1][:, None, None] + sample_grid = torch.stack( + [2.0 * (center_col + offset_col[None]) / (n_col - 1) - 1.0, + 2.0 * (center_row + offset_row[None]) / (n_row - 1) - 1.0], dim=-1 + ) + polar_values = F.grid_sample( + patterns[index][:, None], sample_grid, mode="bilinear", + padding_mode="zeros", align_corners=True + )[:, 0] + valid_mask = F.grid_sample( + ones_image.expand(n_chunk, 1, n_row, n_col), sample_grid, mode="bilinear", + padding_mode="zeros", align_corners=True + )[:, 0] > 0.999 + scores[start:start + n_chunk], _ = _local_polar_score( + polar_values, valid_mask, n_phi, ring_weights, min_valid_frac + ) + return scores + + +def _descend_batched(patterns, anchors, offset_row, offset_col, ring_weights, n_phi, + n_radial, device, schedule=(4.0, 2.0, 1.0), sweeps=2): + """Batched coarse descent (shared step schedule, to 1 px) + quadratic sub-pixel. + + Walks each pattern downhill on the angular-uniformity score from its integer + ``anchors`` center (Hooke-Jeeves 8-neighbor search, step halving over ``schedule``), + then fits a 2D quadratic to the local 3x3 for the sub-pixel vertex. Used for both + the global mean-DP anchor (a batch of one) and the per-position pass. + + patterns : (M, H, W); anchors : (M, 2) integer-valued start centers. + Returns (M, 2) sub-pixel (row, col) origins. + """ + n_patterns = patterns.shape[0] + pattern_ids = torch.arange(n_patterns, device=device) + neighbor_steps = torch.tensor(_NEIGHBOR_STEPS_8, device=device) + patch_offsets = torch.tensor(_PATCH_OFFSETS_3X3, dtype=torch.float32, device=device) + pattern_ids_per_neighbor = pattern_ids.repeat_interleave(8) + pattern_ids_per_patch = pattern_ids.repeat_interleave(9) + center = anchors.clone() + + def score_at(pattern_index, centers): + return _local_score_pairs(patterns, pattern_index, centers, offset_row, offset_col, + ring_weights, n_phi, device) + + best_score = score_at(pattern_ids, center) + # Coarse descent: at each step size, move to the best-improving 8-neighbor; when no + # neighbor improves, halve the step. Brackets the minimum to ~1 px. + for step in schedule: + for _ in range(sweeps): + neighbors = (center[:, None, :] + step * neighbor_steps[None]).reshape(n_patterns * 8, 2) + neighbor_scores = score_at(pattern_ids_per_neighbor, neighbors).reshape(n_patterns, 8) + best_neighbor = neighbor_scores.argmin(dim=1) + best_neighbor_score = neighbor_scores.gather(1, best_neighbor[:, None]).squeeze(1) + improved = best_neighbor_score < best_score - 1e-12 + best_neighbor_center = neighbors.reshape(n_patterns, 8, 2)[pattern_ids, best_neighbor] + center = torch.where(improved[:, None], best_neighbor_center, center) + best_score = torch.where(improved, best_neighbor_score, best_score) + # Re-center each 3x3 patch on its integer minimum so the quadratic brackets it. + patch_centers = (center[:, None, :] + patch_offsets[None]).reshape(n_patterns * 9, 2) + patch_scores = score_at(pattern_ids_per_patch, patch_centers).reshape(n_patterns, 9) + center = patch_centers.reshape(n_patterns, 9, 2)[pattern_ids, patch_scores.argmin(dim=1)] + # Sub-pixel vertex from a 2D-quadratic fit to the final 3x3 of scores. + patch_centers = (center[:, None, :] + patch_offsets[None]).reshape(n_patterns * 9, 2) + score_patch = score_at(pattern_ids_per_patch, patch_centers).reshape(n_patterns, 3, 3) + subpixel_offset = _quadratic_subpixel_offset(score_patch).to(torch.float32) + return center + subpixel_offset diff --git a/tests/diffraction/test_ellipse_fit.py b/tests/diffraction/test_ellipse_fit.py new file mode 100644 index 00000000..c32aaf56 --- /dev/null +++ b/tests/diffraction/test_ellipse_fit.py @@ -0,0 +1,210 @@ +"""Tests for elliptical distortion fitting from amorphous ring patterns.""" + +import numpy as np +import pytest +import torch + +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.diffraction.ellipticity import ( + _amorphous_ring_model, + convert_ellipse_params, + convert_ellipse_params_r, + fit_elliptical_distortion, +) +from quantem.diffraction.polar_transform import polar_transform + + +def _make_synthetic_dataset( + n_row: int = 256, + n_col: int = 256, + center: tuple[float, float] = (128.0, 128.0), + a: float = 80.0, + b: float = 60.0, + theta_deg: float = 30.0, + I0: float = 1000.0, + I1: float = 50.0, + sigma0: float = 20.0, + sigma1: float = 10.0, + sigma2: float = 14.0, + c_bkgd: float = 2.0, + noise: bool = True, +) -> Dataset4dstem: + """Generate a synthetic 4D-STEM dataset (1x1 scan) with known elliptical distortion.""" + row_offsets = torch.arange(n_row, dtype=torch.float32) - center[0] + col_offsets = torch.arange(n_col, dtype=torch.float32) - center[1] + offset_row, offset_col = torch.meshgrid(row_offsets, col_offsets, indexing="ij") + + A, B, C = convert_ellipse_params_r(a, b, np.radians(theta_deg)) + params = torch.tensor( + [I0, I1, sigma0, sigma1, sigma2, c_bkgd, 0.0, 0.0, A, B, C], + dtype=torch.float32, + ) + dp = _amorphous_ring_model(params, offset_row.reshape(-1), offset_col.reshape(-1)) + dp = dp.reshape(n_row, n_col).numpy() + if noise: + dp = np.maximum(dp, 0.0) + rng = np.random.default_rng(42) + dp = rng.poisson(dp).astype(np.float32) + + # Wrap as a 1x1 scan 4D-STEM dataset + arr_4d = dp[np.newaxis, np.newaxis, :, :] + return Dataset4dstem.from_array(arr_4d, name="synthetic") + + +class TestConvertEllipseParams: + """Tests for the (a, b, theta) <-> (A, B, C) converters.""" + + @pytest.mark.parametrize( + "a,b,theta_deg", + [(80.0, 60.0, 30.0), (70.0, 70.0, 0.0)], # general ellipse + circle (θ-undefined branch) + ) + def test_round_trip(self, a, b, theta_deg): + """(a, b, theta) -> (A, B, C) -> (a, b, theta) should be the identity.""" + A, B, C = convert_ellipse_params_r(a, b, np.radians(theta_deg)) + a_back, b_back, theta_back = convert_ellipse_params(A, B, C) + assert abs(a_back - a) < 1e-6 * a + assert abs(b_back - b) < 1e-6 * b + if a != b: # theta is undefined for a circle + angle_err = abs(((np.degrees(theta_back) - theta_deg) + 90.0) % 180.0 - 90.0) + assert angle_err < 1e-6 + + +class TestAmorphousRingModel: + """Tests for the _amorphous_ring_model function.""" + + def test_circular_ring_peak_at_radius(self): + """Model should peak at the ring radius for a circular ring.""" + R = 50.0 + A, B, C = convert_ellipse_params_r(R, R, 0.0) + params = torch.tensor( + [0.0, 100.0, 10.0, 5.0, 5.0, 0.0, 0.0, 0.0, A, B, C], + dtype=torch.float32, + ) + # Sample along the column axis at various radii + offset_col = torch.linspace(0, 80, 200) + offset_row = torch.zeros_like(offset_col) + vals = _amorphous_ring_model(params, offset_row, offset_col) + peak_r = float(offset_col[vals.argmax()]) + assert abs(peak_r - R) < 1.0, f"Peak at {peak_r}, expected ~{R}" + + def test_background_far_from_ring(self): + """Far from the ring and center, model should return ~c_bkgd.""" + R = 50.0 + c_bkgd = 5.0 + A, B, C = convert_ellipse_params_r(R, R, 0.0) + params = torch.tensor( + [100.0, 100.0, 5.0, 3.0, 3.0, c_bkgd, 0.0, 0.0, A, B, C], + dtype=torch.float32, + ) + # Very far from center and ring + offset_row = torch.tensor([200.0]) + offset_col = torch.tensor([200.0]) + val = float(_amorphous_ring_model(params, offset_row, offset_col)) + assert abs(val - c_bkgd) < 1.0, f"Expected ~{c_bkgd}, got {val}" + + +class TestFitEllipticalDistortion: + """Tests for fit_elliptical_distortion.""" + + def test_recover_known_ellipse(self): + """Fit should recover known ellipse params from a synthetic pattern.""" + a_true, b_true, theta_true = 80.0, 60.0, 30.0 + center_true = (128.0, 128.0) + ds = _make_synthetic_dataset(a=a_true, b=b_true, theta_deg=theta_true, center=center_true) + # fit_radii brackets the ring (mean_r~70) well outside the central beam + result = fit_elliptical_distortion( + ds, + center=center_true, + fit_radii=(45.0, 100.0), + ) + a_fit, b_fit, theta_fit = result["ellipse_params"] + # Check semiaxes within 10% relative error + assert abs(a_fit - a_true) / a_true < 0.10, f"a: got {a_fit:.2f}, expected {a_true}" + assert abs(b_fit - b_true) / b_true < 0.10, f"b: got {b_fit:.2f}, expected {b_true}" + # Angle within 5 degrees + angle_err = abs(theta_fit - theta_true) + angle_err = min(angle_err, 180.0 - angle_err) # handle wrapping + assert angle_err < 5.0, f"theta: got {theta_fit:.1f}, expected {theta_true}" + + def test_recover_circular(self): + """For a circular pattern (a==b), fitted a and b should be similar.""" + R = 70.0 + ds = _make_synthetic_dataset(a=R, b=R, theta_deg=0.0, noise=False) + result = fit_elliptical_distortion( + ds, + center=(128.0, 128.0), + fit_radii=(45.0, 100.0), + ) + a_fit, b_fit, _ = result["ellipse_params"] + ratio = a_fit / b_fit + assert 0.95 < ratio < 1.05, f"Expected a/b ~ 1 for circular, got {ratio:.3f}" + + def test_refined_center(self): + """Fit should refine center when it's slightly off.""" + center_true = (130.0, 126.0) + ds = _make_synthetic_dataset(center=center_true, noise=False) + # Give a slightly wrong center + result = fit_elliptical_distortion( + ds, + center=(128.0, 128.0), + fit_radii=(45.0, 100.0), + ) + row_fit, col_fit = result["center"] + assert abs(row_fit - center_true[0]) < 3.0, ( + f"row: got {row_fit:.1f}, expected {center_true[0]}" + ) + assert abs(col_fit - center_true[1]) < 3.0, ( + f"col: got {col_fit:.1f}, expected {center_true[1]}" + ) + + def test_empty_annulus_raises(self): + """Should raise if no pixels fall in the annulus.""" + ds = _make_synthetic_dataset() + with pytest.raises(ValueError, match="No pixels"): + fit_elliptical_distortion( + ds, + center=(128.0, 128.0), + fit_radii=(200.0, 210.0), # outside image + ) + + @pytest.mark.parametrize("write", [True, False]) + def test_writes_ellipticity_metadata(self, write): + """metadata['ellipticity'] is stamped iff write_metadata is True.""" + ds = _make_synthetic_dataset(noise=False) + assert ds.metadata["ellipticity"] is None + result = fit_elliptical_distortion( + ds, + center=(128.0, 128.0), + fit_radii=(45.0, 100.0), + max_iter=10, + write_metadata=write, + ) + expected = result["ellipse_params"] if write else None + assert ds.metadata["ellipticity"] == expected + + def test_accepts_tensor_backed_input(self): + """Tensor-backed datasets should fit identically to numpy-backed ones.""" + ds_np = _make_synthetic_dataset(noise=False) + ds_t = Dataset4dstem.from_tensor( + torch.from_numpy(ds_np.array.copy()), name="synthetic_tensor" + ) + kwargs = dict(center=(128.0, 128.0), fit_radii=(45.0, 100.0)) + result_np = fit_elliptical_distortion(ds_np, **kwargs) + result_t = fit_elliptical_distortion(ds_t, **kwargs) + np.testing.assert_allclose( + result_t["ellipse_params"], result_np["ellipse_params"], rtol=1e-4 + ) + + +def test_ellipse_correction_pipeline(): + """fit_elliptical_distortion -> polar_transform: corrected ring is uniform across phi.""" + ds = _make_synthetic_dataset(a=80.0, b=60.0, theta_deg=30.0, noise=False) + fit = fit_elliptical_distortion(ds, center=(128.0, 128.0), fit_radii=(45.0, 100.0)) + polar = polar_transform(ds, ellipse_params=fit["ellipse_params"]) + # In the corrected polar the ring sits at a constant radius, so intensity + # at the peak radius should be ~uniform across phi. + arr = polar.numpy()[0, 0] # (n_phi, n_r), single 1x1 scan position + peak_r = int(arr.mean(axis=0).argmax()) + ring_at_phi = arr[:, peak_r] + cv = float(ring_at_phi.std() / ring_at_phi.mean()) + assert cv < 0.1, f"Ring varies across phi (cv={cv:.3f}); correction not applied." diff --git a/tests/diffraction/test_polar.py b/tests/diffraction/test_polar.py new file mode 100644 index 00000000..7f8e7716 --- /dev/null +++ b/tests/diffraction/test_polar.py @@ -0,0 +1,726 @@ +import numpy as np +import pytest +import torch + +from quantem.core.datastructures.dataset2d import Dataset2d +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.datastructures.polar4dstem import Polar4dstem +from quantem.diffraction.polar import PairDistributionFunction +from quantem.diffraction.polar_transform import ( + find_origin_angular_descent, + find_origin_angular_grid, + polar_transform, +) + +# ============================================================================ +# Fixtures +# ============================================================================ + + +@pytest.fixture +def synthetic_diffraction_pattern(): + """Create a synthetic diffraction pattern with concentric rings.""" + ny, nx = 256, 256 + y, x = np.ogrid[:ny, :nx] + cy, cx = (ny - 1) / 2.0, (nx - 1) / 2.0 + + # Create rings with Gaussian profiles at specific radii + pattern = np.zeros((ny, nx), dtype=np.float32) + ring_radii = [10, 20, 30, 40] + r = np.sqrt((y - cy) ** 2 + (x - cx) ** 2) + for radius in ring_radii: + pattern += 100 * np.exp(-((r - radius) ** 2) / (2 * 2**2)) + # central beam + pattern += 1000 * np.exp(-(r**2) / (2 * 3**2)) + # noise + rng = np.random.default_rng(42) + pattern += rng.poisson(5, size=(ny, nx)) + + return pattern.astype(np.float32) + + +@pytest.fixture +def synthetic_4dstem_dataset(synthetic_diffraction_pattern): + """Create a synthetic 4D-STEM dataset with 3x3 scan.""" + scan_y, scan_x = 3, 3 + ny, nx = synthetic_diffraction_pattern.shape + + array_4d = np.zeros((scan_y, scan_x, ny, nx), dtype=np.float32) + for iy in range(scan_y): + for ix in range(scan_x): + # Add slight variations + rng = np.random.default_rng(42 + iy * scan_x + ix) + variation = 1.0 + 0.1 * rng.standard_normal() + array_4d[iy, ix] = synthetic_diffraction_pattern * variation + + return Dataset4dstem.from_array( + array=array_4d, + name="test_4dstem", + origin=(0, 0, 0, 0), + sampling=(1.0, 1.0, 0.015, 0.015), + units=["nm", "nm", "1/Angstrom", "1/Angstrom"], + signal_units="counts", + ) + + +@pytest.fixture +def synthetic_dataset2d(synthetic_diffraction_pattern): + """Create a synthetic 2D diffraction dataset.""" + return Dataset2d.from_array( + array=synthetic_diffraction_pattern, + name="test_2d_diffraction", + origin=(0, 0), + sampling=(0.015, 0.015), + units=["1/Angstrom", "1/Angstrom"], + signal_units="counts", + ) + + +# ============================================================================ +# Test PairDistributionFunction Construction +# ============================================================================ + + +class TestPairDistributionFunctionConstruction: + """Test PairDistributionFunction initialization from various input types.""" + + def test_from_data_with_dataset4dstem(self, synthetic_4dstem_dataset): + """Test construction from a Dataset4dstem object.""" + pdf = PairDistributionFunction.from_data( + synthetic_4dstem_dataset, + find_origin=False, + ) + assert isinstance(pdf.polar, Polar4dstem) + assert pdf.input_data is synthetic_4dstem_dataset + assert pdf.polar.shape[0] == 3 # scan_y + assert pdf.polar.shape[1] == 3 # scan_x + assert pdf.polar.shape[2] == 180 # num_annular_bins + + def test_direct_init_without_token_raises(self, synthetic_dataset2d): + """Test that direct __init__ without token raises RuntimeError.""" + pdf_valid = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + with pytest.raises(RuntimeError, match="Use PairDistributionFunction.from_data"): + PairDistributionFunction(polar=pdf_valid.polar, device="cpu") + + def test_from_data_origin_method(self, synthetic_4dstem_dataset): + """from_data routes origin_method to the two finders and rejects bad values.""" + for method in ("grid", "descent"): + pdf = PairDistributionFunction.from_data( + synthetic_4dstem_dataset, find_origin=True, origin_method=method + ) + assert pdf.polar.shape[:2] == (3, 3) + with pytest.raises(ValueError, match="origin_method"): + PairDistributionFunction.from_data( + synthetic_4dstem_dataset, find_origin=True, origin_method="bogus" + ) + + +@pytest.mark.parametrize("finder", [find_origin_angular_grid, find_origin_angular_descent]) +class TestOriginFinding: + """Shared coverage for both origin finders (``find_origin_angular_grid`` and ``find_origin_angular_descent``). Each test runs against both.""" + + @staticmethod + def _ring_pattern(ny, nx, cy, cx, radii=(12, 24, 36, 48), beam_sigma=3.0): + """Concentric circular rings + a central beam, centered at (cy, cx).""" + y, x = np.ogrid[:ny, :nx] + r = np.sqrt((y - cy) ** 2 + (x - cx) ** 2) + p = np.zeros((ny, nx), dtype=np.float32) + for radius in radii: + p += 100.0 * np.exp(-((r - radius) ** 2) / (2 * 2.0**2)) + p += 1000.0 * np.exp(-(r**2) / (2 * beam_sigma**2)) + return p.astype(np.float32) + + @staticmethod + def _wrap(arr): + return Dataset4dstem.from_array( + array=arr, + name="origin_test", + origin=(0, 0, 0, 0), + sampling=(1.0, 1.0, 1.0, 1.0), + units=["nm", "nm", "1/Angstrom", "1/Angstrom"], + signal_units="counts", + ) + + def test_recovers_center(self, finder, synthetic_4dstem_dataset): + """Recovers the known center of the shared 3x3 fixture.""" + origins = finder(synthetic_4dstem_dataset, radial_min=4, radial_max=50, device="cpu") + assert origins.shape == (3, 3, 2) + assert np.all(np.abs(origins - 127.5) < 0.5) + + def test_subpixel(self, finder): + """Recovers distinct fractional centers to sub-pixel precision (also + exercises per-position batching).""" + ny = nx = 128 + true_centers = [(63.3, 64.7), (64.6, 63.4), (62.8, 65.2), (65.1, 62.9)] + arr = np.stack([self._ring_pattern(ny, nx, cy, cx) for cy, cx in true_centers]).reshape( + 2, 2, ny, nx + ) + origins = finder(self._wrap(arr), radial_min=4, radial_max=54, device="cpu") + assert origins.shape == (2, 2, 2) + for k, (cy, cx) in enumerate(true_centers): + row, col = origins[k // 2, k % 2] + err = float(np.hypot(row - cy, col - cx)) + assert err < 0.1, f"sub-pixel origin off by {err:.3f} px at position {k}" + + def test_small_detector(self, finder): + """Works on a small (64 px) detector -- the regime where the grid finder's + fixed search margin used to collapse the search annulus.""" + ny = nx = 64 + cy, cx = 30.4, 30.6 + arr = self._ring_pattern(ny, nx, cy, cx, radii=(8, 16, 24))[None, None] + origins = finder(self._wrap(arr), radial_min=3, radial_max=28, device="cpu") + assert origins.shape == (1, 1, 2) + assert float(np.hypot(origins[0, 0, 0] - cy, origins[0, 0, 1] - cx)) < 0.3 + + def test_accepts_tensor_backed_input(self, finder): + """Runs on a tensor-backed dataset (whose .array is None).""" + arr = self._ring_pattern(128, 128, 64.0, 64.0)[None, None] + ds_t = Dataset4dstem.from_tensor(torch.from_numpy(arr)) + assert ds_t.array is None # tensor-backed + origins = finder(ds_t, radial_min=4, radial_max=54, device="cpu") + assert origins.shape == (1, 1, 2) + assert float(np.hypot(origins[0, 0, 0] - 64.0, origins[0, 0, 1] - 64.0)) < 0.3 + + def test_ellipse_params(self, finder): + """ellipse_params is accepted and the center is recovered on an elliptical + ring when the correct correction is supplied.""" + ny = nx = 128 + cy, cx, theta_deg = 64.3, 63.7, 25.0 + theta = np.radians(theta_deg) + a, b, ring_r = 40 * np.sqrt(1.6), 40 / np.sqrt(1.6), 40.0 + y, x = np.ogrid[:ny, :nx] + u_col, u_row = x - cx, y - cy + u_a = u_col * np.cos(theta) + u_row * np.sin(theta) + u_b = -u_col * np.sin(theta) + u_row * np.cos(theta) + r_ell = np.sqrt((u_a / a) ** 2 + (u_b / b) ** 2) + p = 300.0 * np.exp(-(((r_ell - 1) * ring_r) ** 2) / (2 * 4.0**2)) + p += 1000.0 * np.exp(-(u_col**2 + u_row**2) / (2 * 3.0**2)) + ds = self._wrap(p.astype(np.float32)[None, None]) + origins = finder( + ds, ellipse_params=(a, b, theta_deg), radial_min=8, radial_max=58, device="cpu" + ) + err = float(np.hypot(origins[0, 0, 0] - cy, origins[0, 0, 1] - cx)) + assert err < 0.3, f"elliptical origin off by {err:.3f} px" + + +# ============================================================================ +# Test Polar Transform +# ============================================================================ + + +class TestPolarTransform: + """Test polar coordinate transformation.""" + + def test_polar_transform_basic(self, synthetic_4dstem_dataset): + """Test basic polar transformation.""" + polar = polar_transform(synthetic_4dstem_dataset) + assert isinstance(polar, Polar4dstem) + assert polar.shape[0] == 3 # scan_y + assert polar.shape[1] == 3 # scan_x + assert polar.shape[2] == 180 # num_annular_bins + assert polar.shape[3] > 0 # radial bins + + def test_polar_transform_single_origin(self, synthetic_4dstem_dataset): + """Test polar transformation with single origin broadcast to all positions.""" + origin = np.array([128.0, 128.0]) + polar = polar_transform( + synthetic_4dstem_dataset, + origin_array=origin, + ) + assert isinstance(polar, Polar4dstem) + + def test_polar_transform_radial_range(self, synthetic_4dstem_dataset): + """Test polar transformation with custom radial range.""" + polar = polar_transform( + synthetic_4dstem_dataset, + radial_min=5.0, + radial_max=50.0, + radial_step=2.0, + ) + assert isinstance(polar, Polar4dstem) + # Check that radial dimension matches expected size + expected_n_r = int(np.ceil((50.0 - 5.0) / 2.0)) + assert polar.shape[3] == expected_n_r + + def test_polar_transform_scan_pos(self, synthetic_4dstem_dataset): + """Test polar transformation for a single scan position.""" + polar_2d = polar_transform( + synthetic_4dstem_dataset, + scan_pos=(0, 0), + ) + # should return 2D tensor (phi, r) + assert polar_2d.ndim == 2 + assert polar_2d.shape[0] == 180 # num_annular_bins + + +# ============================================================================ +# Test Radial Mean Calculation +# ============================================================================ + + +class TestRadialMeanCalculation: + """Test radial mean intensity calculation.""" + + def test_calculate_radial_mean_with_mask(self, synthetic_4dstem_dataset): + """Test radial mean calculation with real-space mask.""" + pdf = PairDistributionFunction.from_data( + synthetic_4dstem_dataset, + find_origin=False, + ) + mask = np.zeros((3, 3), dtype=bool) + mask[0:2, 0:2] = True + radial_mean = pdf.calculate_radial_mean( + mask_realspace=mask, + returnval=True, + ) + assert radial_mean is not None + + def test_pipeline_torch_native(self, synthetic_4dstem_dataset): + """The pipeline keeps the polar data on-device (tensor-backed) and + gives identical results for numpy-backed vs tensor-backed input.""" + ds_np = synthetic_4dstem_dataset + ds_t = Dataset4dstem.from_tensor(torch.from_numpy(ds_np.array.copy())) + pdf_np = PairDistributionFunction.from_data(ds_np, find_origin=False) + pdf_t = PairDistributionFunction.from_data(ds_t, find_origin=False) + # Polar data is tensor-backed regardless of input backing (no round-trip). + assert pdf_np.polar.array is None + assert pdf_t.polar.array is None + assert isinstance(pdf_t.polar.tensor, torch.Tensor) + # numpy-backed and tensor-backed inputs give identical I(k). + ik_np = pdf_np.calculate_radial_mean(returnval=True).cpu().numpy() + ik_t = pdf_t.calculate_radial_mean(returnval=True).cpu().numpy() + assert np.allclose(ik_np, ik_t, rtol=1e-5, atol=1e-6) + + +# ============================================================================ +# Test Background Fitting +# ============================================================================ + + +class TestBackgroundFitting: + """Test background fitting.""" + + def test_fit_bg_basic(self, synthetic_dataset2d): + """Test basic background fitting.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + Ik = pdf.calculate_radial_mean(returnval=True) + k = np.asarray(pdf.qq) + kmin, kmax = float(k.min()), float(k.max()) + bg, f = pdf.fit_bg(Ik, kmin=kmin * 0.1, kmax=kmax * 0.9) + assert bg.shape == Ik.shape + assert f.shape == Ik.shape + # Check that background is positive + assert (bg >= 0).all() + + def test_fit_bg_recovers_known_background(self, synthetic_4dstem_dataset): + """The torch fit recovers a known smooth background under a structure + modulation, on a realistic high-dynamic-range I(k) / calibrated axis.""" + pdf = PairDistributionFunction.from_data(synthetic_4dstem_dataset, find_origin=False) + nb = len(np.asarray(pdf.qq)) + pdf.polar.sampling[3] = 2.3 / (nb - 1) # realistic 1/A q-axis + k = np.asarray(pdf.qq) + rng = np.random.RandomState(0) + bg_true = ( + 0.5 + 900.0 * np.exp(-(k**2) / (2 * 0.09**2)) + 6.0 * np.exp(-(k**4) / (2 * 0.8**4)) + ) + # add 35% amplitude ripple to the true bg to simulate scattering + modulation = 1.0 + 0.35 * np.sin(2 * np.pi * k / 0.42 + 0.5) * np.exp(-k / 1.2) + Ik = np.clip(bg_true * modulation + np.abs(0.5 * rng.randn(k.size)), 1e-6, None).astype( + np.float32 + ) + # try fitting to this simulated Ik + bg, _ = pdf.fit_bg(Ik.copy()) + bg = bg.cpu().numpy() + assert (bg >= 0).all() + # The fit tracks the modulated data, so pointwise error reflects the + # ~35% structure modulation + # the recovered shape of the smooth background should still correlate strongly with the truth. + corr = float(np.corrcoef(bg, bg_true)[0, 1]) + assert corr > 0.99, f"bg shape far from truth: corr {corr:.4g}" + + def test_fit_bg_batched_matches_per_curve(self, synthetic_4dstem_dataset): + """fit_bg_batched reproduces the per-curve torch fit for a stack of + realistic radial means (the line-scan fits one background per bin).""" + pdf = PairDistributionFunction.from_data(synthetic_4dstem_dataset, find_origin=False) + nb = len(np.asarray(pdf.qq)) + pdf.polar.sampling[3] = 2.3 / (nb - 1) + k = np.asarray(pdf.qq) + # A stack of 8 distinct I(k) (central beam + envelope + a gentle, varying structure modulation). + curves = [] + for seed in range(8): + rng = np.random.RandomState(seed) + bg_true = ( + 0.5 + + (400.0 + 300.0 * rng.rand()) * np.exp(-(k**2) / (2 * 0.09**2)) + + 6.0 * np.exp(-(k**4) / (2 * 0.8**4)) + ) + mod = 1.0 + 0.2 * np.sin( + 2 * np.pi * k / (0.4 + 0.2 * rng.rand()) + rng.rand() + ) * np.exp(-k / 1.2) + curves.append((bg_true * mod).astype(np.float32)) + stack = np.stack(curves) + + bg_batched, _ = pdf.fit_bg_batched(stack, kmin=0.1) + bg_batched = bg_batched.cpu().numpy() + for i in range(stack.shape[0]): + bg_single, _ = pdf.fit_bg(stack[i].copy(), kmin=0.1) + bg_single = bg_single.cpu().numpy() + rel = np.abs(bg_batched[i] - bg_single) / np.clip(np.abs(bg_single), 1e-6, None) + assert rel.max() < 1e-3, f"batched vs per-curve bg mismatch (row {i}): {rel.max():.4g}" + + +# ============================================================================ +# Test PDF Calculation +# ============================================================================ + + +class TestPDFCalculation: + """Test the PDF calculation pipeline.""" + + def test_calculate_Gr_with_bandpass(self, synthetic_dataset2d): + """Test PDF calculation with bandpass filtering.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + pdf.calculate_Gr( + k_min_fit=0.1, + k_max_fit=2.0, + k_lowpass=0.02, + k_highpass=0.001, + ) + assert pdf.reduced_pdf is not None + + def test_calculate_Gr_with_mask(self, synthetic_4dstem_dataset): + """Test PDF calculation with real-space mask.""" + pdf = PairDistributionFunction.from_data( + synthetic_4dstem_dataset, + find_origin=False, + ) + mask = np.zeros((3, 3), dtype=bool) + mask[0:2, 0:2] = True + pdf.calculate_Gr( + k_min_fit=0.1, + k_max_fit=2.0, + mask_realspace=mask, + ) + assert pdf.reduced_pdf is not None + + def test_calculate_gr_requires_Gr(self, synthetic_dataset2d): + """Test that calculate_gr raises if calculate_Gr has not been run.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + with pytest.raises(RuntimeError, match="Reduced PDF not computed"): + pdf.calculate_gr(density=0.05) + + def test_calculate_gr_requires_density(self, synthetic_dataset2d): + """The density is an input: calculate_gr raises without one and never estimates it.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + with pytest.raises(ValueError, match="needs the number density"): + pdf.calculate_gr() + r, gr = pdf.calculate_gr(density=0.05, returnval=True) + assert isinstance(gr, np.ndarray) + assert len(gr) == len(r) + # a density cached by a damped calculate_Gr is reused + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, damp_origin_oscillations=True, density=0.05) + assert pdf.rho0 == 0.05 + assert pdf.calculate_gr(returnval=True) is not None + + def test_damping_requires_density_and_Gr(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + with pytest.raises(RuntimeError, match="depends on Sk, reduced_pdf, and r"): + pdf.damp_origin_oscillations(density=0.05) + with pytest.raises(ValueError, match="requires `density`"): + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, damp_origin_oscillations=True) + + +# ============================================================================ +# Integration Workflows +# ============================================================================ + + +class TestIntegrationWorkflows: + """Test complete end-to-end workflows.""" + + def test_complete_pdf_workflow_2d(self, synthetic_dataset2d): + """Test: 2D diffraction → polar transform → G(r) → g(r).""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + Gr_results = pdf.calculate_Gr( + k_min_fit=0.1, + k_max_fit=2.0, + r_max=10.0, + r_step=0.05, + returnval=True, + ) + assert Gr_results is not None + r, Gr = Gr_results + assert not np.isnan(r).any() + assert not np.isnan(Gr).any() + assert not np.isinf(Gr).any() + assert len(r) > 0 + assert len(Gr) == len(r) + gr_results = pdf.calculate_gr( + density=0.05, + returnval=True, + ) + assert gr_results is not None + r_gr, gr = gr_results + assert not np.isnan(gr).any() + assert not np.isinf(gr).any() + assert len(gr) == len(r_gr) + + def test_complete_pdf_workflow_4dstem(self, synthetic_4dstem_dataset): + """Test: 4D-STEM → origin finding → polar transform → G(r).""" + pdf = PairDistributionFunction.from_data( + synthetic_4dstem_dataset, + find_origin=True, + ) + mask = np.zeros((3, 3), dtype=bool) + mask[0:2, 0:2] = True + pdf.calculate_Gr( + k_min_fit=0.1, + k_max_fit=2.0, + mask_realspace=mask, + ) + assert pdf.reduced_pdf is not None + assert not np.isnan(pdf.reduced_pdf).any() + assert not np.isinf(pdf.reduced_pdf).any() + + def test_polar_transform_input_types(self, synthetic_diffraction_pattern): + """Test from_data works with Dataset2d and Dataset4dstem.""" + # Test with Dataset2d + ds2 = Dataset2d.from_array( + array=synthetic_diffraction_pattern, + name="test", + ) + pdf_ds2 = PairDistributionFunction.from_data( + ds2, + find_origin=False, + ) + assert pdf_ds2.polar.shape[2] == 180 + + # Test with Dataset4dstem + array_4d = synthetic_diffraction_pattern[None, None, :, :] # (1, 1, ny, nx) + ds4 = Dataset4dstem.from_array(array_4d, name="test") + pdf_ds4 = PairDistributionFunction.from_data( + ds4, + find_origin=False, + ) + assert pdf_ds4.polar.shape[2] == 180 + assert pdf_ds2.polar.shape == pdf_ds4.polar.shape + + def test_low_r_correction_workflow(self, synthetic_dataset2d): + """Test: G(r) calculation → low-r correction with a known density → g(r).""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + Fk_damped, G_cor = pdf.damp_origin_oscillations( + density=0.05, + max_iter=5, + tolerance=1e-2, + ) + assert torch.isfinite(G_cor).all() + assert len(G_cor) == len(pdf.r) + results = pdf.calculate_gr( + density=0.05, + returnval=True, + ) + assert results is not None + r, gr = results + assert not np.isnan(gr).any() + + +# ============================================================================ +# Test state handling across repeated calls +# ============================================================================ + + +class TestStateHandling: + """Repeated calls must give the same result as a fresh object.""" + + def test_refit_background_when_fit_range_changes(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + pdf.calculate_Gr(k_min_fit=0.3, k_max_fit=1.5) + + fresh = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + fresh.calculate_Gr(k_min_fit=0.3, k_max_fit=1.5) + + np.testing.assert_allclose(pdf.bg.cpu().numpy(), fresh.bg.cpu().numpy()) + np.testing.assert_allclose(pdf.reduced_pdf, fresh.reduced_pdf) + + def test_unmasked_call_after_masked_recomputes_radial_mean(self, synthetic_4dstem_dataset): + pdf = PairDistributionFunction.from_data(synthetic_4dstem_dataset, find_origin=False) + mask = np.zeros((3, 3), dtype=bool) + mask[0, 0] = True + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, mask_realspace=mask) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + + fresh = PairDistributionFunction.from_data(synthetic_4dstem_dataset, find_origin=False) + fresh.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + + np.testing.assert_allclose(pdf.Ik.cpu().numpy(), fresh.Ik.cpu().numpy()) + np.testing.assert_allclose(pdf.reduced_pdf, fresh.reduced_pdf) + + def test_direct_fit_bg_does_not_poison_calculate_Gr(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + Ik = pdf.calculate_radial_mean(returnval=True) + pdf.fit_bg(Ik, kmin=0.5, kmax=1.0) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + + fresh = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + fresh.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + + np.testing.assert_allclose(pdf.bg.cpu().numpy(), fresh.bg.cpu().numpy()) + np.testing.assert_allclose(pdf.reduced_pdf, fresh.reduced_pdf) + + def test_rerun_clears_stale_pdf(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + assert pdf.Fk_masked is None + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, r_max=10.0, r_step=0.05) + pdf.calculate_gr(density=0.05) + assert pdf.Fk_masked is not None + + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, r_max=5.0, r_step=0.02) + assert pdf.pdf is None + r, gr = pdf.calculate_gr(density=0.05, returnval=True) + assert len(gr) == len(r) == len(pdf.r) + + def test_calculate_gr_new_density_after_damping_matches_fresh_run(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0, damp_origin_oscillations=True, density=0.05) + pdf.calculate_gr(density=0.1) + + fresh = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + fresh.calculate_Gr( + k_min_fit=0.1, k_max_fit=2.0, damp_origin_oscillations=True, density=0.1 + ) + fresh.calculate_gr(density=0.1) + + assert pdf.rho0 == 0.1 + np.testing.assert_allclose( + pdf.reduced_pdf_damped.cpu().numpy(), fresh.reduced_pdf_damped.cpu().numpy() + ) + np.testing.assert_allclose(pdf.pdf, fresh.pdf) + + def test_max_iter_reaches_the_low_r_correction(self, synthetic_dataset2d): + pdf = PairDistributionFunction.from_data(synthetic_dataset2d, find_origin=False) + kw = dict(k_min_fit=0.1, k_max_fit=2.0, damp_origin_oscillations=True, density=0.05) + # no iterations: the corrected G(r) is the uncorrected one + pdf.calculate_Gr(**kw, max_iter=0) + G0 = pdf.reduced_pdf + np.testing.assert_allclose( + pdf.reduced_pdf_damped.cpu().numpy(), G0, rtol=0, atol=1e-4 * np.abs(G0).max() + ) + pdf.calculate_Gr(**kw, max_iter=1) + G1 = pdf.reduced_pdf_damped.cpu().numpy() + pdf.calculate_Gr(**kw, max_iter=6) + assert not np.allclose(G1, pdf.reduced_pdf_damped.cpu().numpy()) + + +# ============================================================================ +# Known-answer tests: analytic hard-sphere S(k) with known density +# ============================================================================ + + +def _percus_yevick_S(k, rho0=0.05, eta=0.45, u=0.15): + """Analytic hard-sphere S(k) (Percus-Yevick, Gaussian damped): physical, known rho0.""" + q = 2 * np.pi * k + sig = (6 * eta / (np.pi * rho0)) ** (1 / 3) + a = (1 + 2 * eta) ** 2 / (1 - eta) ** 4 + b = -6 * eta * (1 + eta / 2) ** 2 / (1 - eta) ** 4 + g = eta * (1 + 2 * eta) ** 2 / (2 * (1 - eta) ** 4) + r = np.linspace(0, sig, 4000) + x = r / sig + c = -(a + b * x + g * x**3) + qr = np.outer(q, r) + j0 = np.ones_like(qr) + nz = qr != 0 + j0[nz] = np.sin(qr[nz]) / qr[nz] + chat = 4 * np.pi * np.trapezoid(r**2 * c * j0, r, axis=1) + S = 1 / (1 - rho0 * chat) + return 1 + (S - 1) * np.exp(-(q**2) * u**2 / 2) + + +class TestKnownAnswer: + """Drive the transform with a known S(k) (background fit bypassed: bg = 0, f = 1).""" + + RHO0 = 0.05 + + @classmethod + def _pdf_with_injected_S(cls): + img = np.random.default_rng(0).random((401, 401)).astype(np.float32) + ds = Dataset2d.from_array( + array=img, + name="hard_spheres", + origin=(0, 0), + sampling=(0.01, 0.01), + units=["1/Angstrom", "1/Angstrom"], + signal_units="counts", + ) + pdf = PairDistributionFunction.from_data(ds, find_origin=False) + k = np.asarray(pdf.qq) + S = _percus_yevick_S(k, rho0=cls.RHO0) + pdf.Ik = torch.tensor(S - 1, dtype=torch.float32) + pdf.bg = torch.zeros(len(k)) + pdf.f = torch.ones(len(k)) + pdf.fit_bg = lambda Ik, kmin=None, kmax=None: (pdf.bg, pdf.f) + return pdf, k, S + + def test_G_matches_windowed_transform_of_true_S(self): + """G(r) = (2/pi) sum F(q) sin(qr) dq with F = 2 pi k [S-1]: catches any 2pi, dk or + r-scaling error in the transform (py4DSTEM's G is 4 pi^2 too small).""" + pdf, k, S = self._pdf_with_injected_S() + pdf.calculate_Gr(k_min_fit=0.01, k_max_fit=float(k[-1]), r_max=10.0, r_step=0.02) + r = pdf.r + wk = pdf._k_window( + torch.tensor(k, dtype=torch.float32), pdf.kmin_window, pdf.kmax_window + ).numpy() + Fk = k * (S - 1) + F = 2 * np.pi * Fk * wk + dk = k[1] - k[0] + G_ref = (2 / np.pi) * dk * 2 * np.pi * (np.sin(2 * np.pi * np.outer(r, k)) @ F) + G_ref[0] = 0.0 + np.testing.assert_allclose(pdf.reduced_pdf, G_ref, rtol=0, atol=2e-3 * np.abs(G_ref).max()) + # first-neighbour peak of the hard-sphere shell (diameter 2.58 A) + assert abs(r[np.argmax(pdf.reduced_pdf)] - 2.58) < 0.35 + + def test_low_r_correction_with_known_density(self): + """With the true density and r_cut below the contact distance, the corrected G(r) + is -4 pi rho0 r there, g(r) is ~0, and the first peak is not moved.""" + pdf, k, S = self._pdf_with_injected_S() + pdf.calculate_Gr(k_min_fit=0.01, k_max_fit=float(k[-1]), r_max=10.0, r_step=0.02) + r = pdf.r + peak_before = r[np.argmax(pdf.reduced_pdf)] + pdf.calculate_Gr( + k_min_fit=0.01, + k_max_fit=float(k[-1]), + r_max=10.0, + r_step=0.02, + damp_origin_oscillations=True, + density=self.RHO0, + r_cut=2.0, # below the hard-sphere contact at 2.58 A + ) + G_cor = pdf.reduced_pdf_damped.cpu().numpy() + m = (r > 0.3) & (r < 1.5) + np.testing.assert_allclose(G_cor[m] / r[m], -4 * np.pi * self.RHO0, rtol=0.05) + r_g, g = pdf.calculate_gr(returnval=True) + assert np.max(np.abs(g[m])) < 0.1 + assert abs(r[np.argmax(G_cor)] - peak_before) < 0.1