de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/index/refine/pattern_center.py

@@ -0,0 +1,356 @@
+"""Calibrate a single pattern center shared across many patterns.
+
+Alternates per-pattern orientation refinement with a global Levenberg-Marquardt
+step on the shared PC, using the band-edge samples of
+:mod:`de_ebsd.index.refine.highres`.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from de_ebsd.core.device import resolve_device
+from de_ebsd.core.rotation import quaternion_to_matrix, rotation_vector_to_matrix
+from de_ebsd.index.refine.highres import (
+    direction_cosine_derivative,
+    direction_cosines_from_pattern_center,
+    edge_importance_samples,
+    levenberg_marquardt_refine,
+    normalized_residual_jacobian,
+    standardize_tiled,
+)
+from de_ebsd.index.refine.levenberg import adapt_shared_damping
+
+__all__ = ["determine_pattern_center"]
+
+
+#: Blur of the master, in radians on the sphere, from the first round to the last.
+#: A real dynamical master is sharp: its band-edge NCC falls from 1 to ~0 within
+#: 0.02 of PC, so a start 0.03 off sees no slope at all. Blurring widens that basin;
+#: the schedule ends unblurred, so the answer is the sharp master's optimum.
+_DEFAULT_BLUR = (0.05, 0.025, 0.012, 0.0)
+
+
+def _blurred(master, sigma: float):
+    """*master* blurred by a Gaussian of *sigma* radians on the sphere.
+
+    Blurred in its square-Lambert grid, which is equal-area, so one pixel is
+    about ``sqrt(2 pi) / n`` radians everywhere; the grid's shape distortion
+    and the seam at each hemisphere's edge only matter in the coarse stages,
+    which the schedule's final unblurred stage corrects. One convolution per
+    stage, so a blurred stage samples as fast as a sharp one.
+    """
+    from de_ebsd.sim.masterpattern import MasterPattern
+
+    data = master.data
+    n = data.shape[-1]
+    sigma_px = sigma * (n - 1) / np.sqrt(2.0 * np.pi)
+    radius = min(int(np.ceil(3.0 * sigma_px)), n - 1)
+    x = torch.arange(-radius, radius + 1, device=data.device, dtype=data.dtype)
+    kernel = torch.exp(-0.5 * (x / sigma_px) ** 2)
+    kernel = kernel / kernel.sum()
+    flat = data.reshape(-1, 1, n, n)
+    flat = torch.nn.functional.pad(flat, (radius, radius, 0, 0), mode="reflect")
+    flat = torch.nn.functional.conv2d(flat, kernel.view(1, 1, 1, -1))
+    flat = torch.nn.functional.pad(flat, (0, 0, radius, radius), mode="reflect")
+    flat = torch.nn.functional.conv2d(flat, kernel.view(1, 1, -1, 1))
+    return MasterPattern(flat.reshape(data.shape), master.energies_kev, dict(master.metadata))
+
+
+def determine_pattern_center(
+    patterns,
+    master,
+    detector,
+    rotations,
+    pattern_center=None,
+    n_rounds: int = 5,
+    n_tiles: int = 16,
+    samples_per_tile: int = 8,
+    refine_orientations: bool = True,
+    n_pattern_center_steps: int = 10,
+    damping: float = 1e-3,
+    keep_best: bool = True,
+    device=None,
+    return_orientations: bool = False,
+    n_orientation_steps: int = 15,
+    blur=_DEFAULT_BLUR,
+):
+    """Calibrate one shared pattern center from many EBSD patterns.
+
+    Fits a single pattern center shared by the whole calibration set by alternating
+    an orientation-only refinement of every pattern at the current PC with a
+    Levenberg-Marquardt step on the 3-parameter PC whose normal equations are
+    accumulated over all patterns. That ensemble accumulation is what breaks the
+    single-pattern R--PC degeneracy. Calibrate once, then refine the full scan at the
+    fixed PC with ``refine_hr(refine_pattern_center=False, pattern_center=<calibrated
+    pattern_center>)``.
+
+    Parameters
+    ----------
+    patterns
+        Calibration patterns ``(N, H, W)``, or a single ``(H, W)`` image. A
+        well-distributed few dozen to few hundred are plenty to fit three
+        parameters. The raw sampled values are used.
+    master, detector
+        The continuous :class:`de_ebsd.sim.MasterPattern` and the
+        :class:`~de_ebsd.sim.detector.EBSDDetector` geometry. ``detector.pattern_center`` is
+        used only when ``pattern_center`` is ``None``.
+    rotations
+        Per-pattern seed orientations, e.g. from coarse indexing.
+    pattern_center
+        Starting shared pattern center ``(3,)``; ``None`` uses ``detector.pattern_center``.
+    n_rounds
+        Alternating (orientation, shared-PC) rounds.
+    n_tiles, samples_per_tile
+        Tiles per detector side and band-edge samples kept per tile; the defaults
+        give 2048 samples per pattern.
+    refine_orientations
+        Re-refine each orientation at the current PC every round (default), or hold
+        the input ``rotations`` fixed and fit only the PC.
+    n_pattern_center_steps
+        Levenberg-Marquardt steps on the shared PC per round.
+    damping
+        Initial LM damping, adapted each step (x0.5 accepted, x3 rejected) and
+        carried across rounds. Also seeds the inner orientation refinement.
+    keep_best
+        Never return a PC / orientation set scoring below the best seen, making the
+        calibration monotone in the total per-tile NCC.
+    device
+        Defaults to ``master.data.device``; otherwise everything is moved there.
+    return_orientations
+        Also return the refined orientations, before ``info``.
+    n_orientation_steps
+        LM steps for the inner orientation refinement each round.
+    blur
+        Gaussian blur of the master (radians on the sphere) per stage, coarse to
+        fine, spread evenly over the rounds. A sharp dynamical master only has a
+        slope within ~0.01 of the true PC; the blurred stages pull a further-off
+        seed into that basin. ``None`` or ``(0,)`` fits the sharp master only.
+
+    Returns
+    -------
+    pattern_center : torch.Tensor
+        The calibrated shared pattern center ``(3,)``.
+    rotations : orix.quaternion.Rotation
+        Only when ``return_orientations``: the orientations at the calibrated PC.
+    info : dict
+        ``"ncc"`` and ``"ncc_seed"`` (mean per-tile NCC at the calibrated and seed
+        geometry), ``"pc_step"`` (L2 magnitude of the final round's PC move -- small
+        means converged), ``"pc_step_history"``, and ``"n_rounds"``.
+    """
+    device = resolve_device(device, master.data.device)
+    dtype = master.data.dtype
+    if master.data.device != device:
+        master = master.to(device)
+
+    patt = torch.as_tensor(patterns, device=device, dtype=dtype)
+    if patt.ndim == 2:
+        patt = patt[None]
+    if patt.ndim != 3:
+        raise ValueError(f"patterns must be (H, W) or (N, H, W), got shape {tuple(patt.shape)}")
+    n_p = patt.shape[0]
+
+    q0 = torch.as_tensor(np.asarray(rotations.data), device=device, dtype=dtype).reshape(-1, 4)
+    if q0.shape[0] != n_p:
+        raise ValueError(f"rotations length {q0.shape[0]} != number of patterns {n_p}")
+    r_mat = quaternion_to_matrix(q0).detach()  # (N, 3, 3)
+
+    if pattern_center is None:
+        pattern_center = detector.pattern_center
+    pattern_center = torch.as_tensor(pattern_center, device=device, dtype=dtype)
+    if tuple(pattern_center.shape) != (3,):
+        raise ValueError(
+            "pattern_center must be a shared 3-vector (3,), "
+            f"got shape {tuple(pattern_center.shape)}"
+        )
+    pattern_center = pattern_center.clone()
+
+    # fixed band-edge samples + per-tile normalization, computed once (patterns never move)
+    rows, cols, flat, tile_ids, counts = edge_importance_samples(patt, n_tiles, samples_per_tile)
+    e_raw = patt.reshape(n_p, -1).gather(1, flat.long())  # (N, K) raw experimental
+    n_total = int(counts.shape[0])
+    n_active = int((counts > 0).sum())
+    counts_ps = counts.index_select(0, tile_ids).to(dtype)  # (K,)
+    _, _, e_hat = standardize_tiled(e_raw, tile_ids, counts_ps, n_total)  # (N, K) fixed
+
+    def _ncc_per_pattern(r_cur: torch.Tensor, pc_vec: torch.Tensor, seen=master) -> torch.Tensor:
+        pc_bc = pc_vec[None].expand(n_p, 3)
+        dc = direction_cosines_from_pattern_center(detector, pc_bc, rows, cols)
+        m = seen.sample(dc @ r_cur.transpose(-1, -2)).to(device=device, dtype=dtype)
+        _, _, m_hat = standardize_tiled(m, tile_ids, counts_ps, n_total)
+        return (e_hat * m_hat).sum(-1) / n_active
+
+    def _total(r_cur: torch.Tensor, pc_vec: torch.Tensor) -> torch.Tensor:
+        return _ncc_per_pattern(r_cur, pc_vec).sum()
+
+    ncc_seed = float(_ncc_per_pattern(r_mat, pattern_center).mean())
+
+    best_pc = pattern_center.clone()
+    best_r = r_mat.clone()
+    best_total = _total(r_mat, pattern_center)
+    damping = float(damping)
+    pc_step_history: list[float] = []
+    blur = tuple(blur) if blur else (0.0,)
+
+    for round_index in range(n_rounds):
+        # spread the blur stages over the rounds, ending on the last (sharpest)
+        sigma = blur[round_index * len(blur) // n_rounds]
+        seen = _blurred(master, sigma) if sigma > 0 else master
+
+        def _seen_total(r_cur, pc_vec, seen=seen):
+            return _ncc_per_pattern(r_cur, pc_vec, seen).sum()
+
+        # 1. orientation block -- refine each pattern at the current shared PC (well-posed)
+        if refine_orientations:
+            pc_bc = pattern_center[None].expand(n_p, 3).contiguous()
+            r_mat, _, _ = levenberg_marquardt_refine(
+                seen,
+                detector,
+                rows,
+                cols,
+                e_raw,
+                tile_ids,
+                counts,
+                r_mat,
+                pc_bc,
+                False,
+                n_orientation_steps,
+                damping,
+                keep_best,
+                device,
+                dtype,
+            )
+
+        # 2. shared-PC block -- LM on the 3-vector, normal equations summed over all
+        # patterns; jointly with the orientations when those are refined, so the PC
+        # step already accounts for how each orientation would compensate it
+        pc_before = pattern_center.clone()
+        pattern_center, r_mat, damping = _shared_pc_lm(
+            seen,
+            detector,
+            rows,
+            cols,
+            r_mat,
+            pattern_center,
+            e_hat,
+            tile_ids,
+            counts_ps,
+            n_total,
+            damping,
+            n_pattern_center_steps,
+            _seen_total,
+            device,
+            dtype,
+            joint=refine_orientations,
+        )
+        pc_step_history.append(float((pattern_center - pc_before).norm()))
+
+        # judged on the sharp master, so a blurred round cannot claim "best"
+        total = _total(r_mat, pattern_center)
+        if bool(total > best_total):
+            best_total, best_pc, best_r = total, pattern_center.clone(), r_mat.clone()
+
+    if keep_best:
+        pattern_center, r_mat = best_pc, best_r
+
+    pattern_center = pattern_center.detach()
+    info = {
+        "ncc": float(_ncc_per_pattern(r_mat, pattern_center).mean()),
+        "ncc_seed": ncc_seed,
+        "pc_step": pc_step_history[-1] if pc_step_history else 0.0,
+        "pc_step_history": pc_step_history,
+        "n_rounds": n_rounds,
+    }
+    if return_orientations:
+        from orix.quaternion import Rotation
+
+        rotation = Rotation.from_matrix(r_mat.cpu().numpy())
+        return pattern_center, rotation, info
+    return pattern_center, info
+
+
+def _shared_pc_lm(
+    master,
+    detector,
+    rows,
+    cols,
+    r_mat,
+    pattern_center,
+    e_hat,
+    tile_ids,
+    counts_ps,
+    n_total,
+    damping,
+    n_steps,
+    total_fn,
+    device,
+    dtype,
+    joint=False,
+):
+    """LM on the shared PC; with ``joint``, on the orientations alongside it.
+
+    The joint system is block-arrow: each pattern's 3x3 orientation block couples
+    only to the shared PC. Eliminating the orientation blocks (a Schur complement)
+    leaves a 3x3 system for the PC that already allows for every orientation's
+    compensating move -- so the PC is not dragged back along the R--PC valley the
+    way it is when the two blocks are refined in turn.
+    """
+    n_p = r_mat.shape[0]
+    cur_total = total_fn(r_mat, pattern_center)
+
+    for _ in range(n_steps):
+        pc_bc = pattern_center[None].expand(n_p, 3).contiguous()
+        dc = direction_cosines_from_pattern_center(
+            detector, pc_bc, rows, cols
+        ).detach()  # (N, K, 3)
+        # one reverse-mode grid_sample pass -> per-point master gradient dM/dv
+        with torch.enable_grad():
+            v = (dc @ r_mat.transpose(-1, -2)).detach().requires_grad_(True)
+            m_raw = master.sample(v)
+            (g_m,) = torch.autograd.grad(m_raw.sum(), v)
+        m = m_raw.detach().to(device=device, dtype=dtype)
+        g_m = g_m.to(device=device, dtype=dtype)
+
+        q = torch.einsum("pka,pab->pkb", g_m, r_mat)  # gM @ R -> (N, K, 3)
+        ddc_dpc = direction_cosine_derivative(detector, pc_bc, rows, cols)  # (N, K, 3, 3)
+        dm_dpc = torch.einsum("pkb,pkbc->pkc", q, ddc_dpc)  # raw dM/dpc, (N, K, 3)
+        if joint:
+            dm_dw = torch.linalg.cross(dc, q, dim=-1)  # raw dM/dw, (N, K, 3)
+            g_raw = torch.cat([dm_dw, dm_dpc], dim=-1)  # (N, K, 6)
+        else:
+            g_raw = dm_dpc
+
+        res, j_r = normalized_residual_jacobian(m, e_hat, g_raw, tile_ids, counts_ps, n_total)
+        with torch.no_grad():
+            jtr = torch.einsum("pk,pki->pi", res, j_r)  # per pattern, (N, P)
+            jtj = torch.einsum("pki,pkj->pij", j_r, j_r)  # per pattern, (N, P, P)
+            if joint:
+                a = jtj[:, :3, :3]
+                a = a + damping * torch.diag_embed(
+                    torch.diagonal(a, dim1=-2, dim2=-1).clamp_min(1e-9)
+                )
+                b = jtj[:, :3, 3:]  # (N, 3, 3) orientation-PC coupling
+                c = jtj[:, 3:, 3:].sum(0)
+                c = c + damping * torch.diag_embed(torch.diagonal(c).clamp_min(1e-9))
+                a_inv_b = torch.linalg.solve(a, b)
+                a_inv_g = torch.linalg.solve(a, jtr[:, :3, None])[..., 0]
+                schur = c - torch.einsum("pji,pjk->ik", b, a_inv_b)
+                rhs = jtr[:, 3:].sum(0) - torch.einsum("pji,pj->i", b, a_inv_g)
+                delta_pc = -torch.linalg.solve(schur, rhs)  # (3,)
+                delta_w = -(a_inv_g + (a_inv_b @ delta_pc[:, None])[..., 0])  # (N, 3)
+                r_new = r_mat @ rotation_vector_to_matrix(delta_w)
+            else:
+                jtr, jtj = jtr.sum(0), jtj.sum(0)  # summed over patterns
+                diag = torch.diagonal(jtj).clamp_min(1e-9)
+                a = jtj + damping * torch.diag_embed(diag)  # Marquardt (scale-invariant) damping
+                delta_pc = -torch.linalg.solve(a, jtr)  # (3,)
+                r_new = r_mat
+            pc_new = pattern_center + delta_pc
+            new_total = total_fn(r_new, pc_new)
+            accepted = bool(new_total > cur_total)
+            if accepted:
+                pattern_center, r_mat = pc_new, r_new
+                cur_total = new_total
+            damping = adapt_shared_damping(damping, accepted)
+    return pattern_center, r_mat, damping

de_ebsd/index/spherical/__init__.py

@@ -0,0 +1,41 @@
+"""Spherical-harmonic and SO(3) machinery behind the spherical indexer."""
+
+from de_ebsd.index.spherical.correlation import so3_cross_correlation
+from de_ebsd.index.spherical.harmonics import (
+    SphericalHarmonicTransform,
+    forward_sht,
+    gauss_legendre_grid,
+    inverse_sht,
+    lm_indices,
+    n_coefficients,
+    spherical_harmonics,
+)
+from de_ebsd.index.spherical.indexer import SphericalIndexer
+from de_ebsd.index.spherical.wigner import (
+    euler_zyz_to_orix,
+    grid_to_rotation,
+    n_orientations,
+    so3_grid,
+    wigner_D,
+    wigner_d,
+    wigner_small_d,
+)
+
+__all__ = [
+    "SphericalHarmonicTransform",
+    "SphericalIndexer",
+    "euler_zyz_to_orix",
+    "forward_sht",
+    "gauss_legendre_grid",
+    "grid_to_rotation",
+    "inverse_sht",
+    "lm_indices",
+    "n_coefficients",
+    "n_orientations",
+    "so3_cross_correlation",
+    "so3_grid",
+    "spherical_harmonics",
+    "wigner_D",
+    "wigner_d",
+    "wigner_small_d",
+]

de_ebsd/index/spherical/correlation.py

@@ -0,0 +1,106 @@
+"""SO(3) cross-correlation of two spherical-harmonic spectra.
+
+Combines the two spectra with the Wigner-d matrices into a coefficient cube,
+then one 3-D inverse FFT gives the correlation over all of SO(3) on the SOFT
+grid.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from de_ebsd.index.spherical.harmonics import lm_indices, n_coefficients
+from de_ebsd.index.spherical.wigner import so3_grid, wigner_d
+
+__all__ = ["so3_cross_correlation"]
+
+
+def _pack_to_grid(coeffs: torch.Tensor, bandwidth: int) -> torch.Tensor:
+    L = int(bandwidth)
+    degrees, orders = lm_indices(L, device=coeffs.device)
+    pos = degrees * (2 * L + 1) + (orders + L)  # flat index into (L+1)*(2L+1)
+    base = coeffs.new_zeros((*coeffs.shape[:-1], (L + 1) * (2 * L + 1)))
+    packed = base.index_add(-1, pos, coeffs)
+    return packed.reshape(*coeffs.shape[:-1], L + 1, 2 * L + 1)
+
+
+def so3_cross_correlation(
+    master_coeffs: torch.Tensor,
+    pattern_coeffs: torch.Tensor,
+    bandwidth: int,
+    wigner_d_cached: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Cross-correlation of two harmonic spectra over all of SO(3), via SOFT.
+
+    Parameters
+    ----------
+    master_coeffs : torch.Tensor
+        Master spectrum ``a_lm``, shape ``(..., (L+1)**2)`` (packed
+        ``c = l*l + l + m``, the :mod:`.harmonics` convention). May carry leading
+        batch dims that broadcast against ``pattern_coeffs``.
+    pattern_coeffs : torch.Tensor
+        Pattern spectrum ``b_lm``, shape ``(..., (L+1)**2)``. The leading dims are
+        the batch of patterns (many indexed at once).
+    bandwidth : int
+        The band limit ``L``; both spectra must have ``(L+1)**2`` slots.
+    wigner_d_cached : torch.Tensor, optional
+        A precomputed Wigner-d block ``wigner_d(L, beta)`` for this bandwidth's SOFT
+        ``beta`` grid, shape ``(N, L+1, 2L+1, 2L+1)`` with ``N = 2L + 2`` (see
+        :func:`de_ebsd.index.spherical.wigner.wigner_d`). This ``O(L^4)`` block is otherwise
+        rebuilt on every call; pass it to reuse one block across a mini-batched scan
+        (:class:`~de_ebsd.index.spherical.indexer.SphericalIndexer` caches it in
+        ``__init__``). The result is byte-identical to omitting it, provided the
+        cached block matches this bandwidth's ``beta`` grid and the real dtype used
+        here (``float64`` for complex128 spectra, else ``float32``).
+
+    Returns
+    -------
+    torch.Tensor
+        Real tensor ``C`` of shape ``(*batch, N, N, N)`` with ``N = 2L + 2``, sampled
+        on the :func:`.so3_grid` ``(alpha, beta, gamma)`` grid (axes in that order),
+        where ``batch`` is the broadcast of the two inputs' leading dims. ``argmax``
+        gives the best-matching orientation index; feed it to
+        :func:`.grid_to_rotation`. Differentiable in both spectra.
+    """
+    L = int(bandwidth)
+    n_coeff = n_coefficients(L)
+    if master_coeffs.shape[-1] != n_coeff or pattern_coeffs.shape[-1] != n_coeff:
+        raise ValueError(
+            f"expected {n_coeff} coefficients for bandwidth {L}, got "
+            f"{master_coeffs.shape[-1]} (master) and {pattern_coeffs.shape[-1]} (pattern)"
+        )
+    a = master_coeffs
+    b = pattern_coeffs
+    cdtype = torch.result_type(a, b)
+    if not cdtype.is_complex:
+        cdtype = torch.complex128 if cdtype == torch.float64 else torch.complex64
+    a = a.to(cdtype)
+    b = b.to(cdtype)
+
+    a_grid = _pack_to_grid(a, L)  # (..., L+1, 2L+1) over m'
+    b_grid = _pack_to_grid(b, L)  # (..., L+1, 2L+1) over m
+    # cross-spectrum f_{l, m, m'} = b_lm conj(a_lm')
+    f = b_grid[..., :, :, None] * a_grid.conj()[..., :, None, :]  # (..., L+1, 2L+1, 2L+1)
+
+    n = 2 * L + 2
+    real_dtype = torch.float64 if cdtype == torch.complex128 else torch.float32
+    alpha, beta, gamma = so3_grid(L, device=a.device, dtype=real_dtype)
+    if wigner_d_cached is not None:
+        d = wigner_d_cached.to(cdtype)  # reuse the precomputed O(L^4) block
+    else:
+        d = wigner_d(L, beta).to(cdtype)  # (N, L+1, 2L+1, 2L+1)
+
+    # Wigner-d contraction over beta: S_{k, m, m'} = sum_l f_{l, m, m'} d^l_{m m'}(beta_k)
+    s = torch.einsum("...lmn,klmn->...kmn", f, d)  # (..., N, 2L+1, 2L+1)
+
+    # inverse FFT over alpha (m) and gamma (m'): pad the (2L+1) order axes to N,
+    # then C = N^2 ifft2(s) with the linear phase that shifts m,m' back to -L..L
+    s = torch.nn.functional.pad(s, (0, 1, 0, 1))  # (..., N, N, N)
+    c = (n * n) * torch.fft.ifft2(s, dim=(-2, -1))  # axis -2 -> alpha, -1 -> gamma
+    phase_a = torch.exp(-1j * L * alpha).to(cdtype).reshape(n, 1)
+    phase_g = torch.exp(-1j * L * gamma).to(cdtype).reshape(1, n)
+    c = c * phase_a * phase_g  # (..., N=beta, N=alpha, N=gamma)
+
+    # reorder (beta, alpha, gamma) -> (alpha, beta, gamma) and take the real part
+    c = c.transpose(-3, -2)
+    return c.real.contiguous()