de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/index/_kernels.py

@@ -0,0 +1,222 @@
+"""Fused CUDA kernels for live indexing: a master pattern sampled at rotated
+detector directions in one pass, compiled at run time by CuPy.
+
+Projecting through torch operations writes every intermediate -- the rotated
+directions, their Lambert coordinates, the sampled values -- to GPU memory and
+reads it back, about 390 M samples/s on a Titan X. Here each thread rotates its
+direction, maps it to the square Lambert grid and interpolates the master in
+registers, reading only the master (small enough to stay in L2) and writing one
+value; :func:`score` does not even write that, reducing each candidate's samples
+straight to its normalized cross-correlation with its pattern.
+
+The arithmetic is :func:`de_ebsd.core.lambert.sphere_to_square` followed by
+``grid_sample(mode="bilinear", align_corners=True, padding_mode="border")``, as
+:meth:`MasterPattern.sample` does it; ``tests/index/test_live.py`` compares them.
+CuPy is optional: without it (or off CUDA) the torch path is used.
+"""
+
+from __future__ import annotations
+
+import functools
+import importlib.util
+
+import torch
+
+_SOURCE = r"""
+#define SQRT_PI_HALF 1.2533141373155003f
+#define SQRT_PI_OVER_2 0.886226925452758f
+#define TWO_OVER_SQRT_PI 1.1283791670955126f
+
+__device__ __forceinline__ float sample_master(
+    const float* __restrict__ texture, int base, int n, float vx, float vy, float vz)
+{
+    float inverse = rsqrtf(fmaxf(vx * vx + vy * vy + vz * vz, 1e-24f));
+    vx *= inverse; vy *= inverse; vz *= inverse;
+    int channel = base + (vz < 0.f ? 1 : 0);
+    float abs_z = fabsf(vz);
+    float lx = 0.f, ly = 0.f;
+    if (abs_z < 1.f) {
+        float root = sqrtf(fmaxf(2.f * (1.f - abs_z), 1e-30f));
+        if (fabsf(vy) <= fabsf(vx)) {
+            float sign = vx >= 0.f ? 1.f : -1.f;
+            float safe = fabsf(vx) < 1e-12f ? 1e-12f : vx;
+            lx = sign * root * SQRT_PI_OVER_2;
+            ly = sign * root * TWO_OVER_SQRT_PI * atanf(vy / safe);
+        } else {
+            float sign = vy >= 0.f ? 1.f : -1.f;
+            float safe = fabsf(vy) < 1e-12f ? 1e-12f : vy;
+            lx = sign * root * TWO_OVER_SQRT_PI * atanf(vx / safe);
+            ly = sign * root * SQRT_PI_OVER_2;
+        }
+    }
+    float last = (float)(n - 1);
+    float px = fminf(fmaxf((lx / SQRT_PI_HALF + 1.f) * 0.5f * last, 0.f), last);
+    float py = fminf(fmaxf((ly / SQRT_PI_HALF + 1.f) * 0.5f * last, 0.f), last);
+    int x0 = (int)floorf(px), y0 = (int)floorf(py);
+    int x1 = min(x0 + 1, n - 1), y1 = min(y0 + 1, n - 1);
+    float fx = px - x0, fy = py - y0;
+    const float* t = texture + (size_t)channel * n * n;
+    float top = t[y0 * n + x0] * (1.f - fx) + t[y0 * n + x1] * fx;
+    float bottom = t[y1 * n + x0] * (1.f - fx) + t[y1 * n + x1] * fx;
+    return top * (1.f - fy) + bottom * fy;
+}
+
+__device__ __forceinline__ float sample_rotated(
+    const float* __restrict__ texture, int base, int n,
+    const float* __restrict__ r, const float* __restrict__ d)
+{
+    float dx = d[0], dy = d[1], dz = d[2];
+    return sample_master(texture, base, n,
+                         r[0] * dx + r[1] * dy + r[2] * dz,
+                         r[3] * dx + r[4] * dy + r[5] * dz,
+                         r[6] * dx + r[7] * dy + r[8] * dz);
+}
+
+// out[m, p]: rotation m's pattern at direction p, from channel pair bases[m].
+extern "C" __global__ void project(
+    const float* __restrict__ texture, int n,
+    const float* __restrict__ rotations, const int* __restrict__ bases,
+    const float* __restrict__ directions, int pixels, long long count,
+    float* __restrict__ out)
+{
+    long long index = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (index >= count) return;
+    long long m = index / pixels;
+    int p = (int)(index - m * pixels);
+    out[index] = sample_rotated(texture, bases[m], n, rotations + m * 9, directions + p * 3);
+}
+
+// scores[m]: the normalized cross-correlation of rotation m's pattern with
+// rows[m / per_row], a row already zero-mean and unit-norm. One block per m.
+extern "C" __global__ void score(
+    const float* __restrict__ texture, int n,
+    const float* __restrict__ rotations, const int* __restrict__ bases,
+    const float* __restrict__ directions, int pixels,
+    const float* __restrict__ rows, int per_row,
+    float* __restrict__ scores)
+{
+    extern __shared__ float shared[];
+    long long m = blockIdx.x;
+    const float* r = rotations + m * 9;
+    const float* row = rows + (m / per_row) * (long long)pixels;
+    int base = bases[m];
+    float sum = 0.f, squares = 0.f, product = 0.f;
+    for (int p = threadIdx.x; p < pixels; p += blockDim.x) {
+        float s = sample_rotated(texture, base, n, r, directions + p * 3);
+        sum += s; squares += s * s; product += s * row[p];
+    }
+    float* a = shared; float* b = shared + blockDim.x; float* c = shared + 2 * blockDim.x;
+    a[threadIdx.x] = sum; b[threadIdx.x] = squares; c[threadIdx.x] = product;
+    __syncthreads();
+    for (int half = blockDim.x / 2; half > 0; half >>= 1) {
+        if (threadIdx.x < half) {
+            a[threadIdx.x] += a[threadIdx.x + half];
+            b[threadIdx.x] += b[threadIdx.x + half];
+            c[threadIdx.x] += c[threadIdx.x + half];
+        }
+        __syncthreads();
+    }
+    if (threadIdx.x == 0) {
+        float variance = b[0] - a[0] * a[0] / pixels;
+        scores[m] = c[0] * rsqrtf(fmaxf(variance, 1e-12f));
+    }
+}
+"""
+
+#: Threads per block of :func:`score`; a power of two for its reduction.
+SCORE_THREADS = 256
+
+
+def available(device) -> bool:
+    """Whether the fused kernels can run on *device*."""
+    return torch.device(device).type == "cuda" and importlib.util.find_spec("cupy") is not None
+
+
+@functools.cache
+def _module(device_index: int):
+    import cupy
+
+    with cupy.cuda.Device(device_index):
+        return cupy.RawModule(code=_SOURCE, options=("--use_fast_math",))
+
+
+def _cupy(tensor: torch.Tensor):
+    import cupy
+
+    return cupy.from_dlpack(tensor.contiguous())
+
+
+def _launch(name, grid, block, arguments, shared_mem=0):
+    import cupy
+
+    device = torch.cuda.current_device()
+    stream = cupy.cuda.ExternalStream(torch.cuda.current_stream().cuda_stream)
+    with cupy.cuda.Device(device), stream:
+        _module(device).get_function(name)(grid, block, arguments, shared_mem=shared_mem)
+
+
+def project(
+    texture: torch.Tensor, rotations: torch.Tensor, bases: torch.Tensor, directions: torch.Tensor
+) -> torch.Tensor:
+    """Patterns ``(M, pixels)``: *texture* ``(1, channels, n, n)`` sampled at
+    *directions* ``(pixels, 3)`` turned by *rotations* ``(M, 3, 3)``, each from the
+    hemisphere pair starting at its channel ``bases[m]``."""
+    import cupy
+
+    m, pixels = len(rotations), len(directions)
+    out = torch.empty(m, pixels, device=rotations.device, dtype=torch.float32)
+    count = m * pixels
+    threads = 256
+    texture = texture.contiguous()
+    _launch(
+        "project",
+        ((count + threads - 1) // threads,),
+        (threads,),
+        (
+            _cupy(texture),
+            cupy.int32(texture.shape[-1]),
+            _cupy(rotations.float()),
+            _cupy(bases.to(torch.int32)),
+            _cupy(directions.float()),
+            cupy.int32(pixels),
+            cupy.int64(count),
+            _cupy(out),
+        ),
+    )
+    return out
+
+
+def score(
+    texture: torch.Tensor,
+    rotations: torch.Tensor,
+    bases: torch.Tensor,
+    directions: torch.Tensor,
+    rows: torch.Tensor,
+    per_row: int,
+) -> torch.Tensor:
+    """Scores ``(M,)``: the normalized cross-correlation of each rotation's
+    pattern (as :func:`project`) with ``rows[m // per_row]``, rows being
+    zero-mean, unit-norm ``(N, pixels)``."""
+    import cupy
+
+    m = len(rotations)
+    out = torch.empty(m, device=rotations.device, dtype=torch.float32)
+    texture = texture.contiguous()
+    _launch(
+        "score",
+        (m,),
+        (SCORE_THREADS,),
+        (
+            _cupy(texture),
+            cupy.int32(texture.shape[-1]),
+            _cupy(rotations.float()),
+            _cupy(bases.to(torch.int32)),
+            _cupy(directions.float()),
+            cupy.int32(len(directions)),
+            _cupy(rows.float()),
+            cupy.int32(per_row),
+            _cupy(out),
+        ),
+        shared_mem=3 * SCORE_THREADS * 4,
+    )
+    return out

de_ebsd/index/backproject.py

@@ -0,0 +1,171 @@
+"""Back-project a detector pattern onto the unit sphere.
+
+The inverse of the forward projection: each sphere direction is mapped to a
+fractional detector pixel and bilinearly sampled, giving E(n) on a supplied
+direction set plus the field-of-view mask. Differentiable in the pattern
+center, which is what makes PC refinement possible.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from de_ebsd.core.geometry import sample_to_detector_tensor
+from de_ebsd.sim.detector import EBSDDetector
+
+__all__ = ["directions_to_pixels", "pattern_to_sphere"]
+
+
+def _pc_components(
+    detector: EBSDDetector,
+    pattern_center: torch.Tensor | None,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if pattern_center is None:
+        pattern_center = torch.as_tensor(detector.pattern_center, device=device, dtype=dtype)
+    else:
+        if pattern_center.shape[-1] != 3:
+            raise ValueError(
+                f"pattern_center must have last dim 3, got shape {tuple(pattern_center.shape)}"
+            )
+        pattern_center = pattern_center.to(device=device, dtype=dtype)
+    return pattern_center[0], pattern_center[1], pattern_center[2]
+
+
+def directions_to_pixels(
+    directions: torch.Tensor,
+    detector: EBSDDetector,
+    pattern_center: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Map sphere directions to fractional detector pixels (inverse of the DC map).
+
+    Parameters
+    ----------
+    directions
+        Unit vectors ``(M, 3)`` in the sample frame (need not be exactly normalized;
+        the gnomonic ratios are scale-invariant).
+    detector
+        The ``de_ebsd.sim`` :class:`~de_ebsd.sim.detector.EBSDDetector` geometry.
+    pattern_center
+        Optional pattern-center override ``(3,)`` (``pcx, pcy, pcz``). When supplied
+        with ``requires_grad``, the returned pixel coordinates are differentiable in
+        the pattern center. Defaults to ``detector.pattern_center``.
+
+    Returns
+    -------
+    col, row
+        Fractional pixel coordinates ``(M,)``, where integer values land on pixel
+        centers (``col = 0`` is the leftmost column, ``row = 0`` the top row) --
+        ready for :func:`torch.nn.functional.grid_sample` with ``align_corners=True``.
+    in_fov
+        Boolean field-of-view mask ``(M,)``: ``True`` where the direction points at
+        the detector (``d_z > 0``) and lands within the gnomonic bounds.
+    """
+    device = directions.device
+    om = sample_to_detector_tensor(detector, device, directions.dtype)
+    dtype = directions.dtype
+    pcx, pcy, pcz = _pc_components(detector, pattern_center, device, dtype)
+
+    # detector-frame components d = (n . X_d, n . Y_d, n . Z_d) = directions @ Om.T
+    comp = directions @ om.transpose(-2, -1)
+    d_x, d_y, d_z = comp[..., 0], comp[..., 1], comp[..., 2]
+    # gnomonic projection back onto the detector plane (pattern-center independent)
+    d_z_safe = torch.where(d_z.abs() < 1e-12, torch.full_like(d_z, 1e-12), d_z)
+    gn_x = d_x / d_z_safe
+    gn_y = d_y / d_z_safe
+
+    # gnomonic bounds and node scales from the pattern center (torch, differentiable),
+    # matching EBSDDetector.gnomonic_bounds / direction_cosines exactly
+    aspect = detector.ncols / detector.nrows
+    x_min = -aspect * pcx / pcz
+    x_max = aspect * (1.0 - pcx) / pcz
+    y_min = -(1.0 - pcy) / pcz
+    y_max = pcy / pcz
+    x_scale = (x_max - x_min) / detector.ncols
+    y_scale = (y_max - y_min) / detector.nrows
+
+    col = (gn_x - x_min) / x_scale - 0.5
+    row = (y_max - gn_y) / y_scale - 0.5
+
+    in_fov = (d_z > 0) & (gn_x >= x_min) & (gn_x <= x_max) & (gn_y >= y_min) & (gn_y <= y_max)
+    return col, row, in_fov
+
+
+def pattern_to_sphere(
+    patterns: torch.Tensor,
+    detector: EBSDDetector,
+    directions: torch.Tensor,
+    pattern_center: torch.Tensor | None = None,
+    padding_mode: str = "border",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Back-project detector ``patterns`` onto ``directions`` on the unit sphere.
+
+    Each sphere direction reads the detector pixel that scatters along it (no crystal
+    orientation is applied), bilinearly interpolating the
+    pattern. The result is ``E(n)`` sampled on ``directions``, ready to feed
+    :meth:`de_ebsd.index.spherical.harmonics.SphericalHarmonicTransform.forward` (reshape the
+    ``(N_theta, N_phi)`` grid directions to ``(M, 3)`` and the samples back).
+
+    Parameters
+    ----------
+    patterns
+        Detector images ``(..., H, W)`` (``H = detector.nrows``, ``W = detector.ncols``),
+        batched over arbitrary leading dimensions. Real floating dtype.
+    detector
+        The ``de_ebsd.sim`` :class:`~de_ebsd.sim.detector.EBSDDetector` geometry
+        (fixed pattern center). Its pixel -> direction map is inverted here.
+    directions
+        Target sphere directions ``(M, 3)`` in the sample/detector frame, e.g.
+        ``SphericalHarmonicTransform.grid_directions().reshape(-1, 3)``.
+    pattern_center
+        Optional pattern-center override ``(3,)`` (``pcx, pcy, pcz``); pass a tensor
+        with ``requires_grad`` to differentiate the samples in the pattern center
+        (I6 refinement). Defaults to ``detector.pattern_center``.
+    padding_mode
+        ``grid_sample`` padding for directions that fall (just) outside the pixel
+        grid; ``"border"`` (default) clamps to the edge. Off-FOV samples are flagged
+        by the returned mask regardless.
+
+    Returns
+    -------
+    samples
+        Sphere samples ``(..., M)`` (same leading dims and dtype as ``patterns``).
+        Entries outside the field of view are the ``padding_mode`` read-off and
+        should be excluded via ``mask``. Differentiable in ``patterns`` (and, via
+        ``pattern_center``, in the pattern center).
+    mask
+        Field-of-view weight ``(M,)`` in ``patterns``' dtype: ``1`` where the
+        direction lands on the detector, ``0`` otherwise (the detector's gnomonic
+        bounds). Shared across the batch since the geometry is.
+    """
+    if patterns.shape[-2:] != (detector.nrows, detector.ncols):
+        raise ValueError(
+            f"patterns trailing shape {tuple(patterns.shape[-2:])} != detector "
+            f"{(detector.nrows, detector.ncols)}"
+        )
+    if directions.shape[-1] != 3 or directions.ndim != 2:
+        raise ValueError(f"directions must have shape (M, 3), got {tuple(directions.shape)}")
+
+    device = patterns.device
+    dirs = directions.to(device=device, dtype=patterns.dtype)
+    col, row, in_fov = directions_to_pixels(dirs, detector, pattern_center)
+
+    # normalized sampling coordinates for grid_sample(align_corners=True):
+    # col in [0, W-1] -> gx in [-1, 1]; row in [0, H-1] -> gy in [-1, 1]
+    gx = 2.0 * col / (detector.ncols - 1) - 1.0
+    gy = 2.0 * row / (detector.nrows - 1) - 1.0
+    grid = torch.stack([gx, gy], dim=-1)  # (M, 2)
+
+    m = dirs.shape[0]
+    leading = patterns.shape[:-2]
+    batch = patterns.reshape(-1, 1, detector.nrows, detector.ncols)
+    n = batch.shape[0]
+    grid = grid.view(1, 1, m, 2).expand(n, 1, m, 2)
+    sampled = F.grid_sample(
+        batch, grid, mode="bilinear", align_corners=True, padding_mode=padding_mode
+    )  # (N, 1, 1, M)
+    samples = sampled.reshape(*leading, m)
+    mask = in_fov.to(patterns.dtype)
+    return samples, mask