de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/diffraction/prism.py

@@ -0,0 +1,204 @@
+"""Scattering-matrix Bloch backend: the batched, differentiable form of ``bloch``.
+
+Replaces the per-direction dynamical matrix with batched solves of the
+scattering matrix, much faster for dense direction grids and equal to
+:class:`~de_ebsd.diffraction.backends.BlochWaveBackend` in double precision.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+import de_ebsd.diffraction.reflections as _R
+from de_ebsd.diffraction.backends import BlochWaveBackend, SimulationInputs
+from de_ebsd.diffraction.reflections import excitation_errors, master_reflection_list
+
+
+class BlochPrismBackend(BlochWaveBackend):
+    """Scattering-matrix (PRISM) Bloch backend: one shared eigen-decomposition.
+
+    ``cdtype`` selects the working precision: ``torch.complex128`` (default,
+    exact against the reference) or ``torch.complex64`` (faster, ~1e-4).
+
+    The whole intensity path is torch (``matrix_exp`` + a depth recurrence), so
+    it is autograd-differentiable: with a ``requires_grad`` structure-factor LUT
+    (``u_lut``) gradients flow through the Bloch solve back to ``U_g`` (matching
+    finite differences; M4). Connecting ``U_g`` to crystal parameters needs torch
+    structure factors -- the remaining M4 piece.
+    """
+
+    name = "bloch_prism"
+
+    def __init__(self, cdtype: torch.dtype = torch.complex128):
+        self.cdtype = cdtype
+        self.rdtype = torch.float64 if cdtype == torch.complex128 else torch.float32
+        # class_chunk bounds the [chunk, Cmax, Cmax] candidate-Bethe tensor; block
+        # is classified up front so buckets span it (large, GPU-efficient batches);
+        # matmul_batch caps the [B, S, S] matrix_exp batch.
+        self.class_chunk = 1024
+        self.block = 65536
+        self.matmul_batch = 4096
+
+    # ---- setup ---------------------------------------------------------- #
+    def prepare(self, inputs: SimulationInputs) -> None:
+        super().prepare(
+            inputs
+        )  # structure_factors, bethe, depth profile, Sgh LUT, kn, gzero, n_atoms, hmax
+        dev = torch.device(self.device if self.device is not None else "cpu")
+        self.tdev = dev
+        self.off = 2 * self.hmax
+        self.u_lut = torch.as_tensor(self.structure_factors.ucg_lut, device=dev, dtype=self.cdtype)
+        self.sgh_t = torch.as_tensor(self.sgh_lut, device=dev, dtype=self.cdtype)
+        self.lam = float(self.structure_factors.wavelength)
+        self.upz0 = float(self.structure_factors.upz0)
+        self.lambda_t = torch.as_tensor(self.lambda_e, device=dev, dtype=self.rdtype)
+        self.n_depths = self.lambda_t.shape[0]
+
+        # fixed, direction-independent tables (g = 0 prepended at index 0)
+        master = master_reflection_list(self.structure_factors)
+        z1 = np.zeros((1, 3), np.int64)
+        self.hkl_aug = torch.as_tensor(np.concatenate([z1, master.hkl]), device=dev)
+        self.n_beams = self.hkl_aug.shape[0]
+        u0 = float(np.abs(self.structure_factors.lookup(np.zeros(3, np.int64))))
+        self.u_mod = torch.as_tensor(
+            np.concatenate([[u0], master.u_mod]), device=dev, dtype=self.rdtype
+        )
+        self.forbidden = torch.as_tensor(np.concatenate([[False], master.forbidden]), device=dev)
+        self.g_cart = torch.as_tensor(
+            self.structure_factors.crystal.g_cartesian(master.hkl), device=dev, dtype=self.rdtype
+        )
+
+    def _lookup(self, lut: torch.Tensor, hkl: torch.Tensor) -> torch.Tensor:
+        i = hkl + self.off
+        return lut[i[..., 0], i[..., 1], i[..., 2]]
+
+    # ---- depth integration via the scattering matrix -------------------- #
+    def _calc_lgh(self, a: torch.Tensor) -> torch.Tensor:
+        phase = 2.0 * np.pi * self.depthstep / (2.0 * self.kn)
+        step = torch.complex(
+            torch.zeros((), device=self.tdev, dtype=self.rdtype),
+            torch.as_tensor(phase, device=self.tdev, dtype=self.rdtype),
+        )
+        sigma = torch.linalg.matrix_exp(step * a)  # exp(i phase A)
+        psi = sigma[:, :, self.gzero].clone()  # psi_1 = Sigma e_gzero
+        lgh = torch.zeros_like(a)
+        for n in range(self.n_depths):
+            lgh = lgh + self.lambda_t[n] * (psi.conj()[:, :, None] * psi[:, None, :])
+            psi = torch.matmul(sigma, psi[:, :, None]).squeeze(-1)
+        return lgh * (self.depthstep / self.thick)
+
+    # ---- one intensity per incident direction --------------------------- #
+    def intensities(self, directions: torch.Tensor) -> torch.Tensor:
+        k_all = directions.to(self.tdev, self.rdtype)
+        total = k_all.shape[0]
+        out = torch.empty(total, device=self.tdev, dtype=self.out_dtype)
+        for m0 in range(0, total, self.block):
+            k = k_all[m0 : m0 + self.block]
+            strong, weak, sg = self._classify(k)
+            counts = strong.sum(1)
+            for s in torch.unique(counts).tolist():
+                sub = torch.nonzero(counts == s, as_tuple=False).squeeze(1)
+                for b0 in range(0, sub.shape[0], self.matmul_batch):
+                    self._bucket(
+                        int(s), sub[b0 : b0 + self.matmul_batch], strong, weak, sg, out, m0
+                    )
+        return out
+
+    def _classify(self, k: torch.Tensor):
+        bethe = self.bethe
+        d = k.shape[0]
+        strong = torch.zeros(d, self.n_beams, dtype=torch.bool, device=self.tdev)
+        weak = torch.zeros_like(strong)
+        sg_all = torch.empty(d, self.n_beams, device=self.tdev, dtype=self.rdtype)
+        for c0 in range(0, d, self.class_chunk):
+            kc = k[c0 : c0 + self.class_chunk]
+            c = kc.shape[0]
+            sg = excitation_errors(self.g_cart, kc, self.lam)
+            sg = torch.cat([torch.zeros(c, 1, device=self.tdev, dtype=self.rdtype), sg], dim=1)
+            absg = sg.abs()
+            cand = (absg / (self.lam * self.u_mod.clamp_min(1e-300))) <= bethe.cutoff
+            cand[:, self.forbidden] = absg[:, self.forbidden] <= bethe.sg_double
+            cand[:, 0] = True
+
+            # gather the candidate beams to a padded [c, Cmax] index list
+            cc = cand.sum(1)
+            cmax = int(cc.max().item())
+            nz = cand.nonzero(as_tuple=False)  # [T, 2], row-sorted
+            rows = nz[:, 0]
+            starts = torch.zeros(c, dtype=torch.long, device=self.tdev)
+            starts[1:] = cc.cumsum(0)[:-1]
+            within = torch.arange(rows.shape[0], device=self.tdev) - starts[rows]
+            cidx = torch.zeros(c, cmax, dtype=torch.long, device=self.tdev)
+            creal = torch.zeros(c, cmax, dtype=torch.bool, device=self.tdev)
+            cidx[rows, within] = nz[:, 1]
+            creal[rows, within] = True
+
+            # Bethe min over candidate h only
+            absg_c = absg.gather(1, cidx)  # [c, Cmax]
+            hkl_c = self.hkl_aug[cidx]  # [c, Cmax, 3]
+            u = (
+                self._lookup(self.u_lut, hkl_c[:, :, None, :] - hkl_c[:, None, :, :])
+                .abs()
+                .to(self.rdtype)
+            )
+            good = (u > _R._FORBIDDEN_TOL) & creal[:, None, :]
+            rh = torch.where(
+                good,
+                (absg_c[:, :, None] / self.lam) / u.clamp_min(1e-300),
+                torch.full_like(u, 1.0e4),
+            )
+            m_c = rh.min(dim=2).values  # [c, Cmax]
+            st_c = (m_c <= bethe.strong) & creal
+            wk_c = (m_c > bethe.strong) & (m_c <= bethe.weak) & creal
+
+            st = torch.zeros(c, self.n_beams, dtype=torch.bool, device=self.tdev)
+            wk = torch.zeros_like(st)
+            st.scatter_(1, cidx, st_c)
+            wk.scatter_(1, cidx, wk_c)
+            st[:, 0] = True
+            st[:, self.forbidden] = False
+            wk[:, self.forbidden] = False
+            wk[:, 0] = False
+            strong[c0 : c0 + c] = st
+            weak[c0 : c0 + c] = wk
+            sg_all[c0 : c0 + c] = sg
+        return strong, weak, sg_all
+
+    def _bucket(self, s, sub, strong, weak, sg, out, c0):
+        b = sub.shape[0]
+        st, wk, sgs = strong[sub], weak[sub], sg[sub]
+        s_idx = st.nonzero(as_tuple=False)[:, 1].reshape(b, s)  # ascending -> g0 first
+        hkl_s = self.hkl_aug[s_idx]  # [b, s, 3]
+        sg_s = sgs.gather(1, s_idx)  # [b, s]
+
+        wc = wk.sum(1)
+        w_max = int(wc.max().item()) if b else 0
+        a = self._lookup(self.u_lut, hkl_s[:, :, None, :] - hkl_s[:, None, :, :]).clone()
+        weak_diag = torch.zeros(b, s, device=self.tdev, dtype=self.rdtype)
+        if w_max > 0:
+            nz = wk.nonzero(as_tuple=False)  # [T, 2], row-sorted
+            rows, cols = nz[:, 0], nz[:, 1]
+            starts = torch.zeros(b, dtype=torch.long, device=self.tdev)
+            starts[1:] = wc.cumsum(0)[:-1]
+            within = torch.arange(rows.shape[0], device=self.tdev) - starts[rows]
+            w_idx = torch.zeros(b, w_max, dtype=torch.long, device=self.tdev)
+            w_ok = torch.zeros(b, w_max, dtype=torch.bool, device=self.tdev)
+            w_idx[rows, within] = cols
+            w_ok[rows, within] = True
+            hkl_w = self.hkl_aug[w_idx]
+            sg_w = sgs.gather(1, w_idx)
+            inv_sgw = torch.where(w_ok, 1.0 / sg_w, torch.zeros_like(sg_w))
+            u_gw = self._lookup(self.u_lut, hkl_s[:, :, None, :] - hkl_w[:, None, :, :])
+            u_wh = self._lookup(self.u_lut, hkl_w[:, :, None, :] - hkl_s[:, None, :, :])
+            a = a - 0.5 * self.lam * torch.matmul(u_gw * inv_sgw[:, None, :], u_wh)
+            weak_diag = 0.5 * self.lam * (u_gw.abs() ** 2 * inv_sgw[:, None, :]).sum(2)
+
+        di = torch.arange(s, device=self.tdev)
+        diag = (2.0 * sg_s / self.lam - weak_diag).to(self.cdtype) + 1j * self.upz0
+        a[:, di, di] = diag
+
+        lgh = self._calc_lgh(a)
+        sgh = self._lookup(self.sgh_t, hkl_s[:, None, :, :] - hkl_s[:, :, None, :])  # [g,h]=h-g
+        inten = (lgh * sgh).sum((1, 2)).real / self.n_atoms
+        out[sub + c0] = inten.to(self.out_dtype)

de_ebsd/diffraction/reflections.py

@@ -0,0 +1,234 @@
+"""Reflection lists, excitation errors, and the strong/weak Bethe partition.
+
+Ports of EMsoft's ``Initialize_ReflectionList``, ``CalcsgDouble`` and
+``Apply_BethePotentials``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from de_ebsd.crystal.structure_factors import StructureFactors
+from de_ebsd.diffraction.bethe import BetheParameters
+
+if TYPE_CHECKING:
+    from anyplotlib import Figure
+
+__all__ = [
+    "STRONG",
+    "WEAK",
+    "IGNORED",
+    "ReflectionList",
+    "master_reflection_list",
+    "excitation_errors",
+    "candidate_reflections",
+    "partition_strong_weak",
+]
+
+# status codes
+STRONG = 0
+WEAK = 1
+IGNORED = 2
+
+_FORBIDDEN_TOL = 1e-4  # |U_g| in nm^-2 below which a reflection counts as forbidden
+
+
+@dataclass
+class ReflectionList:
+    """A list of reflections, optionally classified for one beam direction.
+
+    Attributes
+    ----------
+    hkl, g, d, u_mod
+        Miller indices (N, 3), lengths ``|g|`` (nm^-1), spacings d (nm), and
+        structure-factor moduli ``|U_g|`` (nm^-2).
+    forbidden
+        True where ``|U_g|`` is (numerically) zero; such reflections only matter
+        as double-diffraction channels.
+    k
+        The beam direction (unit vector, Cartesian crystal frame) used for
+        the classification, or None for the master list.
+    sg
+        Excitation errors (nm^-1) for ``k``, or None.
+    status
+        Per-reflection STRONG / WEAK / IGNORED for ``k``, or None.
+    """
+
+    hkl: np.ndarray
+    g: np.ndarray
+    d: np.ndarray = field(repr=False)
+    u_mod: np.ndarray = field(repr=False)
+    forbidden: np.ndarray = field(repr=False)
+    k: np.ndarray | None = None
+    sg: np.ndarray | None = field(default=None, repr=False)
+    status: np.ndarray | None = field(default=None, repr=False)
+
+    def __len__(self) -> int:
+        return len(self.hkl)
+
+    @property
+    def n_strong(self) -> int:
+        if self.status is None:
+            raise ValueError("no classification; call simulator.reflections(k=...)")
+        return int(np.sum(self.status == STRONG))
+
+    @property
+    def n_weak(self) -> int:
+        if self.status is None:
+            raise ValueError("no classification; call simulator.reflections(k=...)")
+        return int(np.sum(self.status == WEAK))
+
+    def strong(self) -> np.ndarray:
+        """Miller indices of the strong reflections."""
+        if self.status is None:
+            raise ValueError("no classification; call simulator.reflections(k=...)")
+        return self.hkl[self.status == STRONG]
+
+    def weak(self) -> np.ndarray:
+        """Miller indices of the weak reflections."""
+        if self.status is None:
+            raise ValueError("no classification; call simulator.reflections(k=...)")
+        return self.hkl[self.status == WEAK]
+
+    def show(self, n: int = 40) -> Figure:
+        """Bar chart of ``|U_g|`` for the first ``n`` reflections.
+
+        With a classification, bars are colored by status (strong / weak /
+        ignored); without one, all bars share one color.
+        """
+        from de_ebsd.viz.figures import bar_figure
+
+        take = slice(0, n)
+        labels = [f"{h} {k} {l_}" for h, k, l_ in self.hkl[take]]
+        values = self.u_mod[take]
+        colors = None
+        title = f"Reflections (of {len(self)})"
+        if self.status is not None:
+            palette = {STRONG: "#4fc3f7", WEAK: "#ffb74d", IGNORED: "#9e9e9e"}
+            colors = [palette[s] for s in self.status[take]]
+            title = (
+                f"Reflections for k = {np.round(self.k, 4)} "
+                f"({self.n_strong} strong / {self.n_weak} weak of {len(self)})"
+            )
+        return bar_figure(values, labels, title=title, colors=colors)
+
+    def __repr__(self) -> str:
+        if self.status is None:
+            return f"ReflectionList({len(self)} reflections, master list)"
+        return (
+            f"ReflectionList({len(self)} candidates for k={np.round(self.k, 4)}: "
+            f"{self.n_strong} strong, {self.n_weak} weak)"
+        )
+
+
+def master_reflection_list(structure_factors: StructureFactors) -> ReflectionList:
+    """All centering-allowed reflections with d >= dmin, sorted by ``|g|``."""
+    u_mod = np.abs(structure_factors.ucg)
+    return ReflectionList(
+        hkl=structure_factors.hkl,
+        g=structure_factors.g,
+        d=structure_factors.d,
+        u_mod=u_mod,
+        forbidden=u_mod < _FORBIDDEN_TOL,
+    )
+
+
+def excitation_errors(
+    g_cart: torch.Tensor,
+    k_hat: torch.Tensor,
+    wavelength: float,
+) -> torch.Tensor:
+    """Excitation errors s_g in nm^-1 (EMsoft ``Calcsg``), batched over beams.
+
+    Parameters
+    ----------
+    g_cart
+        Reflections in the Cartesian crystal frame, shape (G, 3), nm^-1.
+    k_hat
+        Unit incident-beam directions, shape (..., 3). The foil normal is
+        taken along the beam (master pattern geometry).
+    wavelength
+        Electron wavelength in nm (refraction corrected).
+
+    Returns
+    -------
+    torch.Tensor
+        s_g of shape (..., G).
+    """
+    k = k_hat / wavelength  # (..., 3), nm^-1
+    kg = k.unsqueeze(-2) + g_cart  # (..., G, 3)
+    numerator = -(g_cart * (2.0 * k.unsqueeze(-2) + g_cart)).sum(-1)  # -g.(2k+g)
+    denominator = 2.0 * (kg * k_hat.unsqueeze(-2)).sum(-1)  # 2 (k+g).n
+    return numerator / denominator
+
+
+def candidate_reflections(
+    refl: ReflectionList,
+    sg: np.ndarray,
+    wavelength: float,
+    bethe: BetheParameters,
+) -> np.ndarray:
+    """Mask of reflections kept for one beam direction.
+
+    EMsoft ``Initialize_ReflectionList``: keep when
+    ``|s_g| / (lambda |U_g|) <= cutoff``; forbidden reflections are kept as
+    double-diffraction channels when ``|s_g| <= sg_double``.
+    """
+    with np.errstate(divide="ignore"):
+        r_g = np.abs(sg) / (wavelength * np.maximum(refl.u_mod, 1e-300))
+    keep = r_g <= bethe.cutoff
+    keep[refl.forbidden] = np.abs(sg[refl.forbidden]) <= bethe.sg_double
+    return keep
+
+
+def partition_strong_weak(
+    hkl: np.ndarray,
+    sg: np.ndarray,
+    forbidden: np.ndarray,
+    structure_factors: StructureFactors,
+    bethe: BetheParameters,
+) -> np.ndarray:
+    """Classify candidate reflections as STRONG / WEAK / IGNORED.
+
+    Port of EMsoft ``Apply_BethePotentials``: for each reflection g the
+    perturbation strength is ``min_h |s_g| / (lambda |U_{g-h}|)`` over all
+    candidates h; thresholds ``bethe.strong`` and ``bethe.weak`` split the
+    classes. The incident beam (first entry, g = 0) is always strong.
+
+    Systematically absent (double-diffraction) reflections carry no structure
+    factor (``|U_g| ~ 0``) and so cannot be Bloch beams: they are never strong
+    or weak. They stay in the candidate list only as difference partners for
+    genuine double-diffraction paths ``U_{g-h}`` of the *allowed* beams. Without
+    this exclusion the degenerate self difference ``g - g = 0`` (whose
+    coefficient ``U_0`` is the largest in the table) would drive ``m`` down to
+    ``|s_g|/(lambda |U_0|)`` and spuriously promote any near-Bragg forbidden
+    reflection to a strong beam -- adding beams EMsoft never includes and
+    corrupting the dynamical matrix.
+
+    Parameters are the candidate subset (index 0 must be the g = 0 beam).
+    """
+    n = len(hkl)
+    status = np.full(n, IGNORED, dtype=np.int64)
+    status[0] = STRONG
+
+    diff = hkl[:, None, :] - hkl[None, :, :]  # (N, N, 3)
+    u_diff = np.abs(structure_factors.lookup(diff))  # (N, N)
+    inv_lambda = 1.0 / structure_factors.wavelength
+
+    # rh(i, h) = |s_i| / lambda / |U_{g_i - g_h}|; forbidden difference
+    # vectors are excluded from the minimum (EMsoft sets rh = 10000 there).
+    with np.errstate(divide="ignore"):
+        rh = (inv_lambda * np.abs(sg))[:, None] / u_diff
+    rh = np.where(u_diff > _FORBIDDEN_TOL, rh, 1.0e4)
+    m = rh.min(axis=1)
+
+    status[1:][m[1:] <= bethe.strong] = STRONG
+    between = (m[1:] > bethe.strong) & (m[1:] <= bethe.weak)
+    status[1:][between] = WEAK
+    # forbidden reflections are difference partners only, never beams themselves
+    status[forbidden] = IGNORED
+    return status

