lsurf 1.0.0__py3-none-any.whl

lsurf/surfaces/gpu/curved_wave.py

@@ -0,0 +1,483 @@
+# 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.
+
+"""
+Curved Wave Surface (GPU-Capable)
+
+GPU-accelerated ocean wave surface on a curved (spherical) Earth.
+Supports a single wave component for efficient signed distance computation on GPU.
+
+For multi-wave surfaces or CPU-only computation, see the CPU-only
+CurvedWaveSurface in surfaces.cpu.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from ..protocol import Surface, SurfaceRole
+from ..registry import register_surface_type
+
+# Earth parameters
+EARTH_RADIUS = 6.371e6  # Earth's mean radius in meters
+GRAVITY = 9.81  # Gravity for deep water dispersion
+
+
+@dataclass
+class GPUCurvedWaveSurface(Surface):
+    """
+    Curved-earth ocean wave surface with GPU acceleration.
+
+    This is a simplified single-wave surface on a spherical Earth,
+    optimized for GPU computation. The wave is treated as a perturbation
+    on top of Earth's spherical surface.
+
+    For multiple superimposed waves, use the CPU-only CurvedWaveSurface.
+
+    Parameters
+    ----------
+    amplitude : float
+        Wave amplitude in meters.
+    wavelength : float
+        Wave wavelength (crest-to-crest) in meters.
+    direction : tuple of float
+        Wave propagation direction (dx, dy) in local tangent coordinates.
+    earth_center : tuple of float, optional
+        Center of Earth sphere. Default is (0, 0, -EARTH_RADIUS).
+    earth_radius : float, optional
+        Earth radius in meters. Default is EARTH_RADIUS.
+    time : float, optional
+        Animation time in seconds. Default is 0.0.
+    role : SurfaceRole
+        What happens when a ray hits (typically OPTICAL).
+    name : str
+        Human-readable name.
+    material_front : MaterialField, optional
+        Material above surface (atmosphere).
+    material_back : MaterialField, optional
+        Material below surface (ocean water).
+
+    Examples
+    --------
+    >>> from lsurf.surfaces import GPUCurvedWaveSurface, SurfaceRole
+    >>> from lsurf.materials import ExponentialAtmosphere, WATER
+    >>>
+    >>> ocean = GPUCurvedWaveSurface(
+    ...     amplitude=1.5,
+    ...     wavelength=50.0,
+    ...     direction=(1.0, 0.0),
+    ...     role=SurfaceRole.OPTICAL,
+    ...     name="ocean",
+    ...     material_front=ExponentialAtmosphere(),
+    ...     material_back=WATER,
+    ... )
+    """
+
+    amplitude: float
+    wavelength: float
+    direction: tuple[float, float]
+    role: SurfaceRole
+    earth_center: tuple[float, float, float] = (0, 0, -EARTH_RADIUS)
+    earth_radius: float = EARTH_RADIUS
+    time: float = 0.0
+    name: str = "gpu_curved_wave"
+    material_front: Any = None
+    material_back: Any = None
+
+    # GPU capability
+    _gpu_capable: bool = field(default=True, init=False, repr=False)
+    _geometry_id: int = field(default=4, init=False, repr=False)  # curved_wave = 4
+
+    # Precomputed values (set in __post_init__)
+    _wave_number: float = field(default=0.0, init=False, repr=False)
+    _dir_normalized: tuple[float, float] = field(
+        default=(1.0, 0.0), init=False, repr=False
+    )
+    _earth_center_arr: npt.NDArray = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if self.amplitude <= 0:
+            raise ValueError("Amplitude must be positive")
+        if self.wavelength <= 0:
+            raise ValueError("Wavelength must be positive")
+        if self.earth_radius <= 0:
+            raise ValueError("Earth radius must be positive")
+
+        # Compute wave number
+        self._wave_number = 2.0 * np.pi / self.wavelength
+
+        # Normalize direction
+        dx, dy = self.direction
+        norm = np.sqrt(dx * dx + dy * dy)
+        if norm < 1e-10:
+            self._dir_normalized = (1.0, 0.0)
+        else:
+            self._dir_normalized = (dx / norm, dy / norm)
+
+        # Store earth center as array
+        self._earth_center_arr = np.array(self.earth_center, dtype=np.float64)
+
+    @property
+    def gpu_capable(self) -> bool:
+        """This surface supports GPU acceleration."""
+        return True
+
+    @property
+    def geometry_id(self) -> int:
+        """GPU geometry type ID (curved_wave = 4)."""
+        return 4
+
+    @property
+    def wave_number(self) -> float:
+        """Wave number k = 2*pi/wavelength."""
+        return self._wave_number
+
+    @property
+    def angular_frequency(self) -> float:
+        """Angular frequency from deep water dispersion: omega = sqrt(g*k)."""
+        return np.sqrt(GRAVITY * self._wave_number)
+
+    def get_gpu_parameters(self) -> tuple:
+        """
+        Return parameters for GPU kernel.
+
+        Parameter layout (geometry_id = 4):
+        - p0: earth_center_x
+        - p1: earth_center_y
+        - p2: earth_center_z
+        - p3: earth_radius
+        - p4: amplitude
+        - p5: wave_number
+        - p6: dir_x
+        - p7: dir_y
+        - p8: time
+        - p9-p11: unused (0.0)
+
+        Returns
+        -------
+        tuple of 12 floats
+        """
+        return (
+            self.earth_center[0],
+            self.earth_center[1],
+            self.earth_center[2],
+            self.earth_radius,
+            self.amplitude,
+            self._wave_number,
+            self._dir_normalized[0],
+            self._dir_normalized[1],
+            self.time,
+            0.0,  # unused
+            0.0,  # unused
+            0.0,  # unused
+        )
+
+    def get_materials(self) -> tuple | None:
+        """Return materials for Fresnel calculation."""
+        if self.role == SurfaceRole.OPTICAL:
+            return (self.material_front, self.material_back)
+        return None
+
+    def _wave_height(self, x: npt.NDArray, y: npt.NDArray) -> npt.NDArray:
+        """Compute wave height at positions (x, y) in local tangent space."""
+        dir_x, dir_y = self._dir_normalized
+        omega = self.angular_frequency
+
+        # Phase at each position (simplified: use x,y as tangent coordinates)
+        dot = dir_x * x + dir_y * y
+        theta = self._wave_number * dot - omega * self.time
+
+        return self.amplitude * np.cos(theta)
+
+    def signed_distance(
+        self,
+        positions: npt.NDArray[np.float32],
+    ) -> npt.NDArray[np.float32]:
+        """
+        Compute signed distance from positions to curved wave surface.
+
+        Parameters
+        ----------
+        positions : ndarray, shape (N, 3)
+            Points to compute distance for
+
+        Returns
+        -------
+        ndarray, shape (N,)
+            Signed distance (positive outside, negative inside Earth+wave)
+        """
+        positions = positions.astype(np.float64)
+
+        # Vector from Earth center to each point
+        to_pos = positions - self._earth_center_arr
+        dist_from_center = np.linalg.norm(to_pos, axis=1)
+
+        # Wave height using x,y as local tangent coordinates
+        x = positions[:, 0]
+        y = positions[:, 1]
+        wave_height = self._wave_height(x, y)
+
+        # Surface is at earth_radius + wave_height from center
+        surface_radius = self.earth_radius + wave_height
+
+        # Signed distance: positive outside, negative inside
+        return (dist_from_center - surface_radius).astype(np.float32)
+
+    def intersect(
+        self,
+        origins: npt.NDArray[np.float32],
+        directions: npt.NDArray[np.float32],
+        min_distance: float = 1e-6,
+        max_iterations: int = 200,
+        tolerance: float = 1e-3,
+        max_distance: float | None = None,
+    ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]:
+        """
+        Find ray-surface intersections using ray marching.
+
+        Parameters
+        ----------
+        origins : ndarray, shape (N, 3)
+            Ray origin positions.
+        directions : ndarray, shape (N, 3)
+            Ray direction unit vectors.
+        min_distance : float
+            Minimum valid intersection distance.
+        max_iterations : int
+            Maximum ray marching iterations.
+        tolerance : float
+            Convergence tolerance in meters.
+        max_distance : float, optional
+            Maximum search distance.
+
+        Returns
+        -------
+        distances : ndarray, shape (N,)
+            Distance to intersection (inf if no hit).
+        hit_mask : ndarray, shape (N,), dtype=bool
+            True for rays that hit the surface.
+        """
+        origins = origins.astype(np.float64)
+        directions = directions.astype(np.float64)
+        n_rays = len(origins)
+
+        distances = np.full(n_rays, np.inf, dtype=np.float64)
+        hit_mask = np.zeros(n_rays, dtype=bool)
+
+        t = np.full(n_rays, min_distance, dtype=np.float64)
+        active = np.ones(n_rays, dtype=bool)
+
+        # Find intersection with outer sphere first
+        outer_radius = self.earth_radius + self.amplitude
+        oc = origins - self._earth_center_arr
+        a = np.sum(directions * directions, axis=1)
+        b = 2.0 * np.sum(directions * oc, axis=1)
+        c_outer = np.sum(oc * oc, axis=1) - outer_radius**2
+
+        discriminant_outer = b**2 - 4 * a * c_outer
+        has_potential_hit = discriminant_outer >= 0
+        active[~has_potential_hit] = False
+
+        sqrt_disc_outer = np.sqrt(np.maximum(discriminant_outer, 0))
+        t1_outer = (-b - sqrt_disc_outer) / (2 * a + 1e-20)
+        t_start = np.where(t1_outer > min_distance, t1_outer, min_distance)
+        t = t_start.copy()
+
+        prev_signed_dist = np.full(n_rays, np.inf)
+        prev_t = t.copy()
+        relaxation = 0.5
+
+        for _ in range(max_iterations):
+            if not np.any(active):
+                break
+
+            positions = origins + t[:, np.newaxis] * directions
+
+            to_pos = positions - self._earth_center_arr
+            dist_from_center = np.linalg.norm(to_pos, axis=1)
+            radial = to_pos / np.maximum(dist_from_center[:, np.newaxis], 1e-10)
+
+            cos_angle = np.abs(np.sum(directions * radial, axis=1))
+            cos_angle = np.maximum(cos_angle, 0.01)
+
+            x = positions[:, 0]
+            y = positions[:, 1]
+            wave_height = self._wave_height(x, y)
+
+            surface_radius = self.earth_radius + wave_height
+            signed_dist = dist_from_center - surface_radius
+
+            converged = np.abs(signed_dist) < tolerance
+            hit_mask[active & converged] = True
+            distances[active & converged] = t[active & converged]
+            active[converged] = False
+
+            # Bisection for sign changes
+            crossed = (
+                active
+                & (signed_dist * prev_signed_dist < 0)
+                & np.isfinite(prev_signed_dist)
+            )
+            if np.any(crossed):
+                t_low = np.where(prev_signed_dist > 0, prev_t, t)
+                t_high = np.where(prev_signed_dist > 0, t, prev_t)
+
+                for _ in range(15):
+                    t_mid = (t_low + t_high) / 2
+                    pos_mid = origins + t_mid[:, np.newaxis] * directions
+                    to_pos_mid = pos_mid - self._earth_center_arr
+                    dist_mid = np.linalg.norm(to_pos_mid, axis=1)
+                    wh_mid = self._wave_height(pos_mid[:, 0], pos_mid[:, 1])
+                    sd_mid = dist_mid - (self.earth_radius + wh_mid)
+
+                    above = sd_mid > 0
+                    t_low = np.where(crossed & above, t_mid, t_low)
+                    t_high = np.where(crossed & ~above, t_mid, t_high)
+
+                t[crossed] = (t_low[crossed] + t_high[crossed]) / 2
+
+                positions = origins + t[:, np.newaxis] * directions
+                to_pos = positions - self._earth_center_arr
+                dist_from_center = np.linalg.norm(to_pos, axis=1)
+                wave_height = self._wave_height(positions[:, 0], positions[:, 1])
+                signed_dist = dist_from_center - (self.earth_radius + wave_height)
+
+                converged = np.abs(signed_dist) < tolerance
+                hit_mask[active & converged] = True
+                distances[active & converged] = t[active & converged]
+                active[converged] = False
+
+            too_far = signed_dist < -self.amplitude - 10
+            active[too_far] = False
+
+            if max_distance is not None:
+                exceeded_max = t > max_distance
+                active[exceeded_max] = False
+
+            prev_signed_dist = signed_dist.copy()
+            prev_t = t.copy()
+
+            step = signed_dist / cos_angle * relaxation
+            step = np.clip(step, -self.amplitude * 2, self.amplitude * 2)
+            t[active] += step[active]
+            t = np.maximum(t, 0)
+
+        return distances.astype(np.float32), hit_mask
+
+    def normal_at(
+        self,
+        positions: npt.NDArray[np.float32],
+        incoming_directions: npt.NDArray[np.float32] | None = None,
+    ) -> npt.NDArray[np.float32]:
+        """
+        Compute surface normal at given positions.
+
+        Parameters
+        ----------
+        positions : ndarray, shape (N, 3)
+            Points on the surface.
+        incoming_directions : ndarray, shape (N, 3), optional
+            Incoming ray directions.
+
+        Returns
+        -------
+        normals : ndarray, shape (N, 3)
+            Unit normal vectors.
+        """
+        positions = positions.astype(np.float64)
+
+        # Get local tangent basis
+        to_pos = positions - self._earth_center_arr
+        dist = np.linalg.norm(to_pos, axis=1, keepdims=True)
+        dist = np.maximum(dist, 1e-10)
+        radial = to_pos / dist
+
+        # Compute tangent vectors (simplified)
+        global_x = np.array([1.0, 0.0, 0.0], dtype=np.float64)
+        dot_x = np.sum(radial * global_x, axis=1, keepdims=True)
+        tangent_x = global_x - dot_x * radial
+        tangent_x_norm = np.linalg.norm(tangent_x, axis=1, keepdims=True)
+        tangent_x = tangent_x / np.maximum(tangent_x_norm, 1e-10)
+
+        global_y = np.array([0.0, 1.0, 0.0], dtype=np.float64)
+        dot_y = np.sum(radial * global_y, axis=1, keepdims=True)
+        tangent_y = global_y - dot_y * radial
+        tangent_y_norm = np.linalg.norm(tangent_y, axis=1, keepdims=True)
+        tangent_y = tangent_y / np.maximum(tangent_y_norm, 1e-10)
+
+        # Compute wave normal in local coordinates
+        x = positions[:, 0]
+        y = positions[:, 1]
+        dir_x, dir_y = self._dir_normalized
+        omega = self.angular_frequency
+
+        dot_val = dir_x * x + dir_y * y
+        theta = self._wave_number * dot_val - omega * self.time
+
+        sin_theta = np.sin(theta)
+        WA = self._wave_number * self.amplitude
+
+        # Local normal components
+        nx_local = dir_x * WA * sin_theta
+        ny_local = dir_y * WA * sin_theta
+        nz_local = 1.0 - 0.5 * WA * sin_theta  # Simplified steepness
+
+        norm = np.sqrt(nx_local**2 + ny_local**2 + nz_local**2)
+        norm = np.maximum(norm, 1e-10)
+        nx_local /= norm
+        ny_local /= norm
+        nz_local /= norm
+
+        # Transform to world coordinates
+        normals = (
+            nx_local[:, np.newaxis] * tangent_x
+            + ny_local[:, np.newaxis] * tangent_y
+            + nz_local[:, np.newaxis] * radial
+        )
+
+        if incoming_directions is not None:
+            dot_products = np.sum(normals * incoming_directions, axis=1)
+            flip_mask = dot_products > 0
+            normals[flip_mask] *= -1
+
+        return normals.astype(np.float32)
+
+    def set_time(self, time: float) -> None:
+        """Update the wave animation time."""
+        self.time = time
+
+
+# Register class with registry
+register_surface_type("gpu_curved_wave", "gpu", 4, GPUCurvedWaveSurface)