lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/signed_distance_handler.py

@@ -0,0 +1,522 @@
+# 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.
+
+"""
+Signed Distance Handler
+
+GPU memory management for signed distance and surface normal kernels.
+This module provides high-level interfaces that handle cuda.to_device(),
+copy_to_host(), and kernel launching for the generic signed distance kernels.
+
+These functions enable GPU-accelerated intersection detection for ALL surface
+types by computing signed distance values that can be used for bisection.
+"""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..kernels.intersection import (
+    kernel_signed_distance,
+    kernel_surface_normal,
+)
+
+# GPU support is optional
+try:
+    from numba import cuda
+
+    _HAS_CUDA = cuda.is_available()
+except ImportError:
+    _HAS_CUDA = False
+
+    class _FakeCuda:
+        @staticmethod
+        def is_available():
+            return False
+
+    cuda = _FakeCuda()  # type: ignore[assignment]
+
+
+# Maximum number of parameters supported per surface
+# 64 params allows up to 8 wave components for multi-wave surfaces
+# Layout for multi-wave: 8 header params + 8 waves * 8 params each = 72, using 64
+MAX_SURFACE_PARAMS = 64
+
+
+def compute_signed_distance_gpu(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: tuple | NDArray[np.float32],
+    threads_per_block: int = 256,
+) -> NDArray[np.float32]:
+    """
+    GPU-accelerated signed distance computation for any surface type.
+
+    Parameters
+    ----------
+    positions : ndarray, shape (N, 3)
+        Points to compute signed distance for
+    geometry_id : int
+        Surface geometry type:
+        1 = Plane
+        2 = Sphere
+        3 = Gerstner wave (flat earth)
+        4 = Curved wave (spherical earth)
+        5 = Multi-wave curved (up to 8 waves)
+        7 = Annular plane (ring-shaped detector)
+    params : tuple or ndarray
+        Surface parameters from surface.get_gpu_parameters()
+        Will be padded to 12 elements if shorter.
+    threads_per_block : int, optional
+        CUDA threads per block (default 256)
+
+    Returns
+    -------
+    signed_distances : ndarray, shape (N,)
+        Signed distance from each position to the surface.
+        Sign convention depends on surface type (typically positive outside).
+    """
+    # Convert params to array and pad to MAX_SURFACE_PARAMS
+    params_arr = _prepare_params(params)
+
+    if _HAS_CUDA:
+        return _compute_signed_distance_cuda(
+            positions, geometry_id, params_arr, threads_per_block
+        )
+    else:
+        return _compute_signed_distance_cpu(positions, geometry_id, params_arr)
+
+
+def compute_surface_normal_gpu(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: tuple | NDArray[np.float32],
+    incoming_directions: NDArray[np.float32] | None = None,
+    threads_per_block: int = 256,
+) -> NDArray[np.float32]:
+    """
+    GPU-accelerated surface normal computation for any surface type.
+
+    Parameters
+    ----------
+    positions : ndarray, shape (N, 3)
+        Points on (or near) the surface
+    geometry_id : int
+        Surface geometry type (1=plane, 2=sphere, 3=gerstner, 4=curved_wave)
+    params : tuple or ndarray
+        Surface parameters from surface.get_gpu_parameters()
+    incoming_directions : ndarray, shape (N, 3), optional
+        If provided, normals will be flipped to face incoming rays
+    threads_per_block : int, optional
+        CUDA threads per block (default 256)
+
+    Returns
+    -------
+    normals : ndarray, shape (N, 3)
+        Unit normal vectors at each position
+    """
+    params_arr = _prepare_params(params)
+
+    if _HAS_CUDA:
+        normals = _compute_surface_normal_cuda(
+            positions, geometry_id, params_arr, threads_per_block
+        )
+    else:
+        normals = _compute_surface_normal_cpu(positions, geometry_id, params_arr)
+
+    # Optionally flip normals to face incoming rays
+    if incoming_directions is not None:
+        dot = np.sum(normals * incoming_directions, axis=1)
+        flip_mask = dot > 0
+        normals[flip_mask] = -normals[flip_mask]
+
+    return normals
+
+
+def _prepare_params(params: tuple | NDArray[np.float32]) -> NDArray[np.float32]:
+    """Convert params to float32 array, padded to MAX_SURFACE_PARAMS."""
+    if isinstance(params, tuple):
+        params_arr = np.array(params, dtype=np.float32)
+    else:
+        params_arr = params.astype(np.float32)
+
+    # Pad to MAX_SURFACE_PARAMS
+    if len(params_arr) < MAX_SURFACE_PARAMS:
+        padded = np.zeros(MAX_SURFACE_PARAMS, dtype=np.float32)
+        padded[: len(params_arr)] = params_arr
+        return padded
+    return params_arr[:MAX_SURFACE_PARAMS]
+
+
+# =============================================================================
+# CUDA Implementations
+# =============================================================================
+
+
+def _compute_signed_distance_cuda(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: NDArray[np.float32],
+    threads_per_block: int,
+) -> NDArray[np.float32]:
+    """CUDA implementation of signed distance computation."""
+    num_points = len(positions)
+
+    # Allocate output
+    signed_dist = np.zeros(num_points, dtype=np.float32)
+
+    # Transfer to GPU
+    d_positions = cuda.to_device(positions.astype(np.float32))
+    d_signed_dist = cuda.to_device(signed_dist)
+    d_params = cuda.to_device(params)
+
+    # Launch kernel
+    blocks = (num_points + threads_per_block - 1) // threads_per_block
+    kernel_signed_distance[blocks, threads_per_block](
+        d_positions,
+        d_signed_dist,
+        geometry_id,
+        d_params,
+    )
+
+    cuda.synchronize()
+    return d_signed_dist.copy_to_host()
+
+
+def _compute_surface_normal_cuda(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: NDArray[np.float32],
+    threads_per_block: int,
+) -> NDArray[np.float32]:
+    """CUDA implementation of surface normal computation."""
+    num_points = len(positions)
+
+    # Allocate output
+    normals = np.zeros((num_points, 3), dtype=np.float32)
+
+    # Transfer to GPU
+    d_positions = cuda.to_device(positions.astype(np.float32))
+    d_normals = cuda.to_device(normals)
+    d_params = cuda.to_device(params)
+
+    # Launch kernel
+    blocks = (num_points + threads_per_block - 1) // threads_per_block
+    kernel_surface_normal[blocks, threads_per_block](
+        d_positions,
+        d_normals,
+        geometry_id,
+        d_params,
+    )
+
+    cuda.synchronize()
+    return d_normals.copy_to_host()
+
+
+# =============================================================================
+# CPU Fallback Implementations
+# =============================================================================
+
+
+def _compute_signed_distance_cpu(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: NDArray[np.float32],
+) -> NDArray[np.float32]:
+    """CPU fallback implementation of signed distance computation."""
+    import math
+
+    num_points = len(positions)
+    signed_dist = np.zeros(num_points, dtype=np.float32)
+
+    x = positions[:, 0]
+    y = positions[:, 1]
+    z = positions[:, 2]
+
+    if geometry_id == 1:  # Plane
+        nx, ny, nz = params[0], params[1], params[2]
+        px, py, pz = params[3], params[4], params[5]
+        signed_dist = nx * (x - px) + ny * (y - py) + nz * (z - pz)
+
+    elif geometry_id == 2:  # Sphere
+        cx, cy, cz = params[0], params[1], params[2]
+        r = params[3]
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist = np.sqrt(dx * dx + dy * dy + dz * dz)
+        signed_dist = dist - abs(r)
+
+    elif geometry_id == 3:  # Gerstner wave
+        A = params[0]
+        k = params[1]
+        dir_x, dir_y = params[2], params[3]
+        ref_z = params[4]
+        phase = params[5]
+        time = params[6]
+
+        GRAVITY = 9.81
+        omega = math.sqrt(GRAVITY * k)
+        phase_val = k * (dir_x * x + dir_y * y) - omega * time + phase
+        wave_height = ref_z + A * np.cos(phase_val)
+        signed_dist = z - wave_height
+
+    elif geometry_id == 4:  # Curved wave
+        cx, cy, cz = params[0], params[1], params[2]
+        earth_radius = params[3]
+        A = params[4]
+        k = params[5]
+        dir_x, dir_y = params[6], params[7]
+        time = params[8]
+
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist_from_center = np.sqrt(dx * dx + dy * dy + dz * dz)
+
+        GRAVITY = 9.81
+        omega = math.sqrt(GRAVITY * k)
+        dot_val = dir_x * x + dir_y * y
+        phase_val = k * dot_val - omega * time
+        wave_height = A * np.cos(phase_val)
+        surface_radius = earth_radius + wave_height
+        signed_dist = dist_from_center - surface_radius
+
+    elif geometry_id == 5:  # Multi-wave curved
+        cx, cy, cz = params[0], params[1], params[2]
+        earth_radius = params[3]
+        time = params[4]
+        num_waves = int(params[5])
+
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist_from_center = np.sqrt(dx * dx + dy * dy + dz * dz)
+
+        GRAVITY = 9.81
+        total_wave_height = np.zeros(num_points, dtype=np.float64)
+        for i in range(num_waves):
+            offset = 8 + i * 8
+            A = params[offset + 0]
+            k = params[offset + 1]
+            dir_x = params[offset + 2]
+            dir_y = params[offset + 3]
+            phase = params[offset + 4]
+
+            omega = math.sqrt(GRAVITY * k)
+            dot_val = dir_x * x + dir_y * y
+            phase_val = k * dot_val - omega * time + phase
+            total_wave_height += A * np.cos(phase_val)
+
+        surface_radius = earth_radius + total_wave_height
+        signed_dist = dist_from_center - surface_radius
+
+    elif geometry_id == 7:  # Annular plane
+        # Signed distance to infinite plane (bounds checking done separately)
+        # params: [nx, ny, nz, cx, cy, cz, ...]
+        nx, ny, nz = params[0], params[1], params[2]
+        cx, cy, cz = params[3], params[4], params[5]
+        signed_dist = nx * (x - cx) + ny * (y - cy) + nz * (z - cz)
+
+    else:
+        # Unknown geometry - return inf
+        signed_dist = np.full(num_points, np.inf, dtype=np.float32)
+
+    return signed_dist.astype(np.float32)
+
+
+def _compute_surface_normal_cpu(
+    positions: NDArray[np.float32],
+    geometry_id: int,
+    params: NDArray[np.float32],
+) -> NDArray[np.float32]:
+    """CPU fallback implementation of surface normal computation."""
+    import math
+
+    num_points = len(positions)
+    normals = np.zeros((num_points, 3), dtype=np.float32)
+
+    x = positions[:, 0]
+    y = positions[:, 1]
+    z = positions[:, 2]
+
+    if geometry_id == 1:  # Plane
+        normals[:, 0] = params[0]
+        normals[:, 1] = params[1]
+        normals[:, 2] = params[2]
+
+    elif geometry_id == 2:  # Sphere
+        cx, cy, cz = params[0], params[1], params[2]
+        r = params[3]
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist = np.sqrt(dx * dx + dy * dy + dz * dz)
+        dist = np.maximum(dist, 1e-12)
+        normals[:, 0] = dx / dist
+        normals[:, 1] = dy / dist
+        normals[:, 2] = dz / dist
+        if r < 0:
+            normals *= -1
+
+    elif geometry_id == 3:  # Gerstner wave
+        A = params[0]
+        k = params[1]
+        dir_x, dir_y = params[2], params[3]
+        phase = params[5]
+        time = params[6]
+
+        GRAVITY = 9.81
+        omega = math.sqrt(GRAVITY * k)
+        phase_val = k * (dir_x * x + dir_y * y) - omega * time + phase
+        sin_theta = np.sin(phase_val)
+
+        dz_dx = -A * k * dir_x * sin_theta
+        dz_dy = -A * k * dir_y * sin_theta
+
+        nx = -dz_dx
+        ny = -dz_dy
+        nz = np.ones(num_points, dtype=np.float32)
+
+        norm = np.sqrt(nx * nx + ny * ny + nz * nz)
+        normals[:, 0] = nx / norm
+        normals[:, 1] = ny / norm
+        normals[:, 2] = nz / norm
+
+    elif geometry_id == 4:  # Curved wave
+        cx, cy, cz = params[0], params[1], params[2]
+        A = params[4]
+        k = params[5]
+        dir_x, dir_y = params[6], params[7]
+        time = params[8]
+
+        # Radial direction
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist = np.sqrt(dx * dx + dy * dy + dz * dz)
+        dist = np.maximum(dist, 1e-12)
+        radial = np.stack([dx / dist, dy / dist, dz / dist], axis=-1)
+
+        # Wave perturbation (simplified)
+        GRAVITY = 9.81
+        omega = math.sqrt(GRAVITY * k)
+        dot_val = dir_x * x + dir_y * y
+        phase_val = k * dot_val - omega * time
+        sin_theta = np.sin(phase_val)
+        WA = k * A
+
+        nx_local = dir_x * WA * sin_theta
+        ny_local = dir_y * WA * sin_theta
+        nz_local = 1.0 - 0.5 * WA * sin_theta
+
+        # Combine radial with perturbation
+        normals[:, 0] = radial[:, 0] + 0.01 * nx_local
+        normals[:, 1] = radial[:, 1] + 0.01 * ny_local
+        normals[:, 2] = radial[:, 2] + 0.01 * nz_local
+
+        # Normalize
+        norm = np.sqrt(normals[:, 0] ** 2 + normals[:, 1] ** 2 + normals[:, 2] ** 2)
+        norm = np.maximum(norm, 1e-12)
+        normals /= norm[:, np.newaxis]
+
+    elif geometry_id == 5:  # Multi-wave curved
+        cx, cy, cz = params[0], params[1], params[2]
+        time = params[4]
+        num_waves = int(params[5])
+
+        # Radial direction
+        dx = x - cx
+        dy = y - cy
+        dz = z - cz
+        dist = np.sqrt(dx * dx + dy * dy + dz * dz)
+        dist = np.maximum(dist, 1e-12)
+        radial = np.stack([dx / dist, dy / dist, dz / dist], axis=-1)
+
+        # Accumulate wave perturbations
+        GRAVITY = 9.81
+        nx_local = np.zeros(num_points, dtype=np.float64)
+        ny_local = np.zeros(num_points, dtype=np.float64)
+        nz_local = np.ones(num_points, dtype=np.float64)
+
+        for i in range(num_waves):
+            offset = 8 + i * 8
+            A = params[offset + 0]
+            k = params[offset + 1]
+            dir_x = params[offset + 2]
+            dir_y = params[offset + 3]
+            phase = params[offset + 4]
+
+            omega = math.sqrt(GRAVITY * k)
+            dot_val = dir_x * x + dir_y * y
+            phase_val = k * dot_val - omega * time + phase
+            sin_theta = np.sin(phase_val)
+            WA = k * A
+
+            nx_local += dir_x * WA * sin_theta
+            ny_local += dir_y * WA * sin_theta
+            nz_local -= 0.5 * WA * sin_theta
+
+        # Normalize local normal
+        local_norm = np.sqrt(nx_local**2 + ny_local**2 + nz_local**2)
+        local_norm = np.maximum(local_norm, 1e-12)
+        nx_local /= local_norm
+        ny_local /= local_norm
+        nz_local /= local_norm
+
+        # Combine radial with perturbation
+        normals[:, 0] = radial[:, 0] + 0.01 * nx_local
+        normals[:, 1] = radial[:, 1] + 0.01 * ny_local
+        normals[:, 2] = radial[:, 2] + 0.01 * nz_local
+
+        # Normalize
+        norm = np.sqrt(normals[:, 0] ** 2 + normals[:, 1] ** 2 + normals[:, 2] ** 2)
+        norm = np.maximum(norm, 1e-12)
+        normals /= norm[:, np.newaxis]
+
+    elif geometry_id == 7:  # Annular plane
+        # Normal is constant everywhere (stored in params)
+        normals[:, 0] = params[0]
+        normals[:, 1] = params[1]
+        normals[:, 2] = params[2]
+
+    else:
+        # Unknown geometry - default to +Z
+        normals[:, 2] = 1.0
+
+    return normals.astype(np.float32)
+
+
+__all__ = [
+    "compute_signed_distance_gpu",
+    "compute_surface_normal_gpu",
+]