de_ebsd/index/__init__.py

@@ -0,0 +1,61 @@
+"""Orientation indexing: the matching counterpart to :mod:`de_ebsd.sim`.
+
+:class:`SphericalIndexer` back-projects an experimental pattern to the sphere and
+cross-correlates it against the master over all of SO(3) via spherical harmonics
+and the Wigner-D transform (after EMSphInx; Lenthe, Singh & De Graef, 2019).
+:class:`DictionaryIndexer` is the brute-force template-matching alternative.
+Both emit orix objects for kikuchipy.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+__all__ = [
+    "DictionaryIndexer",
+    "HREBSDResult",
+    "LiveIndexer",
+    "SphericalIndexer",
+    "determine_pattern_center",
+    "hrebsd",
+    "master_from_emsoft",
+    "phase_from_crystal",
+    "refine_hr",
+    "refine_orientation",
+    "remap_reference",
+    "remove_background",
+    "remove_dynamic_background",
+    "remove_static_background",
+    "to_crystal_map",
+]
+
+_ORIGIN = {
+    "DictionaryIndexer": "de_ebsd.index.dictionary",
+    "SphericalIndexer": "de_ebsd.index.spherical",
+    "determine_pattern_center": "de_ebsd.index.refine",
+    "HREBSDResult": "de_ebsd.index.crosscorr",
+    "LiveIndexer": "de_ebsd.index.live",
+    "hrebsd": "de_ebsd.index.crosscorr",
+    "remap_reference": "de_ebsd.index.crosscorr",
+    "refine_hr": "de_ebsd.index.refine",
+    "refine_orientation": "de_ebsd.index.refine",
+    "master_from_emsoft": "de_ebsd.index.io",
+    "phase_from_crystal": "de_ebsd.index.io",
+    "to_crystal_map": "de_ebsd.index.io",
+    "remove_background": "de_ebsd.index.preprocess",
+    "remove_dynamic_background": "de_ebsd.index.preprocess",
+    "remove_static_background": "de_ebsd.index.preprocess",
+}
+
+
+def __getattr__(name: str):
+    # lazy so ``import de_ebsd.index`` stays light: orix, diffpy and
+    # anyplotlib are only imported when actually used
+    module = _ORIGIN.get(name)
+    if module is None:
+        raise AttributeError(f"module 'de_ebsd.index' has no attribute {name!r}")
+    return getattr(importlib.import_module(module), name)
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)