lsurf 1.0.0__py3-none-any.whl

lsurf/surfaces/cpu/curved_wave.py

@@ -0,0 +1,463 @@
+# 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 (CPU-Only)
+
+Ocean wave surface on a curved (spherical) Earth. Uses ray marching
+for intersection due to complex geometry.
+
+For flat-earth wave surfaces, see GerstnerWaveSurface.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..protocol import Surface, SurfaceRole
+from ..registry import register_surface_type
+from .wave_params import GerstnerWaveParams
+
+# Earth parameters
+EARTH_RADIUS = 6.371e6  # Earth's mean radius in meters
+
+
+@dataclass
+class CurvedWaveSurface(Surface):
+    """
+    Ocean wave surface on a curved (spherical) Earth (CPU-only).
+
+    Combines Earth's spherical curvature with Gerstner wave perturbations
+    applied in local tangent space.
+
+    Parameters
+    ----------
+    wave_params : list of GerstnerWaveParams
+        List of wave components.
+    role : SurfaceRole
+        What happens when a ray hits (typically OPTICAL).
+    name : str
+        Human-readable name.
+    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
+        Time for wave animation in seconds. Default is 0.0.
+    material_front : MaterialField, optional
+        Material above surface (atmosphere).
+    material_back : MaterialField, optional
+        Material below surface (ocean water).
+
+    Examples
+    --------
+    >>> from lsurf.surfaces import CurvedWaveSurface, GerstnerWaveParams, SurfaceRole
+    >>> from lsurf.materials import ExponentialAtmosphere, WATER
+    >>>
+    >>> wave = GerstnerWaveParams(amplitude=1.0, wavelength=50.0)
+    >>> ocean = CurvedWaveSurface(
+    ...     wave_params=[wave],
+    ...     role=SurfaceRole.OPTICAL,
+    ...     name="ocean",
+    ...     material_front=ExponentialAtmosphere(),
+    ...     material_back=WATER,
+    ... )
+    """
+
+    wave_params: list[GerstnerWaveParams]
+    role: SurfaceRole
+    name: str = "curved_wave"
+    earth_center: tuple[float, float, float] = (0, 0, -EARTH_RADIUS)
+    earth_radius: float = EARTH_RADIUS
+    time: float = 0.0
+    material_front: Any = None
+    material_back: Any = None
+
+    # CPU-only surface
+    _gpu_capable: bool = field(default=False, init=False, repr=False)
+    _geometry_id: int = field(
+        default=0, init=False, repr=False
+    )  # CPU-only, no GPU geometry
+
+    # Precomputed wave arrays
+    _wave_amplitudes: NDArray = field(default=None, init=False, repr=False)
+    _wave_numbers: NDArray = field(default=None, init=False, repr=False)
+    _wave_frequencies: NDArray = field(default=None, init=False, repr=False)
+    _wave_directions: NDArray = field(default=None, init=False, repr=False)
+    _wave_phases: NDArray = field(default=None, init=False, repr=False)
+    _wave_steepness: NDArray = field(default=None, init=False, repr=False)
+    _earth_center_arr: NDArray = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self._earth_center_arr = np.array(self.earth_center, dtype=np.float64)
+        self._precompute_wave_params()
+
+    def _precompute_wave_params(self) -> None:
+        """Precompute wave parameters as arrays."""
+        n_waves = len(self.wave_params)
+        if n_waves == 0:
+            self._wave_amplitudes = np.array([], dtype=np.float64)
+            self._wave_numbers = np.array([], dtype=np.float64)
+            self._wave_frequencies = np.array([], dtype=np.float64)
+            self._wave_directions = np.zeros((0, 2), dtype=np.float64)
+            self._wave_phases = np.array([], dtype=np.float64)
+            self._wave_steepness = np.array([], dtype=np.float64)
+            return
+
+        self._wave_amplitudes = np.array(
+            [w.amplitude for w in self.wave_params], dtype=np.float64
+        )
+        self._wave_numbers = np.array(
+            [w.wave_number for w in self.wave_params], dtype=np.float64
+        )
+        self._wave_frequencies = np.array(
+            [w.angular_frequency for w in self.wave_params], dtype=np.float64
+        )
+        self._wave_directions = np.array(
+            [w.direction_normalized for w in self.wave_params], dtype=np.float64
+        )
+        self._wave_phases = np.array(
+            [w.phase for w in self.wave_params], dtype=np.float64
+        )
+        self._wave_steepness = np.array(
+            [w.steepness for w in self.wave_params], dtype=np.float64
+        )
+
+    @property
+    def gpu_capable(self) -> bool:
+        """This surface does NOT support GPU acceleration."""
+        return False
+
+    @property
+    def geometry_id(self) -> int:
+        """GPU geometry type ID (0 = CPU-only)."""
+        return 0
+
+    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 get_max_wave_height(self) -> float:
+        """Get maximum possible wave displacement."""
+        if len(self.wave_params) == 0:
+            return 0.0
+        return float(np.sum(self._wave_amplitudes))
+
+    def _get_local_tangent_basis(
+        self, positions: NDArray[np.float64]
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
+        """Compute local tangent basis at each position on the sphere."""
+        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
+
+        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)
+
+        return radial, tangent_x, tangent_y
+
+    def _compute_wave_displacement(
+        self,
+        local_x: NDArray[np.float64],
+        local_y: NDArray[np.float64],
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
+        """Compute Gerstner wave displacement in local coordinates."""
+        n_points = len(local_x)
+        dx = np.zeros(n_points, dtype=np.float64)
+        dy = np.zeros(n_points, dtype=np.float64)
+        dz = np.zeros(n_points, dtype=np.float64)
+
+        for i in range(len(self.wave_params)):
+            A = self._wave_amplitudes[i]
+            k = self._wave_numbers[i]
+            omega = self._wave_frequencies[i]
+            dir_x, dir_y = self._wave_directions[i]
+            phi = self._wave_phases[i]
+            Q = self._wave_steepness[i]
+
+            phase = k * (dir_x * local_x + dir_y * local_y) - omega * self.time + phi
+
+            cos_phase = np.cos(phase)
+            sin_phase = np.sin(phase)
+
+            dx -= Q * A * dir_x * sin_phase
+            dy -= Q * A * dir_y * sin_phase
+            dz += A * cos_phase
+
+        return dx, dy, dz
+
+    def _compute_wave_normal(
+        self,
+        local_x: NDArray[np.float64],
+        local_y: NDArray[np.float64],
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
+        """Compute wave surface normal in local coordinates."""
+        n_points = len(local_x)
+
+        nx = np.zeros(n_points, dtype=np.float64)
+        ny = np.zeros(n_points, dtype=np.float64)
+        nz = np.ones(n_points, dtype=np.float64)
+
+        for i in range(len(self.wave_params)):
+            A = self._wave_amplitudes[i]
+            k = self._wave_numbers[i]
+            omega = self._wave_frequencies[i]
+            dir_x, dir_y = self._wave_directions[i]
+            phi = self._wave_phases[i]
+            Q = self._wave_steepness[i]
+
+            phase = k * (dir_x * local_x + dir_y * local_y) - omega * self.time + phi
+
+            sin_phase = np.sin(phase)
+
+            WA = k * A
+
+            nx += dir_x * WA * sin_phase
+            ny += dir_y * WA * sin_phase
+            nz -= Q * WA * sin_phase
+
+        norm = np.sqrt(nx**2 + ny**2 + nz**2)
+        norm = np.maximum(norm, 1e-10)
+
+        return nx / norm, ny / norm, nz / norm
+
+    def intersect(
+        self,
+        origins: NDArray[np.float32],
+        directions: NDArray[np.float32],
+        min_distance: float = 1e-6,
+        max_iterations: int = 200,
+        tolerance: float = 1e-3,
+        max_distance: float | None = None,
+    ) -> tuple[NDArray[np.float32], 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)
+
+        max_wave_height = self.get_max_wave_height()
+
+        # Find intersection with outer sphere
+        outer_radius = self.earth_radius + max_wave_height
+        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)
+
+            local_x = positions[:, 0]
+            local_y = positions[:, 1]
+            _, _, wave_height = self._compute_wave_displacement(local_x, local_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
+
+            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._compute_wave_displacement(
+                        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._compute_wave_displacement(
+                    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 < -max_wave_height - 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, -max_wave_height * 2, max_wave_height * 2)
+            t[active] += step[active]
+            t = np.maximum(t, 0)
+
+        return distances.astype(np.float32), hit_mask
+
+    def normal_at(
+        self,
+        positions: NDArray[np.float32],
+        incoming_directions: NDArray[np.float32] | None = None,
+    ) -> 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)
+
+        radial, tangent_x, tangent_y = self._get_local_tangent_basis(positions)
+
+        local_x = positions[:, 0]
+        local_y = positions[:, 1]
+
+        nx_local, ny_local, nz_local = self._compute_wave_normal(local_x, local_y)
+
+        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("curved_wave", "cpu", cls=CurvedWaveSurface)