"""Built-in FFT ocean pass (design D11, RM-E7).
A dedicated forward pass (its own GLSL shader pair, like the water / grid passes)
drawn in the TRANSPARENT slot after the scene colour/depth copy (A7/A8), so it
refracts what is behind it. It reuses the :class:`~simvx.core.water.WaterMaterial`
shading model of the Gerstner water pass (design D16) but replaces the analytic
vertex displacement with a sampled FFT displacement map and the analytic normal
with the FFT slope map (design D11).
The FFT runs entirely on the GPU (:class:`.ocean_compute.OceanCompute`, design
D11 / RM-E7gpu): three compute shaders time-evolve the frozen spectrum, run a
Stockham radix-2 inverse FFT and assemble the per-cascade displacement / slope /
foam fields directly into two sampled texture-2D-arrays (one layer per cascade),
recorded into the frame command buffer BEFORE the main render pass begins
(compute writes to sampled images are illegal inside a render pass). The ocean
shader then samples them in the transparent slot. :mod:`.ocean_fft` stays the
canonical spectrum builder (frozen ``h0`` Phillips spectrum + dispersion,
uploaded once per :meth:`OceanFFT.configure`) and the CPU reference the tests
pin against; only the per-frame FFT hot path lives on the GPU (no per-frame
``numpy`` FFT, no per-frame field upload). The compute chain only runs when a
visible ``OceanSurface3D`` is in the scene, so with no ocean nothing here runs
and the frame is byte-identical (zero-cost when unused).
Set 0 is the shared forward descriptor set (FrameGlobals b13, environment cube
b4, scene colour/depth b14/b15); set 1 is this pass's own camera UBO (b0), the
per-frame surface SSBO (b1) and the displacement / slope texture arrays (b2/b3).
Cascade texture size + count come from the design-D13 quality table
(``QualitySettings.ocean_size`` / ``ocean_cascades``), resolved through
``WorldEnvironment.quality_tier`` so a tier means the same ocean everywhere.
"""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import (
DescriptorWriteBatch,
allocate_descriptor_set,
create_descriptor_set_layout,
create_pool_for_types,
write_ssbo_descriptor,
write_ubo_descriptor,
)
from ..gpu.memory import create_buffer, upload_numpy
from ..gpu.pipeline import PipelineSpec, build_pipeline, create_shader_module
from ..materials.shader_compiler import compile_shader
from .ocean_compute import OceanCompute
from .ocean_fft import OceanFFT
from .ocean_pack import OCEAN_SURFACE_DTYPE, pack_ocean_surface
__all__ = ["OceanPass", "OCEAN_SURFACE_DTYPE"]
log = logging.getLogger(__name__)
MAX_SURFACES = 16
_CAMERA_BYTES = 160 # mat4 view_proj + mat4 inv_view_proj + vec4 cam_pos + vec4 cam_planes
_SURFACE_BYTES = OCEAN_SURFACE_DTYPE.itemsize
# Appearance knobs not carried on WaterMaterial (the FFT-specific ones). Kept as
# module constants for now; promoting them to OceanSurface3D properties is a
# trivial follow-up if a game needs per-ocean control.
_FOAM_THRESHOLD = 0.55
_FOAM_INTENSITY = 1.0
_DISPLACEMENT_SCALE = 1.0
[docs]
class OceanPass:
"""Renders every submitted ``OceanSurface3D`` as an FFT ocean grid."""
def __init__(self, engine: Any):
self._engine = engine
self._pipeline: Any = None
self._pipeline_layout: Any = None
self._vert_module: Any = None
self._frag_module: Any = None
self._set_layout: Any = None
self._pool: Any = None
self._set: Any = None
self._cam_buf: Any = None
self._cam_mem: Any = None
self._surf_buf: Any = None
self._surf_mem: Any = None
self._forward_ssbo_layout: Any = None
self._sampler: Any = None
# CPU spectrum builder + the GPU compute FFT that writes the displacement
# / slope arrays the ocean shader samples (design D11 / RM-E7gpu).
self._fft: OceanFFT | None = None
self._compute = OceanCompute(engine)
self._size = 0
self._cascades = 0
self._ready = False
# Per-frame submissions: (model_matrix, WaterMaterial, subdivisions, size).
self._submissions: list[tuple[np.ndarray, Any, int, tuple[float, float]]] = []
# -------------------------------------------------------------- setup
[docs]
def setup(self, ssbo_layout: Any, render_pass: Any = None) -> None:
"""Create GPU resources. ``ssbo_layout`` is the shared forward set-0 layout."""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
self._forward_ssbo_layout = ssbo_layout
vf = vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT
self._set_layout = create_descriptor_set_layout(
device,
[
(0, vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, vf, 1),
(1, vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, vf, 1),
(2, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, vf, 1), # displacement (vertex + frag)
(3, vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, vk.VK_SHADER_STAGE_FRAGMENT_BIT, 1), # slope
],
)
self._pool = create_pool_for_types(
device,
{
vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER: 1,
vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER: 1,
vk.VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER: 2,
},
max_sets=1,
)
self._set = allocate_descriptor_set(device, self._pool, self._set_layout)
host = vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
self._cam_buf, self._cam_mem = create_buffer(
device, phys, _CAMERA_BYTES, vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, host
)
surf_size = MAX_SURFACES * _SURFACE_BYTES
self._surf_buf, self._surf_mem = create_buffer(
device, phys, surf_size, vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, host
)
upload_numpy(device, self._cam_mem, np.zeros(_CAMERA_BYTES, dtype=np.uint8))
upload_numpy(device, self._surf_mem, np.zeros(surf_size, dtype=np.uint8))
write_ubo_descriptor(device, self._set, 0, self._cam_buf, _CAMERA_BYTES)
write_ssbo_descriptor(device, self._set, 1, self._surf_buf, surf_size)
self._create_sampler()
# GPU compute FFT owns the displacement / slope arrays; a placeholder 1x1
# resize keeps the render descriptor set (b2/b3) valid before the first
# ocean frame resolves the real cascade size/count.
self._compute.setup()
self._compute.resize(OceanFFT(size=1, cascades=1), 1, 1)
self._write_texture_descriptors()
shader_dir = e.shader_dir
self._vert_module = create_shader_module(device, compile_shader(shader_dir / "ocean.vert"))
self._frag_module = create_shader_module(device, compile_shader(shader_dir / "ocean.frag"))
self._create_pipeline(device, render_pass or e.render_pass, e.extent)
self._ready = True
log.debug("Ocean pass initialized")
def _create_sampler(self) -> None:
self._sampler = vk.vkCreateSampler(
self._engine.ctx.device,
vk.VkSamplerCreateInfo(
magFilter=vk.VK_FILTER_LINEAR,
minFilter=vk.VK_FILTER_LINEAR,
addressModeU=vk.VK_SAMPLER_ADDRESS_MODE_REPEAT,
addressModeV=vk.VK_SAMPLER_ADDRESS_MODE_REPEAT,
addressModeW=vk.VK_SAMPLER_ADDRESS_MODE_REPEAT,
anisotropyEnable=vk.VK_FALSE,
unnormalizedCoordinates=vk.VK_FALSE,
mipmapMode=vk.VK_SAMPLER_MIPMAP_MODE_NEAREST,
),
None,
)
def _write_texture_descriptors(self) -> None:
"""Point b2/b3 (displacement + slope) at the compute-owned output arrays."""
with DescriptorWriteBatch(self._engine.ctx.device) as batch:
batch.image(
self._set,
2,
self._compute.disp_view,
self._sampler,
image_layout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
)
batch.image(
self._set,
3,
self._compute.grad_view,
self._sampler,
image_layout=vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
)
def _create_pipeline(self, device: Any, render_pass: Any, extent: tuple[int, int]) -> None:
spec = PipelineSpec(
name="ocean",
topology=vk.VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
vertex_stride=0,
cull_mode=vk.VK_CULL_MODE_NONE,
depth_test=True,
depth_write=False,
depth_compare=vk.VK_COMPARE_OP_LESS_OR_EQUAL,
blend="alpha",
dst_alpha_factor=vk.VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
set_layouts=(self._forward_ssbo_layout, self._set_layout),
push_size=16,
push_stages=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
attachment_count=(2 if self._engine.gbuffer_active else 1),
)
self._pipeline, self._pipeline_layout = build_pipeline(
device,
spec,
render_pass,
extent,
vert_module=self._vert_module,
frag_module=self._frag_module,
)
[docs]
def rebuild_pipeline(self, render_pass: Any) -> None:
if not self._ready:
return
device = self._engine.ctx.device
if self._pipeline:
vk.vkDestroyPipeline(device, self._pipeline, None)
if self._pipeline_layout:
vk.vkDestroyPipelineLayout(device, self._pipeline_layout, None)
self._create_pipeline(device, render_pass, self._engine.extent)
# -------------------------------------------------------------- per-frame
[docs]
def begin_frame(self) -> None:
self._submissions.clear()
[docs]
def submit(self, model_matrix: np.ndarray, material: Any, subdivisions: int, size: tuple[float, float]) -> None:
self._submissions.append((model_matrix, material, int(subdivisions), (float(size[0]), float(size[1]))))
[docs]
@property
def has_surfaces(self) -> bool:
return bool(self._submissions)
def _resolve_quality(self, env: Any) -> tuple[int, int]:
"""Cascade texture ``(size, cascades)`` from the design-D13 quality table."""
from simvx.core.graphics_quality import QualitySettings, resolve_tier
qs = resolve_tier(getattr(env, "quality_tier", "auto")) if env is not None else None
qs = qs or QualitySettings()
return int(qs.ocean_size), int(qs.ocean_cascades)
@staticmethod
def _wind_from_env(env: Any) -> tuple[tuple[float, float], float]:
"""(wind_dir XZ, wind_speed m/s) from the WorldEnvironment wind knobs.
``wind_strength`` is a 0..1 authoring knob (as the Gerstner water reads
it); it maps to a plausible open-sea wind speed so the Phillips fetch is
physically sane without a second knob.
"""
if env is None:
return (1.0, 0.0), 8.0
d = getattr(env, "wind_direction", (1.0, 0.0))
strength = float(getattr(env, "wind_strength", 0.0))
return (float(d[0]), float(d[1])), 4.0 + strength * 16.0
[docs]
def prepare(self, cmd: Any, env: Any, time: float) -> None:
"""Record the GPU FFT compute chain into ``cmd`` (design D11 / RM-E7gpu).
MUST be called before the main render pass begins (the compute chain
writes the sampled displacement / slope arrays, illegal inside a render
pass) and only when :attr:`has_surfaces`. Rebuilds the spectrum on a
quality change, re-uploads the frozen ``h0`` buffer on a wind change, then
dispatches spectrum -> Stockham IFFT -> assemble.
"""
if not self._ready or not self._submissions:
return
size, cascades = self._resolve_quality(env)
if self._fft is None or (size, cascades) != (self._size, self._cascades):
self._fft = OceanFFT(size=size, cascades=cascades)
self._compute.resize(self._fft, size, cascades)
self._write_texture_descriptors()
self._size, self._cascades = size, cascades
wind_dir, wind_speed = self._wind_from_env(env)
mat = self._submissions[0][1]
self._fft.configure(
wind_dir=wind_dir,
wind_speed=wind_speed,
amplitude=float(mat.wave_amplitude),
direction_spread=float(mat.direction_spread),
)
self._compute.upload_spectrum_if_dirty(self._fft)
choppiness = 0.5 + float(mat.wave_steepness)
pool = self._engine.current_timestamp_pool
if pool is not None:
pool.begin(cmd, "ocean_fft")
self._compute.dispatch(cmd, float(time), choppiness=choppiness, foam_threshold=_FOAM_THRESHOLD, foam_decay=0.92)
if pool is not None:
pool.end(cmd, "ocean_fft")
[docs]
def render(
self,
cmd: Any,
view_proj: np.ndarray,
cam_pos: np.ndarray,
near: float,
far: float,
has_skybox: bool,
extent: tuple[int, int],
ssbo_set: Any,
) -> None:
"""Record the ocean draws. ``ssbo_set`` is the shared forward set 0."""
subs = self._submissions
if not self._ready or not subs or self._fft is None:
return
device = self._engine.ctx.device
vp_t = np.ascontiguousarray(view_proj.T, dtype=np.float32)
inv_vp_t = np.ascontiguousarray(np.linalg.inv(view_proj).T, dtype=np.float32)
cam = np.zeros(_CAMERA_BYTES // 4, dtype=np.float32)
cam[0:16] = vp_t.ravel()
cam[16:32] = inv_vp_t.ravel()
cam[32:35] = cam_pos[:3]
cam[36:40] = (float(near), float(far), 1.0 if has_skybox else 0.0, 0.0)
upload_numpy(device, self._cam_mem, cam.view(np.uint8))
count = min(len(subs), MAX_SURFACES)
patch_lengths = self._fft.patch_lengths
packed = np.zeros(count, dtype=OCEAN_SURFACE_DTYPE)
for i in range(count):
model, mat, sub, size = subs[i]
choppiness = 0.5 + float(mat.wave_steepness)
packed[i] = pack_ocean_surface(
model,
mat,
sub,
size,
patch_lengths,
choppiness=choppiness,
displacement_scale=_DISPLACEMENT_SCALE,
foam_intensity=_FOAM_INTENSITY,
)
upload_numpy(device, self._surf_mem, packed.view(np.uint8))
vk_vp = vk.VkViewport(x=0.0, y=0.0, width=float(extent[0]), height=float(extent[1]), minDepth=0.0, maxDepth=1.0)
vk.vkCmdSetViewport(cmd, 0, 1, [vk_vp])
scissor = vk.VkRect2D(offset=vk.VkOffset2D(x=0, y=0), extent=vk.VkExtent2D(width=extent[0], height=extent[1]))
vk.vkCmdSetScissor(cmd, 0, 1, [scissor])
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout, 0, 1, [ssbo_set], 0, None
)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout, 1, 1, [self._set], 0, None
)
ffi = vk.ffi
for i in range(count):
_model, _mat, sub, _size = subs[i]
n = max(int(sub), 1)
pc = np.array([i, 0, 0, 0], dtype=np.int32)
cbuf = ffi.new("char[]", pc.tobytes())
vk._vulkan.lib.vkCmdPushConstants(
cmd,
self._pipeline_layout,
vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
0,
16,
cbuf,
)
vk.vkCmdDraw(cmd, n * n * 6, 1, 0, 0)
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
for obj, fn in [
(self._pipeline, vk.vkDestroyPipeline),
(self._pipeline_layout, vk.vkDestroyPipelineLayout),
(self._vert_module, vk.vkDestroyShaderModule),
(self._frag_module, vk.vkDestroyShaderModule),
(self._set_layout, vk.vkDestroyDescriptorSetLayout),
(self._pool, vk.vkDestroyDescriptorPool),
(self._sampler, vk.vkDestroySampler),
]:
if obj:
fn(device, obj, None)
for buf, mem in [(self._cam_buf, self._cam_mem), (self._surf_buf, self._surf_mem)]:
if buf:
vk.vkDestroyBuffer(device, buf, None)
if mem:
vk.vkFreeMemory(device, mem, None)
self._compute.cleanup()
self._ready = False