diff --git a/src/quantem/core/datastructures/__init__.py b/src/quantem/core/datastructures/__init__.py index dfb5b47a..ac8f3d64 100644 --- a/src/quantem/core/datastructures/__init__.py +++ b/src/quantem/core/datastructures/__init__.py @@ -2,6 +2,7 @@ from quantem.core.datastructures.vector import Vector as Vector from quantem.core.datastructures.dataset4dstem import Dataset4dstem as Dataset4dstem +from quantem.core.datastructures.polar4dstem import Polar4dstem as Polar4dstem from quantem.core.datastructures.dataset4d import Dataset4d as Dataset4d from quantem.core.datastructures.dataset3d import Dataset3d as Dataset3d from quantem.core.datastructures.dataset2d import Dataset2d as Dataset2d diff --git a/src/quantem/core/datastructures/polar4dstem.py b/src/quantem/core/datastructures/polar4dstem.py new file mode 100644 index 00000000..e46dd4f7 --- /dev/null +++ b/src/quantem/core/datastructures/polar4dstem.py @@ -0,0 +1,139 @@ +from typing import Self + +import numpy as np +import torch +from numpy.typing import NDArray + +from quantem.core.datastructures.dataset4d import Dataset4d + + +class Polar4dstem(Dataset4d): + """4D-STEM dataset in polar coordinates (scan_row, scan_col, phi, r_pix).""" + + def __init__( + self, + 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: + metadata = {} + mdata_keys_polar = [ + "polar_radial_min", + "polar_radial_max", + "polar_radial_step", + "polar_num_annular_bins", + "polar_two_fold_rotation_symmetry", + "polar_ellipticity", + ] + for k in mdata_keys_polar: + if k not in metadata: + metadata[k] = None + super().__init__( + array=array, + tensor=tensor, + name=name, + origin=origin, + sampling=sampling, + units=units, + signal_units=signal_units, + metadata=metadata, + _token=_token, + ) + self.origin_array = origin_array + + @classmethod + def from_array( + cls, + 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, + ) -> 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( + 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: + sampling = np.ones(4, dtype=float) + if units is None: + units = ["pixels", "pixels", "deg", "pixels"] + if metadata is None: + metadata = {} + return cls( + array=array, + 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, + ) + + @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: + return int(self.shape[2]) + + @property + def n_r(self) -> int: + 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 e69de29b..2e723c74 100644 --- a/src/quantem/diffraction/__init__.py +++ b/src/quantem/diffraction/__init__.py @@ -0,0 +1 @@ +from quantem.diffraction.polar import PairDistributionFunction as PairDistributionFunction diff --git a/src/quantem/diffraction/polar.py b/src/quantem/diffraction/polar.py new file mode 100644 index 00000000..d2d73eb4 --- /dev/null +++ b/src/quantem/diffraction/polar.py @@ -0,0 +1,1482 @@ +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: elliptical distortion correction in origin finding +# 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 density estimation and Yoshimoto–Omote oscillation damping + - 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_mask : torch.Tensor or None + Window applied to F(k) before the sine transform, combining a bandpass + and a Lorch taper. Set by ``calculate_Gr``. + Fk_damped : torch.Tensor or None + F(k) after iterative low-r oscillation damping, set by ``estimate_density``. + reduced_pdf_damped : torch.Tensor or None + G(r) recomputed from the damped F(k), set by ``estimate_density``. + rho0 : float or None + Estimated atomic number density ρ₀ (atoms/ų), set by ``estimate_density``. + + 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_fit=2.0, r_max=10.0) + >>> rdf.calculate_gr(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_fit=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_mask: torch.Tensor | 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() + + 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, with ``sigma = weights_fit`` (a sin² low-k taper that downweights the + central beam, times a linear factor emphasising higher k). 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). + + Background fit is fit with weighting ``sigma = weights_fit``, where + + mask_low = sin^2( clip((k - kmin) / k_width, 0, 1) * pi/2 ) + weights_fit = (1 / mask_low) * (k[-1] - 0.9*k + dk) + + The fit minimises ``sum[ ((B - I) / weights_fit)^2 ]`` over k in + [kmin, kmax]. 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. ``kmin`` is also the start of the low-k sin^2 taper + + 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] + # set up weighting: sin^2 low-k taper * linear high-k emphasis. + # taper width fixed to default (0.25 1/A) + k_width = 0.25 + ramp = torch.clamp((k_fit - kmin) / k_width, 0.0, 1.0) + mask_low = torch.sin(ramp * (torch.pi / 2.0)) ** 2 + high_k_weight = k_fit[-1] - 0.9 * k_fit + dk + # inv_weights = 1 / weights_fit = mask_low / kfac + # store inverse because can just multiply it later + inv_weights = torch.where( + mask_low > 1e-4, mask_low / high_k_weight, torch.zeros_like(mask_low) + ) # length of fit window (Nf,) + 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 = inv_weights[None, :] # 1 / weights_fit, 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_min: float = 0.0, + 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, + 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 a window in k (low-k sin^2 ramp x Lorch high-k taper). + * 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:`estimate_density` is called + and the corrected F(k)/G(r) are stored as ``self.Fk_damped`` and + ``self.reduced_pdf_damped``. The estimated density is cached in + ``self.rho0`` so that a subsequent :meth:`calculate_gr` call can reuse it. + + 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 + Minimum k (A^-1) for the structure-factor Lorch window. + If None, falls back to ``k_min_fit``. + k_max_window : float or None, optional + Maximum k (A^-1) for the structure-factor Lorch window. + 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_min : float, optional + Minimum r (A) for the real-space grid. + 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:`estimate_density` and store corrected F(k)/G(r). + density : float or None, optional + Known number density (atoms/A^3). If provided together with + ``damp_origin_oscillations=True``, the S(k)/G(r) correction uses + this value instead of estimating it. + r_cut : float, optional + Minimum radial distance (A) for peak search in density estimation. + Forwarded to :meth:`estimate_density`. + 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.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})." + ) + # Recompute the radial mean whenever a real-space mask is given + # only reuse the cache when no mask is passed, otherwise calculate_Gr + # would silently ignore mask_realspace. + if self.Ik is not None and mask_bool is None: + Ik = self.Ik + else: + Ik = self.calculate_radial_mean(mask_realspace=mask_bool, returnval=True) + # Likewise re-fit the background when the region changed. + if self.bg is not None and self.f is not None and mask_bool is None: + bg, f = self.bg, self.f + else: + 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) + # Compute Sk from Fk BEFORE applying the 2pi scaling, + # so that estimate_density corrections are 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 + # damp edges with lorch window + wk = self._lorch_window(k, self.kmin_window, self.kmax_window) + Fk_win = Fk * wk + + r = torch.arange(r_min, r_max, r_step, 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 damped unphysical oscillations near the origin by iteratively estimating density and correcting F(k) + if damp_origin_oscillations: + density_est = self.estimate_density( + density=density, + r_cut=r_cut, + max_iter=20, + tol_percent=1e-1, + ) + self.rho0 = density_est[0] + self.Fk_damped = density_est[1] + self.reduced_pdf_damped = density_est[2] + + 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, + 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 determined by (in priority order): + + 1. The ``density`` argument, if provided. + 2. ``self.rho0``, if already cached from a prior :meth:`estimate_density` call + (e.g. via ``calculate_Gr(damp_origin_oscillations=True)``). + 3. A fresh call to :meth:`estimate_density` (result cached in ``self.rho0``). + + 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``. + + Parameters + ---------- + density : float or None, optional + Number density (atoms/A^3). If None, uses cached or estimated value. + r_cut : float, optional + Minimum radial distance (A) for peak search in density estimation. + Only used when density must be estimated. Forwarded to + :meth:`estimate_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 not None: + rho0 = density + elif self.rho0 is not None: + rho0 = self.rho0 + print(f" Using estimated rho0 = {rho0:.6f} atoms/A^3", flush=True) + else: + # the oscillation correction simultaneously produces a density estimate + # if the user didn't run damping in calculate_Gr, we can still run the density estimation without using the corrected Fk/G(r) + density_est = self.estimate_density( + r_cut=r_cut, + max_iter=20, + tol_percent=1e-1, + ) + self.rho0 = density_est[0] + rho0 = self.rho0 + print(f" Estimated rho0 = {rho0:.6f} atoms/A^3", flush=True) + + # 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 estimate_density( + self, + density: float | None = None, + r_cut: float = 0.8, + max_iter: int = 40, + tol_percent: float = 1e-4, + ) -> tuple[float, torch.Tensor, torch.Tensor]: + """ + Estimate number density rho0 (atoms/A^3) and compute a corrected G(r). + + This method implements an iterative Q-space density estimation by + Yoshimoto & Omote (2022). It uses the structure factor `self.Sk` and + the reduced PDF `self.reduced_pdf` to iteratively update rho0 and a + corrected S(k) so that the implied G(r) is more physically consistent + at low r. + + If ``density`` is provided, the given value is used as a fixed rho0 + for the S(k)/G(r) correction instead of estimating it iteratively. + + This method requires that :meth:`calculate_Gr` has already been run, + because it depends on `self.Sk`, `self.reduced_pdf`, `self.r`, + and the k-window bounds (`self.kmin_fit`, `self.kmin_window`, + `self.kmax_window`). + + Parameters + ---------- + density : float or None, optional + Known number density (atoms/A^3). If provided, used as a fixed + rho0 — the iterative estimation is skipped and only the S(k)/G(r) + correction is performed. + r_cut : float, optional + Minimum radial distance (A) for the peak search used to determine + the correction interval. Peaks below this distance are ignored. + max_iter : int, optional + Maximum number of Q-space iterations. + tol_percent : float, optional + Convergence threshold on the relative change in rho0 (in %), + as defined in Eq. (12) of Yoshimoto & Omote (2022). + + Returns + ------- + rho0 : float + Number density (atoms/A^3), either provided or estimated. + Fk_win_damped : torch.Tensor + Windowed corrected reduced structure function used for the transform. + G_cor : torch.Tensor + Reduced PDF G(r) with dampened oscillations near origin. + """ + # we need the non-reduced structure factor (S(k) = 1 + F(k)/k) for the density estimation correction, + # so we compute it here from the Fk we already have + 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 estimate_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") + + # r_cut sets the minimum r for the peak search used to determine the correction interval + mask_search = self._r >= r_cut + r_search = self._r[mask_search] + G_search = self._reduced_pdf[mask_search] + # find tallest peak and first local minimum to the left of r_peak + ind_max = torch.argmax(G_search) + r_max = r_search[ind_max] + left = self._r < r_max + if not torch.any(left): + # If peak is immediately at cutoff, just use cutoff as rmin + rmin = r_cut + else: + r_left = self._r[left] + G_left = self._reduced_pdf[left] + mins_cond = (G_left[1:-1] < G_left[:-2]) & (G_left[1:-1] < G_left[2:]) + # fix indexing from slicing with +1 + mins_indices = torch.where(mins_cond)[0] + 1 + # minimum closest to the peak, else global min in left interval + if mins_indices.numel() > 0: + rmin = float(r_left[mins_indices[-1]]) + else: + rmin = float(r_left[torch.argmin(G_left)]) + # Restrict r to [0, rmin] for the correction + r_mask = (self._r >= 0.0) & (self._r <= rmin) + 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") + + # Iterative refinement of rho0 and S(k) + fixed_density = density is not None + rho0 = density if fixed_density else 0.0 + rho0_prev = None + Sk_cor = self.Sk.clone() + # calculate lorch function once bc it doesn't change during the iteration + wk = self._lorch_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 + G_beta = G_iter[r_mask] + beta_prev = None + for j in range(max_iter): + if j > 0: + G_beta = G_iter[r_mask] + # calculate alpha/beta for S(k) adjustment + # 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_beta, r_short) + if not fixed_density: + rho0 = float(torch.sum(alpha * beta) / torch.sum(alpha**2)) + if rho0_prev is not None: + Rj = abs(rho0_prev - rho0) / abs(rho0) * 100.0 + if Rj < tol_percent: + break + else: + # fixed density: converge on the S(k) correction magnitude + if beta_prev is not None: + delta = float(torch.max(torch.abs(beta - beta_prev))) + if delta < tol_percent * 1e-2: + 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 + rho0_prev = rho0 + return rho0, 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) + + # Find radial value of primary peak and trough for y-limits + # 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) + + # Find radial value of primary peak + # 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 _lorch_window(self, k: torch.Tensor, kmin: float, kmax: float) -> torch.Tensor: + """ + Construct a combined low-q taper and high-q Lorch window. + + The returned window is: + - 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) + """ + # low q taper + edge_frac_low = 0.1 # 10% of range at low-q + edge_width_low = edge_frac_low * (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, torch.zeros_like(wk), wk) + wk = torch.where(k > kmax, torch.zeros_like(wk), wk) + + # High q taper with Lorch window: w(k) = sin(pi*k/kmax)/(pi*k/kmax) + x = k / kmax + inband = (k >= kmin) & (k <= kmax) + # sinc function: sin(pi*x)/(pi*x) with limit 1 at x=0 + sinc_val = torch.where( + x == 0, + torch.ones_like(x), + torch.sin(torch.pi * x) / (torch.pi * x), + ) + lorch = torch.where(inband, sinc_val, torch.zeros_like(k)) + wk = wk * lorch + return wk + + 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 density estimation. + """ + 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..c3f9239e --- /dev/null +++ b/src/quantem/diffraction/polar_transform.py @@ -0,0 +1,947 @@ +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 + +# 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, +) -> 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. + + 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)) + # 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) + 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) + 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) + dps = dp_batch.repeat_interleave(n_cands, dim=0) + polars = F.grid_sample( + dps, grids, 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 + ] + scores = region.std(dim=2).sum(dim=2) / (region.mean(dim=2).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] + scores_coarse = region_coarse.std(dim=2).sum(dim=2) / ( + region_coarse.mean(dim=2).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 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 + if radial_max is None: + radial_max_eff = float(min(r0, (n_row - 1) - r0, c0, (n_col - 1) - c0)) + 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) + + # Use the global minimum safe radius across all origins so every scan + # position maps to the same-size polar grid (required for a uniform 4D output) + if radial_max is None: + r_row_pos = origins[:, :, 0] + r_row_neg = (n_row - 1) - origins[:, :, 0] + r_col_pos = origins[:, :, 1] + r_col_neg = (n_col - 1) - origins[:, :, 1] + radial_max_eff_array = np.minimum.reduce([r_row_pos, r_row_neg, r_col_pos, r_col_neg]) + radial_max = float(max(radial_max_eff_array.min(), radial_min + radial_step)) + + # 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, +) -> torch.Tensor: + """Score candidate origins by angular std over a mid-radius band. + Lower scores indicate better centering.""" + 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] + return region.std(dim=1).sum(dim=1) / (region.mean(dim=1).sum(dim=1) + 1e-6) + + +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_polar.py b/tests/diffraction/test_polar.py new file mode 100644 index 00000000..b87fd653 --- /dev/null +++ b/tests/diffraction/test_polar.py @@ -0,0 +1,551 @@ +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_estimates_density(self, synthetic_dataset2d): + """Test that calculate_gr estimates density when none is provided.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + results = pdf.calculate_gr(returnval=True) + assert results is not None + r, gr = results + assert isinstance(gr, np.ndarray) + assert len(gr) == len(r) + assert pdf.rho0 > 0 + + def test_estimate_density_requires_Gr(self, synthetic_dataset2d): + """Test that estimate_density requires prior calculate_Gr call.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + with pytest.raises( + RuntimeError, match="depends on Sk, reduced_pdf, and r from calculate_Gr" + ): + pdf.estimate_density() + + +# ============================================================================ +# 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_min=0.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_density_estimation_workflow(self, synthetic_dataset2d): + """Test: G(r) calculation → density estimation → g(r) calculation.""" + pdf = PairDistributionFunction.from_data( + synthetic_dataset2d, + find_origin=False, + ) + pdf.calculate_Gr(k_min_fit=0.1, k_max_fit=2.0) + rho0, Fk_damped, G_cor = pdf.estimate_density( + max_iter=5, + tol_percent=1.0, + ) + assert rho0 > 0 + assert np.isfinite(rho0) + results = pdf.calculate_gr( + density=rho0, + returnval=True, + ) + assert results is not None + r, gr = results + assert not np.isnan(gr).any()