lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/intersection_handler.py

@@ -0,0 +1,327 @@
+# 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.
+
+"""
+Intersection Handler
+
+GPU memory management for ray-surface intersection kernels.
+This module provides high-level interfaces that handle cuda.to_device(),
+copy_to_host(), and kernel launching for the intersection kernels.
+"""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..kernels.intersection import (
+    kernel_plane_intersect,
+    kernel_sphere_intersect,
+)
+
+# 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]
+
+
+def intersect_plane_gpu(
+    origins: NDArray[np.float32],
+    directions: NDArray[np.float32],
+    plane_normal: NDArray[np.float32],
+    plane_point: NDArray[np.float32],
+    threads_per_block: int = 256,
+) -> tuple[NDArray[np.float32], NDArray[np.bool_], NDArray[np.float32]]:
+    """
+    GPU-accelerated ray-plane intersection.
+
+    Parameters
+    ----------
+    origins : ndarray, shape (N, 3)
+        Ray origins
+    directions : ndarray, shape (N, 3)
+        Ray directions (should be normalized)
+    plane_normal : ndarray, shape (3,)
+        Plane normal vector (should be normalized)
+    plane_point : ndarray, shape (3,)
+        A point on the plane
+    threads_per_block : int, optional
+        CUDA threads per block (default 256)
+
+    Returns
+    -------
+    distances : ndarray, shape (N,)
+        Distance to intersection for each ray (inf if no hit)
+    hit_mask : ndarray, shape (N,)
+        Boolean mask indicating which rays hit the plane
+    normals : ndarray, shape (N, 3)
+        Surface normal at intersection (pointing toward ray origin)
+    """
+    num_rays = len(origins)
+
+    # Compute plane equation: n.p + d = 0
+    # d = -n.point
+    plane_d = -np.dot(plane_normal, plane_point)
+
+    if _HAS_CUDA:
+        # Allocate output arrays
+        distances = np.full(num_rays, np.inf, dtype=np.float32)
+        hit_mask = np.zeros(num_rays, dtype=np.bool_)
+        normals = np.zeros((num_rays, 3), dtype=np.float32)
+
+        # Transfer to GPU
+        d_origins = cuda.to_device(origins.astype(np.float32))
+        d_directions = cuda.to_device(directions.astype(np.float32))
+        d_distances = cuda.to_device(distances)
+        d_hit_mask = cuda.to_device(hit_mask)
+        d_normals = cuda.to_device(normals)
+
+        # Launch kernel
+        blocks = (num_rays + threads_per_block - 1) // threads_per_block
+        kernel_plane_intersect[blocks, threads_per_block](
+            d_origins,
+            d_directions,
+            d_distances,
+            d_hit_mask,
+            d_normals,
+            np.float32(plane_normal[0]),
+            np.float32(plane_normal[1]),
+            np.float32(plane_normal[2]),
+            np.float32(plane_d),
+        )
+
+        cuda.synchronize()
+        distances = d_distances.copy_to_host()
+        hit_mask = d_hit_mask.copy_to_host()
+        normals = d_normals.copy_to_host()
+    else:
+        distances, hit_mask, normals = _intersect_plane_cpu(
+            origins, directions, plane_normal, plane_point
+        )
+
+    return distances, hit_mask, normals
+
+
+def intersect_sphere_gpu(
+    origins: NDArray[np.float32],
+    directions: NDArray[np.float32],
+    center: NDArray[np.float32],
+    radius: float,
+    threads_per_block: int = 256,
+) -> tuple[NDArray[np.float32], NDArray[np.bool_], NDArray[np.float32]]:
+    """
+    GPU-accelerated ray-sphere intersection.
+
+    Parameters
+    ----------
+    origins : ndarray, shape (N, 3)
+        Ray origins
+    directions : ndarray, shape (N, 3)
+        Ray directions (should be normalized)
+    center : ndarray, shape (3,)
+        Sphere center position
+    radius : float
+        Sphere radius
+    threads_per_block : int, optional
+        CUDA threads per block (default 256)
+
+    Returns
+    -------
+    distances : ndarray, shape (N,)
+        Distance to intersection for each ray (inf if no hit)
+    hit_mask : ndarray, shape (N,)
+        Boolean mask indicating which rays hit the sphere
+    normals : ndarray, shape (N, 3)
+        Surface normal at intersection (pointing outward)
+    """
+    num_rays = len(origins)
+
+    if _HAS_CUDA:
+        # Allocate output arrays
+        distances = np.full(num_rays, np.inf, dtype=np.float32)
+        hit_mask = np.zeros(num_rays, dtype=np.bool_)
+        normals = np.zeros((num_rays, 3), dtype=np.float32)
+
+        # Transfer to GPU
+        d_origins = cuda.to_device(origins.astype(np.float32))
+        d_directions = cuda.to_device(directions.astype(np.float32))
+        d_distances = cuda.to_device(distances)
+        d_hit_mask = cuda.to_device(hit_mask)
+        d_normals = cuda.to_device(normals)
+
+        # Launch kernel
+        blocks = (num_rays + threads_per_block - 1) // threads_per_block
+        kernel_sphere_intersect[blocks, threads_per_block](
+            d_origins,
+            d_directions,
+            d_distances,
+            d_hit_mask,
+            d_normals,
+            np.float32(center[0]),
+            np.float32(center[1]),
+            np.float32(center[2]),
+            np.float32(radius),
+        )
+
+        cuda.synchronize()
+        distances = d_distances.copy_to_host()
+        hit_mask = d_hit_mask.copy_to_host()
+        normals = d_normals.copy_to_host()
+    else:
+        distances, hit_mask, normals = _intersect_sphere_cpu(
+            origins, directions, center, radius
+        )
+
+    return distances, hit_mask, normals
+
+
+# =============================================================================
+# CPU Fallback Implementations
+# =============================================================================
+
+
+def _intersect_plane_cpu(
+    origins: NDArray[np.float32],
+    directions: NDArray[np.float32],
+    plane_normal: NDArray[np.float32],
+    plane_point: NDArray[np.float32],
+) -> tuple[NDArray[np.float32], NDArray[np.bool_], NDArray[np.float32]]:
+    """CPU implementation of ray-plane intersection."""
+    num_rays = len(origins)
+
+    # Normalize directions
+    d_lens = np.linalg.norm(directions, axis=1, keepdims=True)
+    d_lens = np.where(d_lens < 1e-10, 1.0, d_lens)
+    directions = directions / d_lens
+
+    # Compute plane equation: d = -n.point
+    plane_d = -np.dot(plane_normal, plane_point)
+
+    # Compute denominator: n . direction
+    denom = np.dot(directions, plane_normal)
+
+    # Compute numerator: -(n . origin + d)
+    numer = -(np.dot(origins, plane_normal) + plane_d)
+
+    # Initialize outputs
+    distances = np.full(num_rays, np.inf, dtype=np.float32)
+    hit_mask = np.zeros(num_rays, dtype=np.bool_)
+    normals = np.zeros((num_rays, 3), dtype=np.float32)
+
+    # Find valid intersections
+    valid = np.abs(denom) > 1e-10
+    t = np.where(valid, numer / np.where(valid, denom, 1.0), np.inf)
+    hit = valid & (t > 0)
+
+    distances[hit] = t[hit]
+    hit_mask = hit
+
+    # Set normals (flip if ray hits from behind)
+    for i in range(num_rays):
+        if hit[i]:
+            if denom[i] > 0:
+                normals[i] = -plane_normal
+            else:
+                normals[i] = plane_normal
+
+    return distances.astype(np.float32), hit_mask, normals.astype(np.float32)
+
+
+def _intersect_sphere_cpu(
+    origins: NDArray[np.float32],
+    directions: NDArray[np.float32],
+    center: NDArray[np.float32],
+    radius: float,
+) -> tuple[NDArray[np.float32], NDArray[np.bool_], NDArray[np.float32]]:
+    """CPU implementation of ray-sphere intersection."""
+    num_rays = len(origins)
+
+    # Normalize directions
+    d_lens = np.linalg.norm(directions, axis=1, keepdims=True)
+    d_lens = np.where(d_lens < 1e-10, 1.0, d_lens)
+    directions = directions / d_lens
+
+    # Vector from origin to sphere center
+    oc = origins - center
+
+    # Quadratic coefficients
+    # a = 1 (normalized direction)
+    b = 2.0 * np.sum(directions * oc, axis=1)
+    c = np.sum(oc * oc, axis=1) - radius * radius
+
+    # Discriminant
+    discriminant = b * b - 4.0 * c
+
+    # Initialize outputs
+    distances = np.full(num_rays, np.inf, dtype=np.float32)
+    hit_mask = np.zeros(num_rays, dtype=np.bool_)
+    normals = np.zeros((num_rays, 3), dtype=np.float32)
+
+    # Find rays with valid intersections
+    valid = discriminant >= 0
+    sqrt_disc = np.sqrt(np.maximum(discriminant, 0))
+
+    t1 = np.where(valid, (-b - sqrt_disc) / 2.0, np.inf)
+    t2 = np.where(valid, (-b + sqrt_disc) / 2.0, np.inf)
+
+    # Use nearest positive t
+    t = np.where(t1 > 1e-6, t1, np.where(t2 > 1e-6, t2, np.inf))
+    hit = t < np.inf
+
+    distances[hit] = t[hit]
+    hit_mask = hit
+
+    # Compute normals at intersection points
+    for i in range(num_rays):
+        if hit[i]:
+            p = origins[i] + t[i] * directions[i]
+            n = p - center
+            n_len = np.linalg.norm(n)
+            if n_len > 1e-10:
+                normals[i] = n / n_len
+
+    return distances.astype(np.float32), hit_mask, normals.astype(np.float32)
+
+
+__all__ = [
+    "intersect_plane_gpu",
+    "intersect_sphere_gpu",
+]