lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/kernels/propagation/simple.py

@@ -0,0 +1,460 @@
+# 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 SimpleInhomogeneousModel
+
+CUDA kernels for GPU-accelerated ray propagation through
+radially-symmetric inhomogeneous materials using 1D lookup table 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."""
+
+        class devicearray:
+            """Fake devicearray submodule."""
+
+            DeviceNDArray = object
+
+        @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_lut_interpolate(
+    value: float,
+    lut: cuda.devicearray.DeviceNDArray,
+    min_val: float,
+    delta: float,
+    lut_size: int,
+) -> float:
+    """
+    Linear interpolation in 1D lookup table.
+
+    Parameters
+    ----------
+    value : float
+        Value to interpolate at
+    lut : device array
+        1D lookup table
+    min_val : float
+        Minimum value in table
+    delta : float
+        Spacing between table entries
+    lut_size : int
+        Number of entries in table
+
+    Returns
+    -------
+    float
+        Interpolated value
+    """
+    idx_float = (value - min_val) / delta
+
+    if idx_float <= 0.0:
+        return lut[0]
+    if idx_float >= lut_size - 1:
+        return lut[lut_size - 1]
+
+    idx = int(idx_float)
+    frac = idx_float - idx
+
+    return lut[idx] + frac * (lut[idx + 1] - lut[idx])
+
+
+@cuda.jit(device=True)
+def _device_simple_n_and_gradient(
+    x: float,
+    y: float,
+    z: float,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    ref_radius: float,
+    min_alt: float,
+    delta_h: float,
+    lut_size: int,
+    lut_n: cuda.devicearray.DeviceNDArray,
+    lut_dn_dh: cuda.devicearray.DeviceNDArray,
+) -> tuple[float, float, float, float]:
+    """Compute n and gradient for SimpleInhomogeneousModel."""
+    dx = x - center_x
+    dy = y - center_y
+    dz = z - center_z
+
+    r = math.sqrt(dx * dx + dy * dy + dz * dz)
+    altitude = r - ref_radius
+
+    if altitude < 0.0:
+        altitude = 0.0
+
+    n = _device_lut_interpolate(altitude, lut_n, min_alt, delta_h, lut_size)
+    dn_dh = _device_lut_interpolate(altitude, lut_dn_dh, min_alt, delta_h, lut_size)
+
+    if r < 1e-10:
+        r_hat_x, r_hat_y, r_hat_z = 0.0, 0.0, 1.0
+    else:
+        r_hat_x = dx / r
+        r_hat_y = dy / r
+        r_hat_z = dz / r
+
+    grad_x = dn_dh * r_hat_x
+    grad_y = dn_dh * r_hat_y
+    grad_z = dn_dh * r_hat_z
+
+    return n, grad_x, grad_y, grad_z
+
+
+@cuda.jit(device=True)
+def _device_simple_euler_step(
+    x: float,
+    y: float,
+    z: float,
+    dir_x: float,
+    dir_y: float,
+    dir_z: float,
+    step_size: float,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    ref_radius: float,
+    min_alt: float,
+    delta_h: float,
+    lut_size: int,
+    lut_n: cuda.devicearray.DeviceNDArray,
+    lut_dn_dh: cuda.devicearray.DeviceNDArray,
+) -> tuple[float, float, float, float, float, float, float]:
+    """Euler step for SimpleInhomogeneousModel."""
+    n, grad_x, grad_y, grad_z = _device_simple_n_and_gradient(
+        x,
+        y,
+        z,
+        center_x,
+        center_y,
+        center_z,
+        ref_radius,
+        min_alt,
+        delta_h,
+        lut_size,
+        lut_n,
+        lut_dn_dh,
+    )
+
+    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_simple_rk4_step(
+    x: float,
+    y: float,
+    z: float,
+    dir_x: float,
+    dir_y: float,
+    dir_z: float,
+    step_size: float,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    ref_radius: float,
+    min_alt: float,
+    delta_h: float,
+    lut_size: int,
+    lut_n: cuda.devicearray.DeviceNDArray,
+    lut_dn_dh: cuda.devicearray.DeviceNDArray,
+) -> tuple[float, float, float, float, float, float, float]:
+    """RK4 step for SimpleInhomogeneousModel."""
+    h = step_size
+    h2 = h / 2.0
+
+    def get_n_and_kappa(px, py, pz, dx, dy, dz):
+        n, gx, gy, gz = _device_simple_n_and_gradient(
+            px,
+            py,
+            pz,
+            center_x,
+            center_y,
+            center_z,
+            ref_radius,
+            min_alt,
+            delta_h,
+            lut_size,
+            lut_n,
+            lut_dn_dh,
+        )
+        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.SIMPLE_EULER)
+@cuda.jit
+def _kernel_simple_inhomogeneous_euler(
+    positions,
+    directions,
+    active,
+    geo_path,
+    opt_path,
+    acc_time,
+    step_size: float,
+    num_steps: int,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    ref_radius: float,
+    min_alt: float,
+    delta_h: float,
+    lut_size: int,
+    lut_n,
+    lut_dn_dh,
+    wavelength: float,
+):
+    """GPU kernel for SimpleInhomogeneousModel 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_simple_euler_step(
+            x,
+            y,
+            z,
+            dx,
+            dy,
+            dz,
+            step_size,
+            center_x,
+            center_y,
+            center_z,
+            ref_radius,
+            min_alt,
+            delta_h,
+            lut_size,
+            lut_n,
+            lut_dn_dh,
+        )
+        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.SIMPLE_RK4)
+@cuda.jit
+def _kernel_simple_inhomogeneous_rk4(
+    positions,
+    directions,
+    active,
+    geo_path,
+    opt_path,
+    acc_time,
+    step_size: float,
+    num_steps: int,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    ref_radius: float,
+    min_alt: float,
+    delta_h: float,
+    lut_size: int,
+    lut_n,
+    lut_dn_dh,
+    wavelength: float,
+):
+    """GPU kernel for SimpleInhomogeneousModel 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_simple_rk4_step(
+            x,
+            y,
+            z,
+            dx,
+            dy,
+            dz,
+            step_size,
+            center_x,
+            center_y,
+            center_z,
+            ref_radius,
+            min_alt,
+            delta_h,
+            lut_size,
+            lut_n,
+            lut_dn_dh,
+        )
+        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