"""FFT ocean spectrum + inverse transform (design D11, RM-E7).
Dependency-free (numpy only, no Vulkan / no core-node imports) so BOTH the
desktop ocean pass (``ocean_pass.py``) and the future browser twin (Pyodide,
where ``vulkan`` is unavailable) drive the identical Tessendorf spectrum. One
canonical spectrum + inverse-FFT, producing per-cascade displacement, slope and
foam fields that the ocean shader (``ocean.vert`` / ``ocean.frag``, later
``ocean3d.wgsl``) samples.
Model (Tessendorf 2001, the film/game-standard FFT ocean):
* A Phillips spectrum with directional spreading and an against-wind damp seeds
a complex, Hermitian-symmetric height spectrum ``h0(k)`` per cascade. Each
cascade tiles a square patch of ``patch_length`` metres at ``size``x``size``
resolution, so short cascades add crisp chop and long cascades add slow swell.
* Per frame the spectrum is evolved to time ``t`` with the deep-water dispersion
``omega = sqrt(g |k|)`` and inverse-FFT'd to a real height field, horizontal
choppy displacement, and surface slope (for the normal).
* Foam comes from the Jacobian of the horizontal displacement: where the surface
folds onto itself (Jacobian below a threshold) foam is deposited and then
decays, so crests keep a lace of surf that dissolves in the troughs.
The wind (direction + strength) is taken from the renderer's FrameGlobals wind
slot, exactly like the built-in Gerstner :class:`WaterMaterial`; the FFT ocean
replaces the analytic Gerstner displacement with a sampled displacement map and
feeds the derived normal + foam into the same water shading (design D16).
Design D11 specifies a GPU Stockham radix-2 compute FFT. On desktop Vulkan the
per-frame hot path (time-evolve the spectrum + the 2D inverse FFT + assemble
displacement / slope / Jacobian-foam) runs entirely on the GPU
(``ocean_compute.py`` + ``ocean_spectrum.comp`` / ``ocean_fft.comp`` /
``ocean_assemble.comp``). This module stays the canonical spectrum builder: it
constructs the frozen ``h0`` Phillips spectrum + the dispersion / k-vectors once
per :meth:`configure` (uploaded once to the GPU via :meth:`gpu_static_spectrum`),
and keeps a pure-``numpy`` reference of the exact per-frame maths the shaders
run (:func:`stockham_ifft2` + :meth:`gpu_reference`) so the GPU FFT is verifiable
headlessly against ``numpy.fft`` without an interactive GPU debug loop. The
legacy CPU :meth:`evaluate` is retained as the reference transform the unit
tests pin against.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
__all__ = ["GRAVITY", "OceanCascade", "OceanFFT", "cascade_patch_lengths", "stockham_ifft2"]
GRAVITY = 9.81
# Default cascade patch lengths in metres (design D11: ~250 / 60 / 15 m). The
# first ``n`` are used for an n-cascade config, so "low" (2 cascades) keeps the
# big swell + the mid chop and drops the finest ripple band.
_PATCH_LENGTHS: tuple[float, ...] = (250.0, 60.0, 15.0)
# Per-cascade RMS height weighting: the long swell carries most of the height,
# finer cascades add proportionally less so the sum does not blow out.
_CASCADE_WEIGHTS: tuple[float, ...] = (1.0, 0.5, 0.28)
[docs]
def cascade_patch_lengths(cascades: int) -> tuple[float, ...]:
"""The first ``cascades`` default patch lengths (metres), longest first."""
n = max(1, min(int(cascades), len(_PATCH_LENGTHS)))
return _PATCH_LENGTHS[:n]
# ------------------------------------------------------- Stockham reference FFT
def _stockham_1d(x: np.ndarray, axis: int) -> np.ndarray:
"""Radix-2 Stockham autosort *inverse* transform along ``axis`` (unscaled).
The exact butterfly the GPU ``ocean_fft.comp`` runs, kept in numpy so the
shader is verifiable against ``numpy.fft`` headlessly. Operates on a 2D
complex array; ``N/2`` butterflies per stage, ``log2(N)`` stages, ping-pong
between two buffers (no bit-reversal pass: autosort leaves natural order).
The ``1/N`` inverse scale is applied by the caller.
"""
a = np.moveaxis(np.asarray(x, dtype=np.complex128), axis, 0).copy()
n = a.shape[0]
half = n >> 1
t = int(np.log2(n))
b = np.empty_like(a)
for stage in range(t):
span = 1 << stage
span2 = span << 1
idx = np.arange(half)
j = idx & (span - 1)
k = idx >> stage
ang = 2.0 * np.pi * j / span2 # inverse transform: +i
w = (np.cos(ang) + 1j * np.sin(ang)).reshape((half,) + (1,) * (a.ndim - 1))
av = a[k * span + j]
tv = w * a[k * span + j + half]
b[k * span2 + j] = av + tv
b[k * span2 + j + span] = av - tv
a, b = b, a
return np.moveaxis(a, 0, axis)
[docs]
def stockham_ifft2(spec: np.ndarray) -> np.ndarray:
"""2D inverse FFT via the Stockham butterfly (rows then columns), ``1/N^2``.
Numerically matches :func:`numpy.fft.ifft2` to float round-off; it is the
CPU twin of the GPU ``ocean_fft.comp`` row-pass + column-pass chain.
"""
n = spec.shape[0]
rows = _stockham_1d(spec, axis=1)
cols = _stockham_1d(rows, axis=0)
return cols / float(n * n)
[docs]
@dataclass
class OceanCascade:
"""One FFT cascade: a square ``size``x``size`` patch of ``patch_length`` m.
``h0`` / ``h0_mk_conj`` are the time-invariant spectrum (recomputed only when
the wind/config changes); the per-cascade ``foam`` field accumulates across
``evaluate`` calls so surf lingers on crests and fades in troughs.
"""
size: int
patch_length: float
weight: float
# Angular wave-vector grid (numpy fft ordering) and derived helpers.
kx: np.ndarray = field(default_factory=lambda: np.zeros(0))
kz: np.ndarray = field(default_factory=lambda: np.zeros(0))
kmag: np.ndarray = field(default_factory=lambda: np.zeros(0))
kx_unit: np.ndarray = field(default_factory=lambda: np.zeros(0))
kz_unit: np.ndarray = field(default_factory=lambda: np.zeros(0))
omega: np.ndarray = field(default_factory=lambda: np.zeros(0))
h0: np.ndarray = field(default_factory=lambda: np.zeros(0))
h0_mk_conj: np.ndarray = field(default_factory=lambda: np.zeros(0))
gain: float = 1.0
foam: np.ndarray = field(default_factory=lambda: np.zeros(0))
[docs]
class OceanFFT:
"""Multi-cascade Tessendorf ocean spectrum + inverse transform.
Construct once for a given ``(size, cascades)`` quality config (design D13
dials), call :meth:`configure` whenever the wind/appearance changes, then
:meth:`evaluate` once per frame to get the sampled surface fields.
"""
def __init__(self, size: int, cascades: int, *, seed: int = 1337) -> None:
self.size = int(size)
self.n_cascades = max(1, min(int(cascades), len(_PATCH_LENGTHS)))
self._seed = int(seed)
self._last_time: float | None = None
self._config_key: tuple | None = None
# Bumped every time the frozen spectrum is (re)built; the GPU ocean pass
# compares it to know when to re-upload the static h0 buffer.
self.spectrum_generation = 0
lengths = cascade_patch_lengths(self.n_cascades)
self.cascades: list[OceanCascade] = [
OceanCascade(size=self.size, patch_length=lengths[c], weight=_CASCADE_WEIGHTS[c])
for c in range(self.n_cascades)
]
for cas in self.cascades:
self._build_grid(cas)
cas.foam = np.zeros((self.size, self.size), dtype=np.float32)
# ------------------------------------------------------------------ grid
def _build_grid(self, cas: OceanCascade) -> None:
n = cas.size
# Angular wave numbers in numpy-fft ordering: 2*pi*m/L for m in fftfreq.
m = np.fft.fftfreq(n, d=1.0 / n) # 0,1,..,n/2-1,-n/2,..,-1
base = 2.0 * np.pi / cas.patch_length
kx, kz = np.meshgrid(base * m, base * m) # kx varies over columns, kz over rows
kmag = np.sqrt(kx * kx + kz * kz)
safe = np.where(kmag > 1e-6, kmag, 1.0)
cas.kx, cas.kz, cas.kmag = kx.astype(np.float64), kz.astype(np.float64), kmag
cas.kx_unit = np.where(kmag > 1e-6, kx / safe, 0.0)
cas.kz_unit = np.where(kmag > 1e-6, kz / safe, 0.0)
cas.omega = np.sqrt(GRAVITY * kmag)
# ------------------------------------------------------------- spectrum
def _build_spectrum(
self,
cas: OceanCascade,
wdir: np.ndarray,
wind_speed: float,
amplitude: float,
spread: float,
rng: np.random.Generator,
) -> None:
kmag = cas.kmag
k2 = kmag * kmag
big = kmag > 1e-6
fetch = max(wind_speed, 0.5) ** 2 / GRAVITY # Phillips L = V^2/g
# Directional term: |k.w|^(2*spread), with an against-wind damp so the
# sea travels downwind. Squared cosine keeps the base P(k)=P(-k) even
# part; the sign-based damp is what the Hermitian construction repairs.
dotp = cas.kx_unit * wdir[0] + cas.kz_unit * wdir[1]
direction = np.abs(dotp) ** (2.0 * max(spread, 0.1))
phillips = np.zeros_like(kmag)
with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
spec = np.exp(-1.0 / (k2 * fetch * fetch)) / (k2 * k2) * direction
spec *= np.exp(-k2 * (fetch * 1e-3) ** 2) # suppress the tiniest ripples
phillips[big] = spec[big]
phillips = np.where(dotp < 0.0, phillips * 0.07, phillips) # against-wind damp
phillips = np.nan_to_num(phillips, nan=0.0, posinf=0.0, neginf=0.0)
xr = rng.standard_normal((cas.size, cas.size))
xi = rng.standard_normal((cas.size, cas.size))
h0 = (xr + 1j * xi) / np.sqrt(2.0) * np.sqrt(phillips)
# h0 sampled at -k (index negation, numpy-fft wrap): the Hermitian pair.
idx = np.mod(-np.arange(cas.size), cas.size)
h0_mk = h0[np.ix_(idx, idx)]
cas.h0 = h0
cas.h0_mk_conj = np.conj(h0_mk)
# Fix a per-cascade gain so the t=0 height RMS matches the weighted target
# (stable across frames: gain is not recomputed per frame, so the sea does
# not shimmer in overall scale).
raw = self._height_field(cas, 0.0)
rms = float(np.sqrt(np.mean(raw * raw))) + 1e-9
target = amplitude * cas.weight
cas.gain = target / rms
cas.foam = np.zeros((cas.size, cas.size), dtype=np.float32)
def _time_spectrum(self, cas: OceanCascade, t: float) -> np.ndarray:
"""Evolve h0 to time ``t`` (Hermitian, so the inverse FFT is real)."""
phase = cas.omega * t
expp = np.cos(phase) + 1j * np.sin(phase)
return cas.h0 * expp + cas.h0_mk_conj * np.conj(expp)
def _height_field(self, cas: OceanCascade, t: float) -> np.ndarray:
return np.real(np.fft.ifft2(self._time_spectrum(cas, t)))
# ------------------------------------------------------------- per-frame
[docs]
def evaluate(
self,
time: float,
*,
choppiness: float = 1.0,
foam_threshold: float = 0.6,
foam_decay: float = 0.92,
) -> tuple[np.ndarray, np.ndarray]:
"""Return ``(disp, grad)`` float16 arrays for the current frame.
``disp`` is ``(cascades, size, size, 4)`` = ``[Dx, height, Dz, foam]``;
``grad`` is ``(cascades, size, size, 4)`` = ``[slope_x, slope_z, 0, 0]``.
Both are contiguous, ready to upload straight into a texture-2D-array
(one layer per cascade). Foam accumulates across calls; a repeated
``time`` is a no-op that returns the last result (so a double-eval in a
pipelined frame does not double-advance the surf).
"""
c = float(choppiness)
disp = np.zeros((self.n_cascades, self.size, self.size, 4), dtype=np.float32)
grad = np.zeros((self.n_cascades, self.size, self.size, 4), dtype=np.float32)
advance = self._last_time is None or time != self._last_time
for ci, cas in enumerate(self.cascades):
hk = self._time_spectrum(cas, time) * cas.gain
kx, kz, kxu, kzu = cas.kx, cas.kz, cas.kx_unit, cas.kz_unit
height = np.real(np.fft.ifft2(hk))
dx = np.real(np.fft.ifft2(-1j * kxu * hk)) * c
dz = np.real(np.fft.ifft2(-1j * kzu * hk)) * c
slope_x = np.real(np.fft.ifft2(1j * kx * hk))
slope_z = np.real(np.fft.ifft2(1j * kz * hk))
# Jacobian of the horizontal displacement -> folding -> foam.
jxx = np.real(np.fft.ifft2((kx * kxu) * hk)) * c
jzz = np.real(np.fft.ifft2((kz * kzu) * hk)) * c
jxz = np.real(np.fft.ifft2((kx * kzu) * hk)) * c
jac = (1.0 + jxx) * (1.0 + jzz) - jxz * jxz
fresh = np.clip(foam_threshold - jac, 0.0, 1.0).astype(np.float32)
if advance:
cas.foam = np.maximum(cas.foam * foam_decay, fresh)
disp[ci, :, :, 0] = dx
disp[ci, :, :, 1] = height
disp[ci, :, :, 2] = dz
disp[ci, :, :, 3] = cas.foam
grad[ci, :, :, 0] = slope_x
grad[ci, :, :, 1] = slope_z
if advance:
self._last_time = time
return disp.astype(np.float16), grad.astype(np.float16)
[docs]
@property
def patch_lengths(self) -> tuple[float, ...]:
return tuple(cas.patch_length for cas in self.cascades)
# ----------------------------------------------------------------- GPU path
[docs]
def gpu_static_spectrum(self) -> np.ndarray:
"""The frozen per-cascade spectrum, packed for the GPU compute FFT.
Returns a contiguous ``(cascades, N, N, 8)`` float32 array = two vec4 per
texel: ``[h0.re, h0.im, h0_mk_conj.re, h0_mk_conj.im]`` then
``[kx, kz, 0, 0]``. The per-cascade ``gain`` is baked into ``h0`` /
``h0_mk_conj`` (so the shader-side ``hk`` already carries the target RMS),
exactly as :meth:`evaluate` multiplies ``self._time_spectrum * gain``.
``kx_unit`` / ``kz_unit`` / ``omega`` are re-derived in the shader from
``(kx, kz)`` with the identical formulae (dispersion + safe normalise).
"""
out = np.zeros((self.n_cascades, self.size, self.size, 8), dtype=np.float32)
for ci, cas in enumerate(self.cascades):
h0 = cas.h0 * cas.gain
h0c = cas.h0_mk_conj * cas.gain
out[ci, :, :, 0] = h0.real
out[ci, :, :, 1] = h0.imag
out[ci, :, :, 2] = h0c.real
out[ci, :, :, 3] = h0c.imag
out[ci, :, :, 4] = cas.kx
out[ci, :, :, 5] = cas.kz
return np.ascontiguousarray(out)
[docs]
def gpu_reference(
self,
time: float,
*,
choppiness: float = 1.0,
foam_threshold: float = 0.6,
foam_decay: float = 0.92,
) -> tuple[np.ndarray, np.ndarray]:
"""CPU twin of the GPU compute chain (spectrum -> Stockham IFFT -> assemble).
Mirrors ``ocean_spectrum.comp`` (the 2-real-per-complex packing), the
``ocean_fft.comp`` Stockham row/column IFFT and ``ocean_assemble.comp``
field assembly, using the SAME derived ``kx_unit`` / ``omega`` the shader
computes from the uploaded ``(kx, kz)``. Returns float32 ``(disp, grad)``
with the same layout as :meth:`evaluate`; used by the unit tests to pin
the shader maths against ``numpy`` headlessly (foam accumulation excluded,
it is an identical post-step in both paths).
"""
c = float(choppiness)
disp = np.zeros((self.n_cascades, self.size, self.size, 4), dtype=np.float32)
grad = np.zeros((self.n_cascades, self.size, self.size, 4), dtype=np.float32)
for ci, cas in enumerate(self.cascades):
h0 = cas.h0 * cas.gain
h0c = cas.h0_mk_conj * cas.gain
kx, kz = cas.kx, cas.kz
kmag = np.sqrt(kx * kx + kz * kz)
big = kmag > 1e-6
safe = np.where(big, kmag, 1.0)
kxu = np.where(big, kx / safe, 0.0)
kzu = np.where(big, kz / safe, 0.0)
omega = np.sqrt(GRAVITY * kmag)
expp = np.cos(omega * time) + 1j * np.sin(omega * time)
hk = h0 * expp + h0c * np.conj(expp)
# Pack two Hermitian spectra per complex FFT: Z = A + i*B -> real=IFFT(A),
# imag=IFFT(B). Four IFFTs cover the eight real fields.
c0 = stockham_ifft2(hk + 1j * (-1j * kxu * hk))
c1 = stockham_ifft2((-1j * kzu * hk) + 1j * (1j * kx * hk))
c2 = stockham_ifft2((1j * kz * hk) + 1j * (kx * kxu * hk))
c3 = stockham_ifft2((kz * kzu * hk) + 1j * (kx * kzu * hk))
height, dx = c0.real, c0.imag
dz, slope_x = c1.real, c1.imag
slope_z, jxx = c2.real, c2.imag
jzz, jxz = c3.real, c3.imag
jac = (1.0 + jxx * c) * (1.0 + jzz * c) - (jxz * c) * (jxz * c)
fresh = np.clip(foam_threshold - jac, 0.0, 1.0)
disp[ci, :, :, 0] = dx * c
disp[ci, :, :, 1] = height
disp[ci, :, :, 2] = dz * c
disp[ci, :, :, 3] = fresh
grad[ci, :, :, 0] = slope_x
grad[ci, :, :, 1] = slope_z
return disp, grad