Source code for lsurf.surfaces.gpu.curved_wave

# 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


[docs] @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)
[docs] 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 )
[docs] 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)
[docs] 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)
[docs] 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
[docs] 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)
[docs] 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)