de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/index/spherical/harmonics.py

@@ -0,0 +1,311 @@
+"""Spherical harmonic transform on a Gauss-Legendre x equiangular grid.
+
+Forward and inverse SHT for band-limited functions on the sphere, with the
+associated Legendre functions built by stable upward recursion.
+:class:`SphericalHarmonicTransform` caches the grid and basis so the master and
+every pattern are expanded identically.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+
+from de_ebsd.core.device import resolve_device
+
+__all__ = [
+    "SphericalHarmonicTransform",
+    "forward_sht",
+    "gauss_legendre_grid",
+    "inverse_sht",
+    "lm_indices",
+    "n_coefficients",
+    "spherical_harmonics",
+]
+
+_INV_SQRT_4PI = 1.0 / math.sqrt(4.0 * math.pi)
+
+
+def n_coefficients(bandwidth: int) -> int:
+    """Number of packed coefficients for a transform of the given bandwidth ``L``."""
+    return (bandwidth + 1) ** 2
+
+
+def lm_indices(
+    bandwidth: int, device: torch.device | str | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Degree and order of each packed coefficient slot.
+
+    Returns ``(degrees, orders)``, each a 1-D ``long`` tensor of length
+    ``(L + 1) ** 2``, so that slot ``c`` holds ``a_{l m}`` with ``l = degrees[c]``
+    and ``m = orders[c]``, ordered ``c = l * l + l + m``.
+    """
+    degrees = []
+    orders = []
+    for ell in range(bandwidth + 1):
+        for m in range(-ell, ell + 1):
+            degrees.append(ell)
+            orders.append(m)
+    dl = torch.tensor(degrees, dtype=torch.long, device=device)
+    dm = torch.tensor(orders, dtype=torch.long, device=device)
+    return dl, dm
+
+
+def _gauss_legendre_nodes(
+    n: int, device: torch.device | str | None, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    x, w = np.polynomial.legendre.leggauss(n)
+    xt = torch.as_tensor(x, device=device, dtype=dtype)
+    wt = torch.as_tensor(w, device=device, dtype=dtype)
+    return xt, wt
+
+
+def gauss_legendre_grid(
+    bandwidth: int,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype = torch.float64,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build the exact quadrature grid for bandwidth ``L``.
+
+    Returns ``(theta, phi, weights)``:
+
+    - ``theta`` -- ``(L + 1,)`` colatitude nodes, ``arccos`` of the Gauss-Legendre
+      nodes (in ``[0, pi]``).
+    - ``phi`` -- ``(2 L + 2,)`` equiangular azimuth nodes ``2 pi k / N_phi``.
+    - ``weights`` -- ``(L + 1, 2 L + 2)`` quadrature weights ``w_theta[j] * (2 pi /
+      N_phi)`` such that ``sum weights[j, k] f(theta_j, phi_k)`` integrates any
+      band-limited ``f`` over the sphere exactly.
+    """
+    n_theta = bandwidth + 1
+    n_phi = 2 * bandwidth + 2
+    x, w_theta = _gauss_legendre_nodes(n_theta, device, dtype)
+    theta = torch.arccos(x)
+    phi = 2.0 * math.pi * torch.arange(n_phi, device=device, dtype=dtype) / n_phi
+    weights = w_theta[:, None] * (2.0 * math.pi / n_phi)
+    weights = weights.expand(n_theta, n_phi).contiguous()
+    return theta, phi, weights
+
+
+def _normalized_alf(bandwidth: int, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
+    L = bandwidth
+    p = x.new_zeros((L + 1, L + 1, *x.shape))
+    p[0, 0] = _INV_SQRT_4PI
+    # sectoral terms P_m^m (Condon-Shortley phase accumulates through the -1)
+    for m in range(1, L + 1):
+        p[m, m] = -math.sqrt((2 * m + 1) / (2 * m)) * u * p[m - 1, m - 1]
+    # ascend in degree l at fixed order m
+    for m in range(L + 1):
+        if m + 1 <= L:
+            # l = m + 1: the P_{l-2} term vanishes (P_{m-1}^m = 0)
+            a = math.sqrt((2 * (m + 1) - 1) * (2 * (m + 1) + 1) / ((m + 1 - m) * (m + 1 + m)))
+            p[m + 1, m] = a * x * p[m, m]
+        for ell in range(m + 2, L + 1):
+            a = math.sqrt((2 * ell - 1) * (2 * ell + 1) / ((ell - m) * (ell + m)))
+            b = math.sqrt(
+                (2 * ell + 1)
+                * (ell + m - 1)
+                * (ell - m - 1)
+                / ((2 * ell - 3) * (ell - m) * (ell + m))
+            )
+            p[ell, m] = a * x * p[ell - 1, m] - b * p[ell - 2, m]
+    return p
+
+
+def spherical_harmonics(bandwidth: int, theta: torch.Tensor, phi: torch.Tensor) -> torch.Tensor:
+    """Dense complex spherical-harmonic basis ``Y_lm`` on a grid.
+
+    ``theta`` (``(N_theta,)`` colatitudes) and ``phi`` (``(N_phi,)`` azimuths) are
+    the 1-D node vectors of a separable grid. Returns the complex design matrix of
+    shape ``((L + 1) ** 2, N_theta, N_phi)`` with slot ``c`` (ordering of
+    :func:`lm_indices`) holding ``Y_{l m}`` evaluated at every ``(theta_j, phi_k)``.
+    Matches ``scipy.special.sph_harm_y(l, m, theta, phi)``.
+    """
+    x = torch.cos(theta)
+    u = torch.sin(theta).clamp_min(0.0)
+    p = _normalized_alf(bandwidth, x, u)  # (L+1, L+1, N_theta)
+    degrees, orders = lm_indices(bandwidth, device=theta.device)
+    # P_l^{|m|} for each coefficient slot -> (n_coeff, N_theta)
+    p_lm = p[degrees, orders.abs()]
+    # Condon-Shortley: Y_l^{-m} = (-1)^m conj(Y_l^m); fold the sign of negative m
+    parity = 1.0 - 2.0 * (orders.abs() % 2).to(theta.dtype)  # (-1)^|m|
+    sign = torch.where(orders < 0, parity, torch.ones_like(parity))
+    m_real = orders.to(theta.dtype)
+    phase = torch.exp(1j * (m_real[:, None] * phi[None, :]))  # (n_coeff, N_phi)
+    y = (sign[:, None].to(phase.dtype) * p_lm.to(phase.dtype))[:, :, None] * phase[:, None, :]
+    return y
+
+
+def forward_sht(
+    values: torch.Tensor,
+    theta: torch.Tensor,
+    phi: torch.Tensor,
+    bandwidth: int,
+    weights: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Stateless forward transform on a separable grid (dense design matrix).
+
+    ``values`` has shape ``(..., N_theta, N_phi)`` (real or complex), sampled at the
+    ``theta``/``phi`` nodes. Returns complex coefficients ``(..., (L + 1) ** 2)``.
+    ``weights`` are the quadrature weights ``(N_theta, N_phi)``; if omitted the
+    Gauss-Legendre weights for this bandwidth are used, which assumes ``theta`` /
+    ``phi`` are the canonical grid from :func:`gauss_legendre_grid`. For repeated
+    transforms prefer :class:`SphericalHarmonicTransform`.
+    """
+    y = spherical_harmonics(bandwidth, theta, phi)  # (n_coeff, N_theta, N_phi)
+    if weights is None:
+        _, _, weights = gauss_legendre_grid(bandwidth, theta.device, theta.dtype)
+    weighted = (values * weights).to(y.dtype)
+    return torch.einsum("...tp,ctp->...c", weighted, y.conj())
+
+
+def inverse_sht(coeffs: torch.Tensor, theta: torch.Tensor, phi: torch.Tensor) -> torch.Tensor:
+    """Stateless inverse transform: coefficients ``a_lm`` -> samples on a grid.
+
+    ``coeffs`` has shape ``(..., (L + 1) ** 2)``; the bandwidth is inferred from its
+    length. Returns complex samples ``(..., N_theta, N_phi)`` at the ``theta``/``phi``
+    nodes (the dense synthesis ``sum_lm a_lm Y_lm``).
+    """
+    n_coeff = coeffs.shape[-1]
+    bandwidth = int(round(math.sqrt(n_coeff))) - 1
+    y = spherical_harmonics(bandwidth, theta, phi)
+    return torch.einsum("...c,ctp->...tp", coeffs.to(y.dtype), y)
+
+
+class SphericalHarmonicTransform:
+    """A cached, separable spherical harmonic transform for a fixed bandwidth.
+
+    Builds the exact Gauss-Legendre x equiangular quadrature grid and the
+    associated-Legendre design matrices once, then transforms via an FFT over phi
+    and a per-order matmul over theta. Both directions are batched over arbitrary
+    leading dimensions and are autograd-differentiable in the samples/coefficients.
+
+    Parameters
+    ----------
+    bandwidth : int
+        The band limit ``L``. Sets the quadrature grid (``N_theta = L + 1``,
+        ``N_phi = 2 L + 2``) and the packed coefficient count ``(L + 1) ** 2``.
+    device : torch.device or str, optional
+        Device the grid and cached design matrices are built on (default CPU).
+    dtype : torch.dtype, optional
+        Real floating dtype of the transform; must be ``torch.float32`` or
+        ``torch.float64`` (default ``torch.float64``). The complex coefficient dtype
+        is the matching ``complex64`` / ``complex128``.
+
+    Attributes
+    ----------
+    bandwidth : int
+        The band limit ``L``.
+    theta, phi : torch.Tensor
+        1-D quadrature nodes (``(L + 1,)`` and ``(2 L + 2,)``).
+    weights : torch.Tensor
+        ``(L + 1, 2 L + 2)`` quadrature weights.
+    n_coeff : int
+        Number of packed coefficients ``(L + 1) ** 2``.
+    """
+
+    def __init__(
+        self,
+        bandwidth: int,
+        *,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float64,
+    ) -> None:
+        if dtype not in (torch.float32, torch.float64):
+            raise ValueError("dtype must be a real floating type (float32 or float64)")
+        self.bandwidth = int(bandwidth)
+        self.device = resolve_device(device)
+        self.dtype = dtype
+        self.cdtype = torch.complex128 if dtype == torch.float64 else torch.complex64
+
+        L = self.bandwidth
+        self.n_theta = L + 1
+        self.n_phi = 2 * L + 2
+        self.n_coeff = n_coefficients(L)
+
+        theta, phi, weights = gauss_legendre_grid(L, self.device, dtype)
+        self.theta = theta
+        self.phi = phi
+        self.weights = weights
+
+        degrees, orders = lm_indices(L, device=self.device)
+        self._degrees = degrees
+        self._orders = orders
+        # FFT bin holding Fourier order m (negative m wraps to N_phi + m)
+        self._freq_bins = torch.remainder(orders, self.n_phi)
+
+        with torch.no_grad():
+            x = torch.cos(theta)
+            u = torch.sin(theta).clamp_min(0.0)
+            _, w_theta = _gauss_legendre_nodes(self.n_theta, self.device, dtype)
+            p = _normalized_alf(L, x, u)  # (L+1, L+1, N_theta)
+            p_lm = p[degrees, orders.abs()]  # (n_coeff, N_theta)
+            # fold the Condon-Shortley sign of negative orders into the design matrix
+            parity = 1.0 - 2.0 * (orders.abs() % 2).to(dtype)  # (-1)^|m|
+            sign = torch.where(orders < 0, parity, torch.ones_like(parity))
+            design = sign[:, None] * p_lm  # A[c, j] = s(m) P_l^{|m|}(cos theta_j)
+            self._design = design.contiguous()
+            # forward weights: (2 pi / N_phi) w_theta[j] A[c, j]
+            self._fwd = ((2.0 * math.pi / self.n_phi) * w_theta[None, :] * design).contiguous()
+
+    def to(
+        self,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> SphericalHarmonicTransform:
+        """Return a plan on a new device / real dtype (cheap to rebuild)."""
+        return SphericalHarmonicTransform(
+            self.bandwidth,
+            device=device if device is not None else self.device,
+            dtype=dtype if dtype is not None else self.dtype,
+        )
+
+    def grid_directions(self) -> torch.Tensor:
+        """Cartesian unit vectors ``(N_theta, N_phi, 3)`` of every quadrature node.
+
+        ``n = (sin th cos ph, sin th sin ph, cos th)`` -- sample the sphere here
+        (e.g. resample a master pattern onto these directions), then ``forward``.
+        """
+        st = torch.sin(self.theta)
+        ct = torch.cos(self.theta)
+        cp = torch.cos(self.phi)
+        sp = torch.sin(self.phi)
+        nx = st[:, None] * cp[None, :]
+        ny = st[:, None] * sp[None, :]
+        nz = ct[:, None].expand(self.n_theta, self.n_phi)
+        return torch.stack([nx, ny, nz], dim=-1)
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        """Samples ``(..., N_theta, N_phi)`` -> coefficients ``(..., n_coeff)``.
+
+        ``values`` may be real or complex; the result is complex. Differentiable
+        in ``values``.
+        """
+        if values.shape[-2:] != (self.n_theta, self.n_phi):
+            raise ValueError(
+                f"values trailing shape {tuple(values.shape[-2:])} != grid "
+                f"{(self.n_theta, self.n_phi)}"
+            )
+        # FFT over phi: fhat[..., j, m] = sum_k values[..., j, k] exp(-i m phi_k)
+        fhat = torch.fft.fft(values.to(self.cdtype), dim=-1)
+        # pick the bin of each coefficient's order -> (..., N_theta, n_coeff)
+        fhat_c = fhat.index_select(-1, self._freq_bins)
+        # a_c = sum_j fwd[c, j] fhat_c[..., j, c]
+        return torch.einsum("cj,...jc->...c", self._fwd.to(self.cdtype), fhat_c)
+
+    def inverse(self, coeffs: torch.Tensor) -> torch.Tensor:
+        """Coefficients ``(..., n_coeff)`` -> samples ``(..., N_theta, N_phi)``.
+
+        Returns complex samples (take ``.real`` for a real-valued function).
+        Differentiable in ``coeffs``.
+        """
+        if coeffs.shape[-1] != self.n_coeff:
+            raise ValueError(f"coeffs last dim {coeffs.shape[-1]} != n_coeff {self.n_coeff}")
+        coeffs = coeffs.to(self.cdtype)
+        # per-order theta synthesis, scattered onto FFT bins: g[..., j, bin]
+        contrib = torch.einsum("...c,cj->...jc", coeffs, self._design.to(self.cdtype))
+        g = coeffs.new_zeros((*coeffs.shape[:-1], self.n_theta, self.n_phi))
+        g.index_add_(-1, self._freq_bins, contrib)
+        # inverse FFT (unnormalized): f[..., j, k] = sum_bin g[..., j, bin] exp(+i bin phi_k)
+        return self.n_phi * torch.fft.ifft(g, dim=-1)

de_ebsd/index/spherical/indexer.py

@@ -0,0 +1,345 @@
+"""End-to-end spherical indexer.
+
+Back-projects each pattern to the sphere, expands it in spherical harmonics,
+cross-correlates against the precomputed master spectrum over all of SO(3), and
+reads the orientation off the peak. ``normalize=True`` (the default) divides out
+the rotation-dependent master energy inside the field of view, which keeps a
+bright master region from outscoring the true orientation.
+
+``grid_to_rotation(argmax)`` is directly the sample orientation passed to
+``MasterPattern.project`` -- no inverse or transpose.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+from de_ebsd.core.device import resolve_device
+from de_ebsd.index.backproject import pattern_to_sphere
+from de_ebsd.index.spherical.correlation import so3_cross_correlation
+from de_ebsd.index.spherical.harmonics import SphericalHarmonicTransform
+from de_ebsd.index.spherical.wigner import grid_to_rotation, so3_grid, wigner_d
+
+
+class SphericalIndexer:
+    """Index EBSD patterns against a master pattern by spherical cross-correlation.
+
+    Build once (the master spectrum is precomputed here), then call :meth:`index`
+    on a single pattern or a stack.
+
+    Parameters
+    ----------
+    master
+        A :class:`de_ebsd.sim.MasterPattern`, sampled via ``master.sample``.
+    detector
+        :class:`~de_ebsd.sim.detector.EBSDDetector` geometry, inverted to
+        back-project each pattern to the sphere.
+    bandwidth
+        Band limit ``L``, the accuracy/speed knob: the SOFT grid step is
+        ``360 / (2L + 2)`` degrees. Cost grows ~``O(L^4)`` in both time and memory,
+        since :func:`~de_ebsd.index.spherical.wigner.wigner_d` materializes a
+        ``(N, L+1, 2L+1, 2L+1)`` tensor.
+    device
+        Resolved as explicit argument > ``master.device`` > CPU.
+    energy_kev
+        Energy bin to sample the master at; ``None`` (default) selects the highest.
+
+    Attributes
+    ----------
+    sht : SphericalHarmonicTransform
+        The shared transform plan; master and patterns both go through it.
+    a_lm, a2_lm : torch.Tensor
+        Spectra of the master and of ``M^2``, complex, shape ``((L+1)**2,)``.
+    """
+
+    def __init__(self, master, detector, bandwidth=53, device=None, energy_kev=None):
+        self.master = master
+        self.detector = detector
+        self.bandwidth = int(bandwidth)
+        self.energy_kev = energy_kev
+
+        self.device = resolve_device(device, getattr(master, "device", "cpu"))
+
+        # pin the plan to the master's real dtype so master and patterns transform
+        # identically; the plan requires a real float type
+        master_dtype = getattr(getattr(master, "data", None), "dtype", torch.float64)
+        self.dtype = (
+            master_dtype if master_dtype in (torch.float32, torch.float64) else (torch.float64)
+        )
+        self._eps = torch.finfo(self.dtype).eps
+
+        # the plan does NOT auto-move from its CPU/float64 default, so pin it here
+        self.sht = SphericalHarmonicTransform(self.bandwidth, device=self.device, dtype=self.dtype)
+
+        # quadrature directions M = (L+1)*(2L+2)
+        self._dirs = self.sht.grid_directions().reshape(-1, 3)
+
+        m_samples = master.sample(self._dirs, energy_kev=self.energy_kev)
+        m_samples = m_samples.to(device=self.device, dtype=self.dtype)
+        grid_shape = (self.sht.n_theta, self.sht.n_phi)
+        self.a_lm = self.sht.forward(m_samples.reshape(grid_shape))
+        self.a2_lm = self.sht.forward((m_samples**2).reshape(grid_shape))
+
+        self._weights_flat = self.sht.weights.reshape(-1)
+
+        # cache the O(L^4) Wigner-d block: so3_cross_correlation otherwise rebuilds it
+        # per call, so a mini-batched scan would pay for it per chunk
+        _, beta_grid, _ = so3_grid(self.bandwidth, device=self.device, dtype=self.dtype)
+        self._wig_d = wigner_d(self.bandwidth, beta_grid)
+
+        # normalizer terms, cached on the first normalize=True call; they depend only
+        # on the master and the fixed detector pattern_center, so a variable pattern_center must
+        # invalidate them
+        self._w_lm = None
+        self._a_w = None
+        self._denom_geo = None
+
+    def _geometry_normalizer(self):
+        if self._denom_geo is not None:
+            return self._denom_geo
+        # the FOV mask is fixed by detector geometry: recover it by back-projecting a
+        # zero pattern and keeping only the mask output
+        zero = torch.zeros(
+            (1, self.detector.nrows, self.detector.ncols),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        _, w = pattern_to_sphere(zero, self.detector, self._dirs)
+        w_lm = self.sht.forward(w.reshape(self.sht.n_theta, self.sht.n_phi))
+        # master mean and second moment over the rotated FOV, aligned elementwise with X
+        s1 = so3_cross_correlation(self.a_lm, w_lm, self.bandwidth, wigner_d_cached=self._wig_d)
+        s2 = so3_cross_correlation(self.a2_lm, w_lm, self.bandwidth, wigner_d_cached=self._wig_d)
+        a_w = math.sqrt(4.0 * math.pi) * w_lm[0].real  # FOV spherical area
+        denom = (s2 - s1**2 / a_w.clamp_min(self._eps)).clamp_min(self._eps).sqrt()
+        self._w_lm = w_lm
+        self._a_w = a_w
+        self._denom_geo = denom
+        return denom
+
+    def index(self, patterns, normalize=True, return_surface=False):
+        """Index one pattern or a stack, returning orix orientations and quality.
+
+        Parameters
+        ----------
+        patterns
+            A detector image ``(H, W)`` or stack ``(Np, H, W)``, numpy or torch.
+        normalize
+            Divide out the master's local FOV norm (default), removing the bias
+            toward orientations that place bright master regions in the field of
+            view. ``False`` keeps the raw masked cross-correlation, unbiased only
+            for near-uniform-norm masters.
+        return_surface
+            Also return the SO(3) score surface that was arg-maxed, as a third
+            element. It feeds :func:`de_ebsd.viz.maps.show_correlation` directly.
+
+        Returns
+        -------
+        rotations : orix.quaternion.Rotation
+            Best-matching sample orientation per pattern, length ``Np``.
+        quality : torch.Tensor
+            Peak score per pattern. Normalized, this is roughly in ``[-1, 1]`` and
+            comparable across patterns; unnormalized it has arbitrary scale.
+        surface : torch.Tensor
+            Only when ``return_surface``: shape ``(Np, N, N, N)``, ``N = 2L + 2``,
+            on the ``(alpha, beta, gamma)`` grid. Detached.
+        """
+        patt = torch.as_tensor(patterns, device=self.device, dtype=self.dtype)
+        single = patt.ndim == 2
+        if single:
+            patt = patt[None]
+        if patt.ndim != 3:
+            raise ValueError(
+                f"patterns must be (H, W) or (Np, H, W), got shape {tuple(patt.shape)}"
+            )
+
+        samples, mask = pattern_to_sphere(patt, self.detector, self._dirs)
+
+        # standardize over the FOV and zero outside it. Zeroing the FOV mean makes E's
+        # l=0 harmonic vanish, which drops the DC term from the masked correlation.
+        w = mask
+        denom = w.sum().clamp_min(self._eps)
+        mean = (samples * w).sum(-1, keepdim=True) / denom
+        var = ((samples - mean) ** 2 * w).sum(-1, keepdim=True) / denom
+        e = ((samples - mean) / var.clamp_min(self._eps).sqrt()) * w
+        b_lm = self.sht.forward(e.reshape(*e.shape[:-1], self.sht.n_theta, self.sht.n_phi))
+
+        c = so3_cross_correlation(self.a_lm, b_lm, self.bandwidth, wigner_d_cached=self._wig_d)
+        n_grid = 2 * self.bandwidth + 2
+
+        if normalize:
+            # rho = X / (sqrt(P) * denom_geo). P and A_W are g-independent, so they
+            # scale the surface but not the argmax.
+            denom_geo = self._geometry_normalizer()
+            p_energy = (e**2 * self._weights_flat).sum(-1)
+            scale = p_energy.clamp_min(self._eps).sqrt().reshape(-1, 1, 1, 1)
+            surface = c / (scale * denom_geo[None])
+        else:
+            surface = c
+
+        flat = surface.reshape(surface.shape[0], -1)
+        quality, peak_flat = flat.max(dim=-1)
+        j, k, p = (idx.cpu() for idx in torch.unravel_index(peak_flat, (n_grid, n_grid, n_grid)))
+
+        rotations = grid_to_rotation(self.bandwidth, j, k, p)
+        if return_surface:
+            return rotations, quality.detach(), surface.detach()
+        return rotations, quality.detach()
+
+    def refine(
+        self,
+        patterns,
+        rotations=None,
+        pattern_center=None,
+        refine_pattern_center=False,
+        n_iterations=100,
+        learning_rate=0.02,
+        pattern_center_learning_rate=None,
+        keep_best=True,
+        return_history=False,
+    ):
+        """Gradient-refine indexed orientations (and optionally the pattern center).
+
+        Wraps :func:`de_ebsd.index.refine.orientation.refine_orientation` with this
+        indexer's master, detector, device and SHT quadrature directions, so
+        refinement scores over the same grid the coarse index used. With
+        ``rotations=None`` the seeds come from :meth:`index`, taking raw patterns
+        straight to sharpened orientations.
+
+        ``pattern_center`` defaults to ``detector.pattern_center``; the remaining arguments pass
+        through.
+        Returns ``(rotation, score)``, with ``pattern_center`` appended when
+        ``refine_pattern_center`` and
+        ``history`` when ``return_history``.
+        """
+        from de_ebsd.index.refine.orientation import refine_orientation
+
+        if rotations is None:
+            rotations, _ = self.index(patterns)
+        return refine_orientation(
+            patterns,
+            self.master,
+            self.detector,
+            rotations,
+            pattern_center=pattern_center,
+            refine_pattern_center=refine_pattern_center,
+            n_iterations=n_iterations,
+            learning_rate=learning_rate,
+            pattern_center_learning_rate=pattern_center_learning_rate,
+            directions=self._dirs,
+            device=self.device,
+            keep_best=keep_best,
+            return_history=return_history,
+        )
+
+    def index_scan(
+        self,
+        patterns,
+        shape=None,
+        step=1.0,
+        phase=None,
+        batch_size=None,
+        normalize=True,
+        refine=False,
+        refine_iters=100,
+        refine_lr=0.02,
+    ):
+        """Index a whole scan and assemble it into an orix ``CrystalMap``.
+
+        Parameters
+        ----------
+        patterns
+            Either ``(rows, cols, H, W)``, from which the scan grid is inferred, or
+            a flat row-major ``(N, H, W)``, which requires ``shape``.
+        shape
+            ``(rows, cols)`` of the scan. Cross-checked against 4-D input.
+        step
+            Scan step size, written onto the map's ``x``/``y``.
+        phase
+            An orix :class:`~orix.crystal_map.Phase` or a de_ebsd
+            :class:`~de_ebsd.crystal.Crystal`. If ``None``, derived from
+            ``master.metadata["crystal"]``.
+        batch_size
+            Patterns per mini-batch (default 64). The SO(3) surface is
+            ``(batch_size, N, N, N)``, which dominates memory at production
+            bandwidth. Results are identical to a single batch.
+        refine
+            Gradient-refine each mini-batch before assembly, making
+            ``prop["quality"]`` the refined NCC. Orientation only -- a scan shares
+            one detector geometry, and joint pattern-center fitting over a whole
+            scan is not done here.
+        refine_iters, refine_lr
+            Adam iterations and rotation learning rate used when ``refine``.
+
+        Returns
+        -------
+        orix.crystal_map.CrystalMap
+            One orientation + ``prop["quality"]`` per scan point, single phase.
+        """
+        import numpy as np
+        from orix.quaternion import Rotation
+
+        from de_ebsd.index.io import phase_from_crystal, to_crystal_map
+
+        # shape inference and reshaping stay on the stack's own device (view-only);
+        # only each mini-batch moves to self.device, so a 10^6-pattern scan never
+        # lands wholesale in GPU memory
+        patt = torch.as_tensor(patterns, dtype=self.dtype)
+        if patt.ndim == 4:
+            rows, cols = int(patt.shape[0]), int(patt.shape[1])
+            if shape is not None and (int(shape[0]), int(shape[1])) != (rows, cols):
+                raise ValueError(
+                    f"shape {tuple(shape)} does not match the (rows, cols, H, W) "
+                    f"pattern grid {(rows, cols)}"
+                )
+            patt = patt.reshape(rows * cols, patt.shape[2], patt.shape[3])
+        elif patt.ndim == 3:
+            if shape is None:
+                raise ValueError(
+                    "shape=(rows, cols) is required for (N, H, W) patterns "
+                    "(pass (rows, cols, H, W) to infer it instead)"
+                )
+            rows, cols = int(shape[0]), int(shape[1])
+            if rows * cols != patt.shape[0]:
+                raise ValueError(
+                    f"shape {(rows, cols)} implies {rows * cols} patterns but got {patt.shape[0]}"
+                )
+        else:
+            raise ValueError(
+                f"patterns must be (rows, cols, H, W) or (N, H, W), got shape {tuple(patt.shape)}"
+            )
+
+        n = rows * cols
+        bs = int(batch_size) if batch_size is not None else 64
+        if bs < 1:
+            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
+
+        rot_chunks = []
+        qual_chunks = []
+        for start in range(0, n, bs):
+            chunk = patt[start : start + bs].to(self.device)
+            rot_c, qual_c = self.index(chunk, normalize=normalize)
+            if refine:
+                rot_c, qual_c = self.refine(
+                    chunk, rotations=rot_c, n_iterations=refine_iters, learning_rate=refine_lr
+                )
+            rot_chunks.append(np.asarray(rot_c.data))
+            qual_chunks.append(qual_c)
+        rotations = Rotation(np.concatenate(rot_chunks, axis=0))
+        quality = torch.cat(qual_chunks)
+
+        if phase is None:
+            crystal = None
+            meta = getattr(self.master, "metadata", None)
+            if isinstance(meta, dict):
+                crystal = meta.get("crystal")
+            if crystal is None:
+                raise ValueError(
+                    "no phase given and master.metadata has no 'crystal' to derive "
+                    "one from; pass phase=<orix Phase or de_ebsd Crystal> to "
+                    "index_scan (the Phase gives the CrystalMap its crystal symmetry)"
+                )
+            phase = phase_from_crystal(crystal)
+
+        return to_crystal_map(rotations, quality, (rows, cols), phase, step=step)