de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/index/live/indexer.py

@@ -0,0 +1,386 @@
+"""Live indexing: a buffer of camera frames indexed and refined where it already is.
+
+Built for a camera streaming thousands of patterns a second into GPU memory. A
+buffer of frames is indexed in a fixed number of batched steps -- nothing loops
+over patterns or phases -- and only each pattern's phase, orientation and score
+leave the device.
+
+1. **Prepare** (:meth:`LiveIndexer.prepare`): each frame averaged over the whole
+   sensor down to the refine shape, divided by the scan's static background (the
+   running mean of its frames, :meth:`LiveIndexer.observe`, started from a search's
+   and blurred for a single crystal, :meth:`LiveIndexer.restart_static`), the
+   per-pattern dynamic background removed, then averaged again to the search shape.
+2. **Search** (:meth:`LiveIndexer.search`): one matrix product against every
+   phase's dictionary at once. The dictionaries are the phases' dynamical master
+   patterns projected over their fundamental zones, stacked into one matrix, so the
+   best entry gives the phase and a seed orientation together; a low symmetry's
+   large dictionary is scored a chunk of entries at a time.
+3. **Refine** (:meth:`LiveIndexer.refine`): a local dictionary around each seed --
+   the seed and the six rotations of one step about each axis -- projected from the
+   master on the fly at the refine shape. A parabola through each axis's three
+   scores moves the orientation; the step shrinks each level. Every level is the
+   same seven projections per pattern, so its cost is fixed and known.
+4. **Pattern center** (:meth:`LiveIndexer.refine_pattern_center`): the center
+   shared by a buffer's best patterns, refined the same way -- seven centers a
+   level, each scored with every pattern's orientation refined again there.
+
+Orientations are rotation matrices ``g`` taking a sample-frame direction to the
+crystal frame (``v_crystal = g @ v_sample``), as :meth:`MasterPattern.project`
+uses them.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import replace
+
+import numpy as np
+import torch
+
+from de_ebsd.core.device import resolve_device
+from de_ebsd.core.lambert import sample_texture
+from de_ebsd.core.rotation import quaternion_to_matrix, rotation_vector_to_matrix
+from de_ebsd.index import _kernels
+from de_ebsd.index.live import calibration, deformation
+from de_ebsd.index.live.background import ScanBackground
+from de_ebsd.index.live.common import BATCH, STENCIL, average_to, parabola_offset
+from de_ebsd.index.preprocess import gaussian_blur, normalize_rows
+
+__all__ = ["LiveIndexer"]
+
+
+#: Refine steps, in degrees, each half the last. On a twin scan rendered from a
+#: dynamical master (1024 patterns, 33 grains, 3° dictionary, 128² refine) these
+#: gave a median of 0.14° from the truth, against 0.23° for (1.5, 0.5, 0.17) and
+#: 0.19° with a fourth level of 0.06.
+REFINE_STEPS_DEG = (2.0, 1.0, 0.5, 0.25, 0.12)
+
+#: How much the search's copy of the masters is blurred (Gaussian sigma, degrees on
+#: the sphere). A sharp dynamical pattern's score falls off so steeply away from
+#: its orientation that at a 3° dictionary no entry is near enough for some
+#: grains and a distant wrong one wins: 6-9% of a twin scan's patterns landed
+#: 40-60° off unblurred, none at 2°.
+SEARCH_BLUR_DEG = 2.0
+
+
+#: Dictionary entries scored at a time by :meth:`LiveIndexer.search`, which bounds
+#: its scores at (patterns, this): 64 MB for a batch. A cubic phase's 3° dictionary
+#: (30,443 entries) is one chunk; a triclinic one's (729,000) would otherwise
+#: score into 750 MB per batch.
+SEARCH_CHUNK = 65_536
+
+
+def texture_of(masters, device, blur_deg: float = 0.0) -> torch.Tensor:
+    """Every phase's hemispheres as channels of one texture ``(1, 2 * phases, n, n)``:
+    one sampling pass serves all phases, and each pattern keeps its own phase's."""
+    if len({m.npx for m in masters}) != 1:
+        raise ValueError("every master must have the same npx")
+    texture = torch.cat([m.data[m.energy_index(None)].to(device, torch.float32) for m in masters])
+    # Near the center of the square Lambert grid a pixel spans 90° / npx.
+    per_degree = masters[0].npx / 90.0
+    if blur_deg:
+        texture = gaussian_blur(texture, blur_deg * per_degree, 4.0)
+    return texture[None].contiguous()
+
+
+class LiveIndexer:
+    """Index and refine EBSD patterns of several phases on the device.
+
+    Parameters
+    ----------
+    masters
+        One :class:`~de_ebsd.sim.MasterPattern` per phase, all at the same ``npx``.
+        Each is used at its highest energy bin.
+    point_groups
+        Each phase's orix point group, whose fundamental zone its dictionary samples.
+    detector
+        The camera's :class:`~de_ebsd.sim.detector.EBSDDetector` (its pattern center
+        and angles); frames may come at any resolution, since they are averaged to
+        ``search_shape`` and ``refine_shape``.
+    resolution
+        The dictionary's sampling, in degrees.
+    search_shape, refine_shape
+        Detector shapes the search and the refine compare patterns at.
+    refine_steps_deg
+        The refine levels' steps; empty for no refine.
+    search_blur_deg
+        How much the search's copy of the masters is blurred, see
+        :data:`SEARCH_BLUR_DEG`.
+    device
+        Defaults to the first master's.
+    dtype
+        The search dictionary's. ``None`` picks float16 on GPUs with tensor cores
+        (compute capability 7.0 and up) and float32 elsewhere: Pascal multiplies
+        float16 at 1/64 of its float32 rate, and a CPU has no float16 product. Half
+        precision halves the dictionary, which is what grows with low symmetry; the
+        search only picks seeds, so the refine sets the accuracy either way.
+    """
+
+    def __init__(
+        self,
+        masters,
+        point_groups,
+        detector,
+        *,
+        resolution: float = 3.0,
+        search_shape: tuple[int, int] = (32, 32),
+        refine_shape: tuple[int, int] = (128, 128),
+        refine_steps_deg=REFINE_STEPS_DEG,
+        search_blur_deg: float = SEARCH_BLUR_DEG,
+        device=None,
+        dtype=None,
+    ) -> None:
+        from orix.sampling import get_sample_fundamental
+
+        if len(masters) != len(point_groups):
+            raise ValueError("one point group per master")
+        self.device = resolve_device(device, masters[0].device)
+        if dtype is None:
+            tensor_cores = (
+                self.device.type == "cuda" and torch.cuda.get_device_capability(self.device)[0] >= 7
+            )
+            dtype = torch.float16 if tensor_cores else torch.float32
+        self.dtype = dtype
+        self.dictionary = None
+        self.detector = detector
+        self.search_shape = tuple(search_shape)
+        self.refine_shape = tuple(refine_shape)
+        self.refine_steps = tuple(math.radians(s) for s in refine_steps_deg)
+        self.fused = _kernels.available(self.device)
+        self.texture = texture_of(masters, self.device)
+        self.search_texture = texture_of(masters, self.device, search_blur_deg)
+        self.stencil = torch.tensor(STENCIL, device=self.device, dtype=torch.float32)
+
+        rotations, phases = [], []
+        for number, point_group in enumerate(point_groups):
+            q = np.asarray(
+                get_sample_fundamental(resolution=resolution, point_group=point_group).data
+            )
+            rotations.append(
+                quaternion_to_matrix(torch.as_tensor(q, device=self.device, dtype=torch.float32))
+            )
+            phases.append(torch.full((len(q),), number, device=self.device, dtype=torch.int64))
+        self.rotations = torch.cat(rotations)
+        self.phase_of = torch.cat(phases)
+        self.sizes = [len(r) for r in rotations]
+        self.set_pattern_center(detector.pattern_center)
+        self.background = ScanBackground(self.refine_shape, self.device)
+
+    def directions_at(self, pattern_center, shape) -> torch.Tensor:
+        detector = replace(
+            self.detector, pattern_center=tuple(float(v) for v in pattern_center), shape=shape
+        )
+        return detector.direction_cosines(device=self.device)
+
+    def set_pattern_center(self, pattern_center) -> None:
+        """Index at pattern center *pattern_center* from now on: the dictionary is projected
+        again, which takes milliseconds with the fused kernels.
+
+        Each chunk of entries is written straight into the dictionary, and a new
+        center overwrites the old one's in place: the dictionary is the largest
+        thing on the device, and building it beside a copy of itself (or a list of
+        its chunks) took two to three times its size."""
+        self.detector = replace(
+            self.detector, pattern_center=tuple(float(v) for v in pattern_center)
+        )
+        self.search_directions = self.directions_at(pattern_center, self.search_shape)
+        self.refine_directions = self.directions_at(pattern_center, self.refine_shape)
+        shape = (len(self.search_directions), len(self.rotations))
+        if self.dictionary is None or tuple(self.dictionary.shape) != shape:
+            #: ``(pixels, entries)``, so the search's product reads it at full rate.
+            self.dictionary = torch.empty(shape, device=self.device, dtype=self.dtype)
+        for lo in range(0, len(self.rotations), BATCH * 8):
+            entries = normalize_rows(
+                self.project(
+                    self.search_texture,
+                    self.rotations[lo : lo + BATCH * 8],
+                    self.phase_of[lo : lo + BATCH * 8],
+                    self.search_directions,
+                )
+            )
+            self.dictionary[:, lo : lo + len(entries)] = entries.T
+
+    def project(
+        self,
+        texture: torch.Tensor,
+        rotations: torch.Tensor,
+        phases: torch.Tensor,
+        directions: torch.Tensor,
+    ) -> torch.Tensor:
+        """The patterns ``(M, pixels)`` of *rotations* ``(M, 3, 3)``, each from its
+        own phase's master in *texture*, over *directions* ``(pixels, 3)``."""
+        if self.fused:
+            return _kernels.project(texture, rotations, phases * 2, directions)
+        v = torch.einsum("mij,pj->mpi", rotations, directions)
+        return sample_texture(texture, v, phases[:, None])
+
+    def score_candidates(
+        self, candidates: torch.Tensor, phases: torch.Tensor, rows: torch.Tensor, directions
+    ) -> torch.Tensor:
+        """Scores ``(N, K)`` of candidates ``(N, K, 3, 3)`` against their patterns'
+        normalized *rows* ``(N, pixels)``, sampled at *directions*."""
+        n, k = candidates.shape[:2]
+        flat = candidates.reshape(-1, 3, 3)
+        bases = phases.repeat_interleave(k) * 2
+        if self.fused:
+            return _kernels.score(self.texture, flat, bases, directions, rows, k).view(n, k)
+        simulated = normalize_rows(self.project(self.texture, flat, bases // 2, directions))
+        return torch.einsum("nkp,np->nk", simulated.view(n, k, -1), rows)
+
+    # --- the scan's backgrounds ------------------------------------------------
+
+    @property
+    def static(self):
+        """The scan's static background, ``None`` until enough frames have been seen."""
+        return self.background.static
+
+    @property
+    def ghost(self):
+        """The scan's grains as the static background shows them, or ``None``; see
+        :attr:`ScanBackground.ghost`."""
+        return self.background.ghost
+
+    def pooled(self, frames) -> torch.Tensor:
+        """*frames* on the device as float ``(N, H, W)`` averaged to the refine shape."""
+        t = torch.as_tensor(frames, device=self.device)
+        t = t[None] if t.ndim == 2 else t
+        return average_to(t.to(torch.float32), self.refine_shape)
+
+    def restart_static(self, frames=None, *, blur: float | None = None) -> None:
+        """Forget the scan, for a new one: its static background and its ghost.
+
+        *frames* from all over the same field, such as a search's, start the new
+        static background, so the scan's first patterns are divided by one that
+        holds no grain's bands rather than by the mean of its first few frames.
+
+        With *blur* (a Gaussian sigma, as a share of the pattern's width) the
+        static background is the frames' mean blurred, for a single crystal: its
+        frames all carry its bands, so their mean would divide them out. It then
+        has no ghost."""
+        self.background.restart(blur)
+        if frames is not None:
+            self.observe(frames)
+
+    def observe(self, frames) -> None:
+        """Add *frames* to the scan's static background."""
+        self.background.observe(self.pooled(frames))
+
+    def add_fits(self, phase, g, center_maps=None) -> None:
+        """Add the simulations of fitted orientations *g* ``(N, 3, 3)`` to the
+        scan's :attr:`ghost`."""
+        if self.background.blur:
+            return
+        if center_maps is not None:
+            g = g @ center_maps
+        self.background.add_fits(self.project(self.texture, g, phase, self.refine_directions))
+
+    def corrected(self, frames) -> torch.Tensor:
+        """*frames* ``(N, H, W)`` with the backgrounds removed as :meth:`prepare` removes
+        them, at their own resolution: what the indexer matches, to be looked at."""
+        t = torch.as_tensor(frames, device=self.device)
+        t = t[None] if t.ndim == 2 else t
+        return self.background.remove(t.to(torch.float32))
+
+    def prepare(self, frames) -> tuple[torch.Tensor, torch.Tensor]:
+        """*frames* ``(N, H, W)`` corrected and normalized, as ``(search, refine)``
+        rows ``(N, pixels)`` at the two shapes."""
+        t = self.background.remove(self.pooled(frames))
+        search = average_to(t, self.search_shape)
+        return normalize_rows(search.flatten(1)), normalize_rows(t.flatten(1))
+
+    def search(self, search_rows: torch.Tensor):
+        """The best dictionary entry for each pattern: ``(phase, g, score, entry)``,
+        *entry* indexing :attr:`rotations` (every phase's entries, in phase order).
+
+        The dictionary is scored :data:`SEARCH_CHUNK` entries at a time, keeping
+        each pattern's best so far, so the scores never outgrow one chunk."""
+        rows = search_rows.to(self.dtype)
+        score = torch.full((len(rows),), -math.inf, device=self.device)
+        entry = torch.zeros(len(rows), dtype=torch.int64, device=self.device)
+        for lo in range(0, self.dictionary.shape[1], SEARCH_CHUNK):
+            chunk_score, chunk_entry = (rows @ self.dictionary[:, lo : lo + SEARCH_CHUNK]).max(1)
+            better = chunk_score.float() > score
+            score = torch.where(better, chunk_score.float(), score)
+            entry = torch.where(better, chunk_entry + lo, entry)
+        return self.phase_of[entry], self.rotations[entry], score, entry
+
+    def refine(
+        self,
+        refine_rows: torch.Tensor,
+        phase: torch.Tensor,
+        g: torch.Tensor,
+        *,
+        steps=None,
+        directions=None,
+    ):
+        """Each orientation *g* ``(N, 3, 3)`` refined against its pattern:
+        ``(g, score)``, the score being the last level's at the seed. *steps*
+        (radians) default to the indexer's, *directions* to its pattern center's."""
+        steps = self.refine_steps if steps is None else steps
+        directions = self.refine_directions if directions is None else directions
+        score = torch.zeros(len(g), device=self.device)
+        for step in steps:
+            candidates = g[:, None] @ rotation_vector_to_matrix(self.stencil * step)
+            scores = self.score_candidates(candidates, phase, refine_rows, directions)
+            g = g @ rotation_vector_to_matrix(parabola_offset(scores) * step)
+            score = scores[:, 0]
+        return g, score
+
+    # --- the pattern center and the strain --------------------------------------
+
+    def refine_pattern_center(self, frames, **kwargs) -> tuple[float, float, float]:
+        """The pattern center that best fits the best-scored patterns of *frames*
+        ``(N, H, W)``, which the indexer then uses; see
+        :func:`~de_ebsd.index.live.calibration.refine_pattern_center`."""
+        return calibration.refine_pattern_center(self, frames, **kwargs)
+
+    def center_maps(self, pattern_centers) -> torch.Tensor:
+        """One ``(3, 3)`` map per pattern center; see
+        :func:`~de_ebsd.index.live.deformation.center_maps`."""
+        detector = replace(self.detector, shape=self.refine_shape)
+        return deformation.center_maps(detector, pattern_centers, self.device)
+
+    def refine_deformation(self, refine_rows, phase, g, **kwargs):
+        """Each orientation refined together with its pattern's strain; see
+        :func:`~de_ebsd.index.live.deformation.refine_deformation`."""
+        return deformation.refine_deformation(self, refine_rows, phase, g, **kwargs)
+
+    def strain_from_reference(self, frames, reference, **kwargs):
+        """The strain of *frames* relative to one real *reference* frame; see
+        :func:`~de_ebsd.index.live.deformation.strain_from_reference`."""
+        return deformation.strain_from_reference(self, frames, reference, **kwargs)
+
+    def index_with_strain(self, frames, pattern_centers=None):
+        """:meth:`index`, then each orientation refined together with its pattern's
+        strain (:meth:`refine_deformation`): ``(phase, g, strain, score, entry)``.
+        *pattern_centers* ``(N, 3)``, one per frame, default to the indexer's.
+
+        Each buffer's orientations join the scan's :attr:`ghost` before its
+        strains are fitted, as its frames joined the static background."""
+        maps = None if pattern_centers is None else self.center_maps(pattern_centers)
+        out = []
+        for lo in range(0, len(frames), BATCH):
+            search_rows, refine_rows = self.prepare(frames[lo : lo + BATCH])
+            phase, g, _, entry = self.search(search_rows)
+            g, _ = self.refine(refine_rows, phase, g)
+            batch_maps = None if maps is None else maps[lo : lo + BATCH]
+            self.add_fits(phase, g, batch_maps)
+            g, strain, score = self.refine_deformation(
+                refine_rows, phase, g, center_maps=batch_maps, ghost=self.ghost
+            )
+            out.append((phase, g, strain, score, entry))
+        return tuple(torch.cat(column) for column in zip(*out, strict=True))
+
+    def index(self, frames):
+        """*frames* ``(N, H, W)`` on the device, indexed: ``(phase, g, score, entry)``.
+        *entry* is the dictionary entry the search seeded the refine from (see
+        :meth:`search`): anything worked out once per entry, such as a color, stands
+        for the refined orientation to within the dictionary's resolution."""
+        out = []
+        for lo in range(0, len(frames), BATCH):
+            search_rows, refine_rows = self.prepare(frames[lo : lo + BATCH])
+            phase, g, score, entry = self.search(search_rows)
+            if self.refine_steps:
+                g, score = self.refine(refine_rows, phase, g)
+            out.append((phase, g, score, entry))
+        return tuple(torch.cat(column) for column in zip(*out, strict=True))

de_ebsd/index/preprocess.py

@@ -0,0 +1,195 @@
+"""Background removal for experimental patterns.
+
+Static (mean over the scan) and dynamic (Gaussian-blurred per pattern)
+background correction, batched on the GPU.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from de_ebsd.core.constants import DIVISION_FLOOR, SINGLE_PRECISION_FLOOR
+
+__all__ = [
+    "remove_static_background",
+    "remove_dynamic_background",
+    "remove_background",
+]
+
+
+def _as_stack(patterns, device, dtype):
+    x = torch.as_tensor(patterns, device=device, dtype=dtype)
+    if x.ndim == 2:
+        lead = ()
+    elif x.ndim == 3:
+        lead = (x.shape[0],)
+    elif x.ndim == 4:
+        lead = (x.shape[0], x.shape[1])
+    else:
+        raise ValueError(
+            f"patterns must be (H,W), (N,H,W) or (rows,cols,H,W); got shape {tuple(x.shape)}"
+        )
+    h, w = x.shape[-2], x.shape[-1]
+    stack = x.reshape(-1, h, w)
+
+    def restore(processed):
+        return processed.reshape(*lead, h, w)
+
+    return stack, restore
+
+
+def _rescale01(p):
+    lo = p.amin(dim=(-2, -1), keepdim=True)
+    hi = p.amax(dim=(-2, -1), keepdim=True)
+    return (p - lo) / (hi - lo).clamp_min(DIVISION_FLOOR)
+
+
+def _apply(pattern, background, operation):
+    if operation == "subtract":
+        return pattern - background
+    if operation == "divide":
+        return pattern / background.clamp_min(DIVISION_FLOOR)
+    raise ValueError(f"operation must be 'subtract' or 'divide', got {operation!r}")
+
+
+def normalize_rows(x: torch.Tensor) -> torch.Tensor:
+    """Each row of *x* shifted to zero mean and scaled to unit length, so a dot product
+    of two rows is their normalized cross-correlation."""
+    x = x - x.mean(-1, keepdim=True)
+    return x / x.norm(dim=-1, keepdim=True).clamp_min(SINGLE_PRECISION_FLOOR)
+
+
+def _gaussian_kernel_1d(sigma, truncate, device, dtype):
+    radius = int(truncate * float(sigma) + 0.5)
+    x = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
+    k = torch.exp(-0.5 * (x / float(sigma)) ** 2)
+    return k / k.sum(), radius
+
+
+def gaussian_blur(stack, sigma, truncate):
+    k, r = _gaussian_kernel_1d(sigma, truncate, stack.device, stack.dtype)
+    x = stack[:, None]  # (N, 1, H, W)
+    x = F.conv2d(F.pad(x, (r, r, 0, 0), mode="reflect"), k.view(1, 1, 1, -1))
+    x = F.conv2d(F.pad(x, (0, 0, r, r), mode="reflect"), k.view(1, 1, -1, 1))
+    return x[:, 0]
+
+
+def remove_static_background(
+    patterns, static_bg, operation="subtract", device=None, dtype=torch.float32
+):
+    """Remove the fixed detector/beam background from a stack of patterns.
+
+    Parameters
+    ----------
+    patterns
+        Detector images, ``(H, W)`` / ``(N, H, W)`` / ``(rows, cols, H, W)``, numpy or
+        torch. ``H``, ``W`` are the detector shape.
+    static_bg
+        The static background image ``(H, W)`` (e.g. a kikuchipy signal's
+        ``static_background``, or the scan mean), numpy or torch.
+    operation
+        ``"subtract"`` (default) or ``"divide"``.
+    device
+        Torch device to run on. ``None`` (default) keeps the patterns' device
+        (``"cuda"`` recommended for a whole scan).
+    dtype
+        Working floating dtype (default ``float32``).
+
+    Returns
+    -------
+    torch.Tensor
+        The cleaned patterns, same leading shape as the input, each rescaled to
+        ``[0, 1]`` (indexing is invariant to the per-pattern affine scale).
+    """
+    if device is None:
+        device = torch.as_tensor(patterns).device if torch.is_tensor(patterns) else "cpu"
+    stack, restore = _as_stack(patterns, device, dtype)
+    bg = torch.as_tensor(static_bg, device=stack.device, dtype=dtype)
+    if bg.shape != stack.shape[-2:]:
+        raise ValueError(f"static_bg shape {tuple(bg.shape)} != detector {tuple(stack.shape[-2:])}")
+    out = _rescale01(_apply(stack, bg[None], operation))
+    return restore(out)
+
+
+def remove_dynamic_background(
+    patterns, sigma=None, operation="subtract", truncate=4.0, device=None, dtype=torch.float32
+):
+    """Remove the per-pattern dynamic background (a Gaussian blur of each pattern).
+
+    Parameters
+    ----------
+    patterns
+        Detector images, ``(H, W)`` / ``(N, H, W)`` / ``(rows, cols, H, W)``, numpy or
+        torch.
+    sigma
+        Standard deviation of the Gaussian window. ``None`` (default) uses
+        ``width / 8`` (kikuchipy's default), the width being the number of columns.
+    operation
+        ``"subtract"`` (default) or ``"divide"``.
+    truncate
+        Truncate the Gaussian window at this many standard deviations (default ``4.0``,
+        matching scipy/kikuchipy).
+    device, dtype
+        As in :func:`remove_static_background`.
+
+    Returns
+    -------
+    torch.Tensor
+        The cleaned patterns, same leading shape as the input, each rescaled to
+        ``[0, 1]``.
+    """
+    if device is None:
+        device = torch.as_tensor(patterns).device if torch.is_tensor(patterns) else "cpu"
+    stack, restore = _as_stack(patterns, device, dtype)
+    if sigma is None:
+        sigma = stack.shape[-1] / 8.0
+    bg = gaussian_blur(stack, sigma, truncate)
+    out = _rescale01(_apply(stack, bg, operation))
+    return restore(out)
+
+
+def remove_background(
+    patterns,
+    static_bg=None,
+    dynamic=True,
+    operation="subtract",
+    sigma=None,
+    truncate=4.0,
+    device=None,
+    dtype=torch.float32,
+):
+    """Full EBSD preprocessing: static then dynamic background removal, on the GPU.
+
+    Convenience wrapper that runs :func:`remove_static_background` (when ``static_bg``
+    is given) followed by :func:`remove_dynamic_background` (when ``dynamic``), the
+    usual pre-indexing pipeline, all on ``device`` in one pass over the scan.
+
+    Parameters
+    ----------
+    patterns
+        The raw scan, ``(H, W)`` / ``(N, H, W)`` / ``(rows, cols, H, W)``, numpy or torch.
+    static_bg
+        The static background image ``(H, W)``; if ``None``, the static step is skipped.
+    dynamic
+        Whether to also remove the dynamic (per-pattern Gaussian) background
+        (default ``True``).
+    operation, sigma, truncate, device, dtype
+        Passed through to the two steps (see those functions).
+
+    Returns
+    -------
+    torch.Tensor
+        The cleaned patterns, same leading shape as the input, rescaled to ``[0, 1]``,
+        ready to hand to :meth:`~de_ebsd.index.dictionary.DictionaryIndexer.index`.
+    """
+    out = patterns
+    if static_bg is not None:
+        out = remove_static_background(
+            out, static_bg, operation=operation, device=device, dtype=dtype
+        )
+    if dynamic:
+        out = remove_dynamic_background(
+            out, sigma=sigma, operation=operation, truncate=truncate, device=device, dtype=dtype
+        )
+    return out

de_ebsd/index/refine/__init__.py

@@ -0,0 +1,7 @@
+"""Gradient refinement of orientation and pattern center."""
+
+from de_ebsd.index.refine.highres import refine_hr
+from de_ebsd.index.refine.orientation import refine_orientation
+from de_ebsd.index.refine.pattern_center import determine_pattern_center
+
+__all__ = ["determine_pattern_center", "refine_hr", "refine_orientation"]