lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/kernels/propagation/grid.py

@@ -0,0 +1,527 @@
+# The Clear BSD License
+#
+# Copyright (c) 2026 Tobias Heibges
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+#
+#      * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#
+#      * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#
+#      * Neither the name of the copyright holder nor the names of its
+#      contributors may be used to endorse or promote products derived from this
+#      software without specific prior written permission.
+#
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
+# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+GPU Kernels for GridInhomogeneousModel
+
+CUDA kernels for GPU-accelerated ray propagation through
+3D grid-based inhomogeneous materials using trilinear interpolation.
+"""
+
+import math
+
+# GPU support is optional
+try:
+    from numba import cuda
+
+    HAS_CUDA = True
+except ImportError:
+
+    class _FakeCuda:
+        """Fake cuda module for when numba is not installed."""
+
+        @staticmethod
+        def jit(*args, **kwargs):
+            """Return a no-op decorator."""
+
+            def decorator(func):
+                return func
+
+            if args and callable(args[0]):
+                return args[0]
+            return decorator
+
+        @staticmethod
+        def is_available():
+            return False
+
+        @staticmethod
+        def grid(n):
+            return 0
+
+        @staticmethod
+        def synchronize():
+            pass
+
+    cuda = _FakeCuda()  # type: ignore[assignment]
+    HAS_CUDA = False
+
+from ..device_functions import device_euler_step
+from ..registry import PropagationKernelID, register_kernel
+
+SPEED_OF_LIGHT = 299792458.0
+
+
+@cuda.jit(device=True)
+def _device_trilinear_interpolate(
+    x: float,
+    y: float,
+    z: float,
+    grid,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    dx: float,
+    dy: float,
+    dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+) -> float:
+    """Trilinear interpolation in 3D grid."""
+    xn = (x - x_min) / dx
+    yn = (y - y_min) / dy
+    zn = (z - z_min) / dz
+
+    xn = max(0.0, min(xn, nx - 1.001))
+    yn = max(0.0, min(yn, ny - 1.001))
+    zn = max(0.0, min(zn, nz - 1.001))
+
+    x0 = int(xn)
+    y0 = int(yn)
+    z0 = int(zn)
+
+    x1 = min(x0 + 1, nx - 1)
+    y1 = min(y0 + 1, ny - 1)
+    z1 = min(z0 + 1, nz - 1)
+
+    xd = xn - x0
+    yd = yn - y0
+    zd = zn - z0
+
+    c000 = grid[x0, y0, z0]
+    c001 = grid[x0, y0, z1]
+    c010 = grid[x0, y1, z0]
+    c011 = grid[x0, y1, z1]
+    c100 = grid[x1, y0, z0]
+    c101 = grid[x1, y0, z1]
+    c110 = grid[x1, y1, z0]
+    c111 = grid[x1, y1, z1]
+
+    c00 = c000 * (1 - xd) + c100 * xd
+    c01 = c001 * (1 - xd) + c101 * xd
+    c10 = c010 * (1 - xd) + c110 * xd
+    c11 = c011 * (1 - xd) + c111 * xd
+
+    c0 = c00 * (1 - yd) + c10 * yd
+    c1 = c01 * (1 - yd) + c11 * yd
+
+    return c0 * (1 - zd) + c1 * zd
+
+
+@cuda.jit(device=True)
+def _device_grid_n_and_gradient(
+    x: float,
+    y: float,
+    z: float,
+    n_grid,
+    grad_grid,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    dx: float,
+    dy: float,
+    dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+) -> tuple[float, float, float, float]:
+    """Get n and gradient from GridInhomogeneousModel."""
+    n = _device_trilinear_interpolate(
+        x, y, z, n_grid, x_min, y_min, z_min, dx, dy, dz, nx, ny, nz
+    )
+
+    # Normalize coordinates for gradient interpolation
+    xn = (x - x_min) / dx
+    yn = (y - y_min) / dy
+    zn = (z - z_min) / dz
+
+    xn = max(0.0, min(xn, nx - 1.001))
+    yn = max(0.0, min(yn, ny - 1.001))
+    zn = max(0.0, min(zn, nz - 1.001))
+
+    x0 = int(xn)
+    y0 = int(yn)
+    z0 = int(zn)
+
+    x1 = min(x0 + 1, nx - 1)
+    y1 = min(y0 + 1, ny - 1)
+    z1 = min(z0 + 1, nz - 1)
+
+    xd = xn - x0
+    yd = yn - y0
+    zd = zn - z0
+
+    # Interpolate grad_x
+    c000 = grad_grid[x0, y0, z0, 0]
+    c001 = grad_grid[x0, y0, z1, 0]
+    c010 = grad_grid[x0, y1, z0, 0]
+    c011 = grad_grid[x0, y1, z1, 0]
+    c100 = grad_grid[x1, y0, z0, 0]
+    c101 = grad_grid[x1, y0, z1, 0]
+    c110 = grad_grid[x1, y1, z0, 0]
+    c111 = grad_grid[x1, y1, z1, 0]
+    c00 = c000 * (1 - xd) + c100 * xd
+    c01 = c001 * (1 - xd) + c101 * xd
+    c10 = c010 * (1 - xd) + c110 * xd
+    c11 = c011 * (1 - xd) + c111 * xd
+    c0 = c00 * (1 - yd) + c10 * yd
+    c1 = c01 * (1 - yd) + c11 * yd
+    grad_x = c0 * (1 - zd) + c1 * zd
+
+    # Interpolate grad_y
+    c000 = grad_grid[x0, y0, z0, 1]
+    c001 = grad_grid[x0, y0, z1, 1]
+    c010 = grad_grid[x0, y1, z0, 1]
+    c011 = grad_grid[x0, y1, z1, 1]
+    c100 = grad_grid[x1, y0, z0, 1]
+    c101 = grad_grid[x1, y0, z1, 1]
+    c110 = grad_grid[x1, y1, z0, 1]
+    c111 = grad_grid[x1, y1, z1, 1]
+    c00 = c000 * (1 - xd) + c100 * xd
+    c01 = c001 * (1 - xd) + c101 * xd
+    c10 = c010 * (1 - xd) + c110 * xd
+    c11 = c011 * (1 - xd) + c111 * xd
+    c0 = c00 * (1 - yd) + c10 * yd
+    c1 = c01 * (1 - yd) + c11 * yd
+    grad_y = c0 * (1 - zd) + c1 * zd
+
+    # Interpolate grad_z
+    c000 = grad_grid[x0, y0, z0, 2]
+    c001 = grad_grid[x0, y0, z1, 2]
+    c010 = grad_grid[x0, y1, z0, 2]
+    c011 = grad_grid[x0, y1, z1, 2]
+    c100 = grad_grid[x1, y0, z0, 2]
+    c101 = grad_grid[x1, y0, z1, 2]
+    c110 = grad_grid[x1, y1, z0, 2]
+    c111 = grad_grid[x1, y1, z1, 2]
+    c00 = c000 * (1 - xd) + c100 * xd
+    c01 = c001 * (1 - xd) + c101 * xd
+    c10 = c010 * (1 - xd) + c110 * xd
+    c11 = c011 * (1 - xd) + c111 * xd
+    c0 = c00 * (1 - yd) + c10 * yd
+    c1 = c01 * (1 - yd) + c11 * yd
+    grad_z = c0 * (1 - zd) + c1 * zd
+
+    return n, grad_x, grad_y, grad_z
+
+
+@cuda.jit(device=True)
+def _device_grid_euler_step(
+    x: float,
+    y: float,
+    z: float,
+    dir_x: float,
+    dir_y: float,
+    dir_z: float,
+    step_size: float,
+    n_grid,
+    grad_grid,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    dx: float,
+    dy: float,
+    dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+) -> tuple[float, float, float, float, float, float, float]:
+    """Euler step for GridInhomogeneousModel."""
+    n, grad_x, grad_y, grad_z = _device_grid_n_and_gradient(
+        x, y, z, n_grid, grad_grid, x_min, y_min, z_min, dx, dy, dz, nx, ny, nz
+    )
+
+    new_x, new_y, new_z, new_dx, new_dy, new_dz = device_euler_step(
+        x, y, z, dir_x, dir_y, dir_z, n, grad_x, grad_y, grad_z, step_size
+    )
+
+    return new_x, new_y, new_z, new_dx, new_dy, new_dz, n
+
+
+@cuda.jit(device=True)
+def _device_grid_rk4_step(
+    x: float,
+    y: float,
+    z: float,
+    dir_x: float,
+    dir_y: float,
+    dir_z: float,
+    step_size: float,
+    n_grid,
+    grad_grid,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    grid_dx: float,
+    grid_dy: float,
+    grid_dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+) -> tuple[float, float, float, float, float, float, float]:
+    """RK4 step for GridInhomogeneousModel."""
+    h = step_size
+    h2 = h / 2.0
+
+    def get_n_and_kappa(px, py, pz, dx, dy, dz):
+        n, gx, gy, gz = _device_grid_n_and_gradient(
+            px,
+            py,
+            pz,
+            n_grid,
+            grad_grid,
+            x_min,
+            y_min,
+            z_min,
+            grid_dx,
+            grid_dy,
+            grid_dz,
+            nx,
+            ny,
+            nz,
+        )
+        dot = dx * gx + dy * gy + dz * gz
+        kx = (gx - dot * dx) / n
+        ky = (gy - dot * dy) / n
+        kz = (gz - dot * dz) / n
+        return n, kx, ky, kz
+
+    def normalize(dx, dy, dz):
+        norm = math.sqrt(dx * dx + dy * dy + dz * dz)
+        if norm < 1e-12:
+            norm = 1.0
+        return dx / norm, dy / norm, dz / norm
+
+    # k1
+    n0, kx1, ky1, kz1 = get_n_and_kappa(x, y, z, dir_x, dir_y, dir_z)
+    k1_rx, k1_ry, k1_rz = dir_x, dir_y, dir_z
+    k1_dx, k1_dy, k1_dz = kx1, ky1, kz1
+
+    # k2
+    px = x + h2 * k1_rx
+    py = y + h2 * k1_ry
+    pz = z + h2 * k1_rz
+    dx, dy, dz = normalize(dir_x + h2 * k1_dx, dir_y + h2 * k1_dy, dir_z + h2 * k1_dz)
+    n1, kx2, ky2, kz2 = get_n_and_kappa(px, py, pz, dx, dy, dz)
+    k2_rx, k2_ry, k2_rz = dx, dy, dz
+    k2_dx, k2_dy, k2_dz = kx2, ky2, kz2
+
+    # k3
+    px = x + h2 * k2_rx
+    py = y + h2 * k2_ry
+    pz = z + h2 * k2_rz
+    dx, dy, dz = normalize(dir_x + h2 * k2_dx, dir_y + h2 * k2_dy, dir_z + h2 * k2_dz)
+    n2, kx3, ky3, kz3 = get_n_and_kappa(px, py, pz, dx, dy, dz)
+    k3_rx, k3_ry, k3_rz = dx, dy, dz
+    k3_dx, k3_dy, k3_dz = kx3, ky3, kz3
+
+    # k4
+    px = x + h * k3_rx
+    py = y + h * k3_ry
+    pz = z + h * k3_rz
+    dx, dy, dz = normalize(dir_x + h * k3_dx, dir_y + h * k3_dy, dir_z + h * k3_dz)
+    n3, kx4, ky4, kz4 = get_n_and_kappa(px, py, pz, dx, dy, dz)
+
+    # Final RK4 combination
+    new_x = x + (h / 6.0) * (k1_rx + 2 * k2_rx + 2 * k3_rx + dx)
+    new_y = y + (h / 6.0) * (k1_ry + 2 * k2_ry + 2 * k3_ry + dy)
+    new_z = z + (h / 6.0) * (k1_rz + 2 * k2_rz + 2 * k3_rz + dz)
+
+    new_dx = dir_x + (h / 6.0) * (k1_dx + 2 * k2_dx + 2 * k3_dx + kx4)
+    new_dy = dir_y + (h / 6.0) * (k1_dy + 2 * k2_dy + 2 * k3_dy + ky4)
+    new_dz = dir_z + (h / 6.0) * (k1_dz + 2 * k2_dz + 2 * k3_dz + kz4)
+
+    new_dx, new_dy, new_dz = normalize(new_dx, new_dy, new_dz)
+
+    n_avg = (n0 + 4 * n1 + n2) / 6.0
+
+    return new_x, new_y, new_z, new_dx, new_dy, new_dz, n_avg
+
+
+@register_kernel(PropagationKernelID.GRID_EULER)
+@cuda.jit
+def _kernel_grid_inhomogeneous_euler(
+    positions,
+    directions,
+    active,
+    geo_path,
+    opt_path,
+    acc_time,
+    step_size: float,
+    num_steps: int,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    grid_dx: float,
+    grid_dy: float,
+    grid_dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+    n_grid,
+    grad_grid,
+    wavelength: float,
+):
+    """GPU kernel for GridInhomogeneousModel using Euler integration."""
+    c = SPEED_OF_LIGHT
+
+    idx = cuda.grid(1)
+    if idx >= positions.shape[0]:
+        return
+    if not active[idx]:
+        return
+
+    x = positions[idx, 0]
+    y = positions[idx, 1]
+    z = positions[idx, 2]
+    dx = directions[idx, 0]
+    dy = directions[idx, 1]
+    dz = directions[idx, 2]
+    gp = geo_path[idx]
+    op = opt_path[idx]
+    at = acc_time[idx]
+
+    for _ in range(num_steps):
+        x, y, z, dx, dy, dz, n = _device_grid_euler_step(
+            x,
+            y,
+            z,
+            dx,
+            dy,
+            dz,
+            step_size,
+            n_grid,
+            grad_grid,
+            x_min,
+            y_min,
+            z_min,
+            grid_dx,
+            grid_dy,
+            grid_dz,
+            nx,
+            ny,
+            nz,
+        )
+        gp += step_size
+        op += n * step_size
+        at += n * step_size / c
+
+    positions[idx, 0] = x
+    positions[idx, 1] = y
+    positions[idx, 2] = z
+    directions[idx, 0] = dx
+    directions[idx, 1] = dy
+    directions[idx, 2] = dz
+    geo_path[idx] = gp
+    opt_path[idx] = op
+    acc_time[idx] = at
+
+
+@register_kernel(PropagationKernelID.GRID_RK4)
+@cuda.jit
+def _kernel_grid_inhomogeneous_rk4(
+    positions,
+    directions,
+    active,
+    geo_path,
+    opt_path,
+    acc_time,
+    step_size: float,
+    num_steps: int,
+    x_min: float,
+    y_min: float,
+    z_min: float,
+    grid_dx: float,
+    grid_dy: float,
+    grid_dz: float,
+    nx: int,
+    ny: int,
+    nz: int,
+    n_grid,
+    grad_grid,
+    wavelength: float,
+):
+    """GPU kernel for GridInhomogeneousModel using RK4 integration."""
+    c = SPEED_OF_LIGHT
+
+    idx = cuda.grid(1)
+    if idx >= positions.shape[0]:
+        return
+    if not active[idx]:
+        return
+
+    x = positions[idx, 0]
+    y = positions[idx, 1]
+    z = positions[idx, 2]
+    dx = directions[idx, 0]
+    dy = directions[idx, 1]
+    dz = directions[idx, 2]
+    gp = geo_path[idx]
+    op = opt_path[idx]
+    at = acc_time[idx]
+
+    for _ in range(num_steps):
+        x, y, z, dx, dy, dz, n_avg = _device_grid_rk4_step(
+            x,
+            y,
+            z,
+            dx,
+            dy,
+            dz,
+            step_size,
+            n_grid,
+            grad_grid,
+            x_min,
+            y_min,
+            z_min,
+            grid_dx,
+            grid_dy,
+            grid_dz,
+            nx,
+            ny,
+            nz,
+        )
+        gp += step_size
+        op += n_avg * step_size
+        at += n_avg * step_size / c
+
+    positions[idx, 0] = x
+    positions[idx, 1] = y
+    positions[idx, 2] = z
+    directions[idx, 0] = dx
+    directions[idx, 1] = dy
+    directions[idx, 2] = dz
+    geo_path[idx] = gp
+    opt_path[idx] = op
+    acc_time[idx] = at

lsurf/propagation/kernels/propagation/protocol.py

@@ -0,0 +1,105 @@
+# The Clear BSD License
+#
+# Copyright (c) 2026 Tobias Heibges
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+#
+#      * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#
+#      * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#
+#      * Neither the name of the copyright holder nor the names of its
+#      contributors may be used to endorse or promote products derived from this
+#      software without specific prior written permission.
+#
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
+# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+Propagation Kernel Protocol
+
+Defines the interface for material propagation kernels that integrate
+rays through inhomogeneous materials.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from numba.cuda.devicearray import DeviceNDArray
+
+
+@runtime_checkable
+class PropagationKernelProtocol(Protocol):
+    """
+    Protocol for material propagation kernels.
+
+    Propagation kernels integrate rays through materials with spatially-varying
+    refractive index. They update ray positions, directions, and accumulated
+    path lengths based on the material's optical properties.
+
+    The kernel signature varies depending on the material type:
+    - Simple: 1D LUT interpolation (altitude only)
+    - Spectral: 2D LUT interpolation (altitude × wavelength)
+    - Grid: 3D trilinear interpolation
+
+    Common Parameters
+    -----------------
+    positions : DeviceNDArray, shape (N, 3)
+        Ray positions in meters (in/out)
+    directions : DeviceNDArray, shape (N, 3)
+        Ray direction unit vectors (in/out)
+    active : DeviceNDArray, shape (N,)
+        Boolean mask of active rays
+    geo_path : DeviceNDArray, shape (N,)
+        Accumulated geometric path length in meters (in/out)
+    opt_path : DeviceNDArray, shape (N,)
+        Accumulated optical path length (n·ds integral) (in/out)
+    acc_time : DeviceNDArray, shape (N,)
+        Accumulated travel time in seconds (in/out)
+    step_size : float
+        Integration step size in meters
+    num_steps : int
+        Number of integration steps to perform
+
+    Notes
+    -----
+    Kernels are registered with the central registry using @register_kernel
+    decorator and PropagationKernelID enum values.
+
+    The actual signature includes material-specific parameters (LUT arrays,
+    center coordinates, etc.) that vary by kernel type.
+    """
+
+    def __call__(
+        self,
+        positions: DeviceNDArray,
+        directions: DeviceNDArray,
+        active: DeviceNDArray,
+        geo_path: DeviceNDArray,
+        opt_path: DeviceNDArray,
+        acc_time: DeviceNDArray,
+        step_size: float,
+        num_steps: int,
+        *args,
+        **kwargs,
+    ) -> None:
+        """Execute the propagation kernel."""
+        ...