lsurf 1.0.0__py3-none-any.whl

lsurf/surfaces/gpu/gerstner_wave.py

@@ -0,0 +1,377 @@
+# 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.
+
+"""
+Gerstner Wave Surface (GPU-Capable)
+
+GPU-accelerated flat-earth Gerstner wave surface. Supports a single wave
+component for efficient signed distance computation on GPU.
+
+For multi-wave surfaces or CPU-only computation, see the CPU-only
+GerstnerWaveSurface 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
+
+# Gravity constant for deep water dispersion
+GRAVITY = 9.81
+
+
+@dataclass
+class GPUGerstnerWaveSurface(Surface):
+    """
+    Flat-earth Gerstner wave surface with GPU acceleration.
+
+    This is a simplified single-wave surface optimized for GPU computation.
+    For multiple superimposed waves, use the CPU-only GerstnerWaveSurface.
+
+    The surface implements a height field: z = reference_z + A*cos(k*dot(d,xy) - omega*t + phi)
+
+    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), will be normalized.
+    reference_z : float
+        Mean sea level z-coordinate in meters.
+    phase : float, optional
+        Initial phase offset in radians. Default is 0.0.
+    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 (air/atmosphere).
+    material_back : MaterialField, optional
+        Material below surface (water).
+
+    Examples
+    --------
+    >>> from lsurf.surfaces import GPUGerstnerWaveSurface, SurfaceRole
+    >>> from lsurf.materials import AIR_STP, WATER
+    >>>
+    >>> ocean = GPUGerstnerWaveSurface(
+    ...     amplitude=1.5,
+    ...     wavelength=50.0,
+    ...     direction=(1.0, 0.0),
+    ...     reference_z=0.0,
+    ...     role=SurfaceRole.OPTICAL,
+    ...     name="ocean",
+    ...     material_front=AIR_STP,
+    ...     material_back=WATER,
+    ... )
+    """
+
+    amplitude: float
+    wavelength: float
+    direction: tuple[float, float]
+    reference_z: float
+    role: SurfaceRole
+    phase: float = 0.0
+    time: float = 0.0
+    name: str = "gpu_gerstner_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=3, init=False, repr=False)  # gerstner_wave = 3
+
+    # 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
+    )
+
+    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")
+
+        # 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)
+
+    @property
+    def gpu_capable(self) -> bool:
+        """This surface supports GPU acceleration."""
+        return True
+
+    @property
+    def geometry_id(self) -> int:
+        """GPU geometry type ID (gerstner_wave = 3)."""
+        return 3
+
+    @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 = 3):
+        - p0: amplitude
+        - p1: wave_number
+        - p2: dir_x
+        - p3: dir_y
+        - p4: reference_z
+        - p5: phase
+        - p6: time
+        - p7-p11: unused (0.0)
+
+        Returns
+        -------
+        tuple of 12 floats
+        """
+        return (
+            self.amplitude,
+            self._wave_number,
+            self._dir_normalized[0],
+            self._dir_normalized[1],
+            self.reference_z,
+            self.phase,
+            self.time,
+            0.0,  # unused
+            0.0,  # unused
+            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)."""
+        dir_x, dir_y = self._dir_normalized
+        omega = self.angular_frequency
+
+        # Phase at each position
+        dot = dir_x * x + dir_y * y
+        theta = self._wave_number * dot - omega * self.time + self.phase
+
+        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 wave surface.
+
+        Parameters
+        ----------
+        positions : ndarray, shape (N, 3)
+            Points to compute distance for
+
+        Returns
+        -------
+        ndarray, shape (N,)
+            Signed distance (positive above surface, negative below)
+        """
+        x = positions[:, 0].astype(np.float64)
+        y = positions[:, 1].astype(np.float64)
+        z = positions[:, 2].astype(np.float64)
+
+        surface_z = self.reference_z + self._wave_height(x, y)
+        return (z - surface_z).astype(np.float32)
+
+    def intersect(
+        self,
+        origins: npt.NDArray[np.float32],
+        directions: npt.NDArray[np.float32],
+        min_distance: float = 1e-6,
+    ) -> 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.
+
+        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)
+
+        tolerance = 1e-4
+        max_iterations = 100
+        max_distance = 10000.0
+
+        for i in range(n_rays):
+            ox, oy, oz = origins[i]
+            dx, dy, dz = directions[i]
+
+            # Initial guess
+            if abs(dz) > 1e-10:
+                t = (self.reference_z - oz) / dz
+            else:
+                t = 0.0
+
+            t = max(t - self.amplitude / max(abs(dz), 0.1), min_distance)
+
+            for _ in range(max_iterations):
+                px = ox + t * dx
+                py = oy + t * dy
+                pz = oz + t * dz
+
+                z_surf = (
+                    self.reference_z
+                    + self._wave_height(np.array([px]), np.array([py]))[0]
+                )
+
+                signed_dist = pz - z_surf
+
+                if abs(signed_dist) < tolerance:
+                    if t > min_distance:
+                        distances[i] = t
+                        hit_mask[i] = True
+                    break
+
+                if abs(dz) > 0.01:
+                    step = signed_dist / abs(dz) * 0.8
+                else:
+                    step = signed_dist * 0.5
+
+                step = np.clip(step, -max_distance * 0.1, max_distance * 0.1)
+                t += step
+
+                if t < 0 or t > max_distance:
+                    break
+
+        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.
+        """
+        x = positions[:, 0].astype(np.float64)
+        y = positions[:, 1].astype(np.float64)
+        n_points = len(x)
+
+        dir_x, dir_y = self._dir_normalized
+        omega = self.angular_frequency
+
+        # Phase at each position
+        dot = dir_x * x + dir_y * y
+        theta = self._wave_number * dot - omega * self.time + self.phase
+
+        # Gradient of height field
+        # dz/dx = A * k * dir_x * (-sin(theta))
+        # dz/dy = A * k * dir_y * (-sin(theta))
+        sin_theta = np.sin(theta)
+        dz_dx = -self.amplitude * self._wave_number * dir_x * sin_theta
+        dz_dy = -self.amplitude * self._wave_number * dir_y * sin_theta
+
+        # Normal = (-dz/dx, -dz/dy, 1) normalized
+        nx = -dz_dx
+        ny = -dz_dy
+        nz = np.ones(n_points, dtype=np.float64)
+
+        norm = np.sqrt(nx * nx + ny * ny + nz * nz)
+        normals = np.stack([nx / norm, ny / norm, nz / norm], axis=-1)
+
+        if incoming_directions is not None:
+            dot_products = np.sum(incoming_directions * normals, 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_gerstner_wave", "gpu", 3, GPUGerstnerWaveSurface)