851 lines
35 KiB
Python
851 lines
35 KiB
Python
# CelerisLab/simulation.py
|
|
"""
|
|
Top-level orchestrator -- assembles LBM field, stepper, body manager,
|
|
and CUDA context into a single coherent simulation.
|
|
|
|
Usage::
|
|
|
|
sim = Simulation("configs/config_lbm.json")
|
|
sim.add_body("circle", center=(50, 50), radius=10)
|
|
sim.initialize()
|
|
sim.run(1000)
|
|
macro = sim.get_macroscopic() # {"rho": ..., "ux": ..., "uy": ...}
|
|
force = sim.read_force(0) # force on body 0
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, Optional, Tuple, Any
|
|
|
|
import os
|
|
import operator
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
from .config import LBMConfig, BodyConfig, load_lbm_config, load_body_config
|
|
from .cuda.context import CudaContext, context_bound
|
|
from .cuda import compiler_v2 as compiler
|
|
from .lbm.field import LBMField
|
|
from .lbm.stepper import LBMStepper
|
|
from .body.objects import SimObject
|
|
from .body.manager import ObjectManager
|
|
from .body.ddf_patch import patch_ddf_for_body_sync, upload_patched_ddf
|
|
|
|
|
|
class Simulation:
|
|
"""High-level simulation handle.
|
|
|
|
LBM field/stepper and body manager live at the same level;
|
|
this class orchestrates them.
|
|
"""
|
|
|
|
def __init__(self,
|
|
lbm_config_path: Optional[str] = None,
|
|
body_config_path: Optional[str] = None,
|
|
device_id: int = 0):
|
|
# Internal stream for async DTOH coordination
|
|
self._stream: Optional[cuda.Stream] = None
|
|
self._build_dir = None
|
|
self._closed = False
|
|
|
|
# Load configs
|
|
self.lbm_cfg = load_lbm_config(lbm_config_path)
|
|
self.lbm_cfg.validate_runtime()
|
|
self.body_cfg = load_body_config(body_config_path)
|
|
|
|
# CUDA context
|
|
self.ctx = CudaContext(device_id)
|
|
try:
|
|
with self.ctx.activate():
|
|
arch = self._resolve_compile_arch()
|
|
self._build_dir = compiler.generate_config(self.lbm_cfg, n_objects=0)
|
|
self._ptx_path = compiler.compile_kernel(arch=arch, build_dir=self._build_dir)
|
|
self._module = compiler.load_module(self._ptx_path)
|
|
self._compiled_n_objects = 0
|
|
self._allocation_signature = self.lbm_cfg.allocation_signature()
|
|
self._compiled_viscosity = self.lbm_cfg.viscosity
|
|
|
|
self.field = LBMField(self.lbm_cfg, self._module)
|
|
self.field.ctx = self.ctx
|
|
self.stepper = LBMStepper(self.field, self._module, self.lbm_cfg)
|
|
self.bodies = ObjectManager(
|
|
self.lbm_cfg.nx, self.lbm_cfg.ny, self.lbm_cfg.nz,
|
|
self.lbm_cfg.nq, self.lbm_cfg,
|
|
)
|
|
self._initialized = False
|
|
self._load_objects_from_body_config()
|
|
if self.bodies.count > 0:
|
|
self.recompile()
|
|
self._assert_object_count_contract(expected_count=self.bodies.count)
|
|
except Exception:
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
# Construction must report its original failure. Cleanup is
|
|
# best effort here; an explicit close failure remains retryable.
|
|
pass
|
|
raise
|
|
|
|
# -- Stream access -------------------------------------------------------
|
|
@property
|
|
@context_bound
|
|
def stream(self) -> cuda.Stream:
|
|
"""Internal CUDA stream for async DTOH coordination."""
|
|
if self._stream is None:
|
|
self._stream = cuda.Stream()
|
|
return self._stream
|
|
|
|
# -- Object management ---------------------------------------------------
|
|
def add_body(self, type: str = "circle", center=None,
|
|
radius: float = 0.0, *, surface_actuation=None, **kwargs) -> int:
|
|
"""Add a simulation body. Returns body_id.
|
|
|
|
Before ``initialize()`` the object is added to the formal registry
|
|
immediately. After ``initialize()`` the object is staged via
|
|
``stage_add()`` and will be committed at the next ``sync_bodies()``.
|
|
|
|
Args:
|
|
type: ``"circle"``, ``"sensor"``, or ``"force_region"``
|
|
(future: ``"polygon"``, ``"mesh"``).
|
|
center: (x, y) center coordinates (lattice units).
|
|
radius: body radius (lattice units).
|
|
surface_actuation: Fixed circle aperture dict/SurfaceActuation:
|
|
angle, half_width (radians), normal_velocity (positive outward).
|
|
D2Q9/FP32/unshifted/DB only; set before initialization or stage
|
|
a new body. No runtime motion or amplitude setter.
|
|
**kwargs: passed to the geometry constructor.
|
|
|
|
Returns:
|
|
int body_id (formal id before initialize; -1 if staged).
|
|
"""
|
|
if self.lbm_cfg.y_wall_bc == "on_site_velocity":
|
|
raise ValueError("on_site_velocity currently requires an object-free domain")
|
|
type_lower = str(type).strip().lower()
|
|
if surface_actuation is not None and type_lower != "circle":
|
|
raise ValueError("Surface actuation is only supported on circles")
|
|
|
|
if type_lower == "circle":
|
|
from .body.geometry.circle import CircleGeometry
|
|
cx, cy = float(center[0]), float(center[1])
|
|
geom = CircleGeometry(cx, cy, radius)
|
|
obj = SimObject(obj_id=-1, geometry=geom,
|
|
center=center, radius=radius,
|
|
is_sensor=False, surface_actuation=surface_actuation)
|
|
elif type_lower == "sensor":
|
|
from .body.geometry.circle import CircleGeometry
|
|
cx, cy = float(center[0]), float(center[1])
|
|
geom = CircleGeometry(cx, cy, radius)
|
|
obj = SimObject(obj_id=-1, geometry=geom,
|
|
center=center, radius=radius,
|
|
is_sensor=True)
|
|
elif type_lower == "force_region":
|
|
from .body.geometry.circle import CircleGeometry
|
|
cx, cy = float(center[0]), float(center[1])
|
|
geom = CircleGeometry(cx, cy, radius)
|
|
obj = SimObject(obj_id=-1, geometry=geom,
|
|
center=center, radius=radius,
|
|
is_force_region=True)
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown body type '{type}'. "
|
|
"Supported: 'circle', 'sensor', 'force_region'."
|
|
)
|
|
|
|
if self._initialized:
|
|
self.bodies.stage_add(obj)
|
|
return -1
|
|
return self.bodies.add(obj)
|
|
|
|
def add_jet(self, jet) -> int:
|
|
"""Register a body-independent Jet before ``initialize()``; returns its index.
|
|
|
|
The jet drives every scene cut link whose true hit point lies in its
|
|
region (overlaps rejected at initialize). Not checkpointed.
|
|
"""
|
|
from .body.jets import Jet
|
|
if not isinstance(jet, Jet):
|
|
raise TypeError("add_jet expects a Jet")
|
|
if self._initialized:
|
|
raise RuntimeError("Jets must be added before initialize()")
|
|
self.bodies.jets.append(jet)
|
|
return len(self.bodies.jets) - 1
|
|
|
|
@context_bound
|
|
def set_jet_amplitude(self, index: int, amplitude: float) -> None:
|
|
"""Change one jet's signed amplitude; re-uploads per-link velocity only.
|
|
|
|
Invalid input changes nothing. A device error leaves the host amplitude
|
|
unchanged but the device column undefined: treat it like any other CUDA
|
|
exception here and close/recreate the Simulation before stepping again.
|
|
"""
|
|
jet = self.bodies.jets[index].with_amplitude(amplitude)
|
|
if self._initialized:
|
|
self.field.curved.set_jet_amplitude(index, amplitude, self.stream)
|
|
self.bodies.jets[index] = jet # later geometry syncs keep the new amplitude
|
|
|
|
def remove_body(self, id: int) -> None:
|
|
"""Remove a body from the simulation.
|
|
|
|
Before ``initialize()`` the object is removed from the formal registry
|
|
immediately. After ``initialize()`` the removal is staged and will be
|
|
committed at the next ``sync_bodies()``. Removal compacts numeric IDs;
|
|
obtain the remaining IDs again after immediate removal or runtime sync.
|
|
|
|
Args:
|
|
id: body_id from ``add_body()``.
|
|
"""
|
|
if self._initialized:
|
|
self.bodies.stage_remove(id)
|
|
else:
|
|
self.bodies.remove(id)
|
|
|
|
# -- Legacy convenience (keep for backward compat) -----------------------
|
|
def add_cylinder(self, center: Tuple[float, ...],
|
|
radius: float) -> int:
|
|
"""Add a cylinder body (convenience wrapper around add_body)."""
|
|
return self.add_body("circle", center=center, radius=radius)
|
|
|
|
def add_sensor(self, center: Tuple[float, ...],
|
|
radius: float) -> int:
|
|
"""Add a sensor body (convenience wrapper around add_body)."""
|
|
return self.add_body("sensor", center=center, radius=radius)
|
|
|
|
def add_object(self, obj: SimObject) -> int:
|
|
"""Add a pre-configured SimObject directly."""
|
|
return self.bodies.add(obj)
|
|
|
|
def _load_objects_from_body_config(self) -> None:
|
|
"""Instantiate simple objects declared in ``config_body.json``."""
|
|
for spec in self.body_cfg.objects:
|
|
obj_type = str(spec.get("type", "")).strip().lower()
|
|
center = tuple(spec.get("center", ()))
|
|
radius = float(spec.get("radius", 0.0))
|
|
omega = spec.get("omega", None)
|
|
if spec.get("surface_actuation") is not None and obj_type != "cylinder":
|
|
raise ValueError("Surface actuation is only supported on cylinders in body config")
|
|
|
|
if obj_type == "":
|
|
raise ValueError(
|
|
"Body config object is missing required key 'type'.")
|
|
if len(center) not in (2, 3):
|
|
raise ValueError(
|
|
f"Body config object '{obj_type}' requires center "
|
|
"with 2 or 3 entries."
|
|
)
|
|
|
|
if obj_type == "cylinder":
|
|
body_id = self.add_body("circle", center=center, radius=radius,
|
|
surface_actuation=spec.get("surface_actuation"))
|
|
if omega is not None:
|
|
self.bodies.set_body_state(body_id, omega=float(omega))
|
|
elif obj_type == "sensor":
|
|
self.add_body("sensor", center=center, radius=radius)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported body config object type '{obj_type}'. "
|
|
"Supported types: cylinder, sensor."
|
|
)
|
|
|
|
# -- Runtime control -----------------------------------------------------
|
|
def set_body(self, id: int, omega: float = None,
|
|
vx: float = None, vy: float = None) -> None:
|
|
"""Set runtime body state (host action array only, no H2D).
|
|
|
|
The GPU action buffer is updated at the next ``run()`` call.
|
|
|
|
Args:
|
|
id: body_id from add_body().
|
|
omega: rotation speed (lattice units). ``None`` = no change.
|
|
vx: translation velocity x (future). ``None`` = no change.
|
|
vy: translation velocity y (future). ``None`` = no change.
|
|
"""
|
|
obj = self.bodies.get(id)
|
|
proposed = dict(vars(obj.state))
|
|
proposed.update({k: float(v) for k, v in (("omega", omega), ("vx", vx), ("vy", vy))
|
|
if v is not None})
|
|
obj.validate_surface_actuation(self.lbm_cfg, proposed)
|
|
self.bodies.validate_jet_motion(id, proposed)
|
|
changed = False
|
|
if omega is not None:
|
|
obj.state.omega = float(omega)
|
|
changed = True
|
|
if vx is not None:
|
|
obj.state.vx = float(vx)
|
|
changed = True
|
|
if vy is not None:
|
|
obj.state.vy = float(vy)
|
|
changed = True
|
|
if changed:
|
|
self.bodies._refresh_action_from_objects()
|
|
|
|
def set_force(self, id: int, fx: float = 0.0, fy: float = 0.0) -> None:
|
|
"""Set force density on a force_region object (host only, no H2D).
|
|
|
|
The GPU action buffer is updated at the next ``run()`` call.
|
|
|
|
Args:
|
|
id: body_id from add_body(type='force_region', ...).
|
|
fx: x-component of force density (lattice units).
|
|
fy: y-component of force density (lattice units).
|
|
"""
|
|
self.bodies.set_force_state(body_id=id, fx=float(fx), fy=float(fy))
|
|
|
|
# -- Telemetry readback --------------------------------------------------
|
|
def read_force(self, id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return the force vector on body *id* from the pinned obs buffer.
|
|
|
|
Args:
|
|
normalize: If True (default), divide by accumulated step count
|
|
to return the average force per step. If False, return the
|
|
raw GPU-accumulated sum.
|
|
"""
|
|
if self.bodies.obs_pinned is None:
|
|
raise RuntimeError("No obs buffer. Call run() first.")
|
|
return self.bodies.read_force(id, normalize=normalize)
|
|
|
|
def read_torque(self, id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return the torque on body *id* from the pinned obs buffer.
|
|
|
|
Args:
|
|
normalize: If True (default), divide by accumulated step count
|
|
to return the average torque per step. If False, return the
|
|
raw GPU-accumulated sum.
|
|
"""
|
|
if self.bodies.obs_pinned is None:
|
|
raise RuntimeError("No obs buffer. Call run() first.")
|
|
return self.bodies.read_torque(id, normalize=normalize)
|
|
|
|
def read_sensor(self, id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return the sensor reading for body *id* from the pinned obs buffer.
|
|
|
|
Area-normalisation (dividing by the sensor footprint cell count) is
|
|
**always applied internally**. When *normalize* is ``True`` (default),
|
|
the result is also divided by the accumulated step count, yielding a
|
|
**per-step area-averaged velocity**.
|
|
|
|
Args:
|
|
normalize: If True, apply time-normalisation (divide by steps).
|
|
If False, return the area-averaged value without time division.
|
|
"""
|
|
if self.bodies.obs_pinned is None:
|
|
raise RuntimeError("No obs buffer. Call run() first.")
|
|
return self.bodies.read_sensor(id, normalize=normalize)
|
|
|
|
@context_bound
|
|
def read_surface_exchange(self, id: int):
|
|
"""Latest step's wall/aperture (Fx,Fy,Tz) exchange, not total nozzle force.
|
|
|
|
Supports SurfaceActuation and Jet-driven bodies, including zero-amplitude
|
|
apertures. Passive bodies without an aperture are rejected.
|
|
Unlike read_body this is not accumulated/time-averaged. It resets after
|
|
geometry sync or checkpoint load; advance before interpreting it.
|
|
"""
|
|
obj = self.bodies.get(id)
|
|
curved = self.field.curved
|
|
has_aperture = (curved.aperture is not None
|
|
and np.any(curved.aperture & (curved.body_id == id)))
|
|
if obj.surface_actuation is None and not has_aperture:
|
|
raise ValueError("Body has no surface actuation")
|
|
self.stream.synchronize()
|
|
return curved.surface_exchange(id)
|
|
|
|
@context_bound
|
|
def read_body(self, id: int, *,
|
|
stream: cuda.Stream | None = None,
|
|
normalize: bool = True):
|
|
"""Return unified telemetry for one body.
|
|
|
|
Args:
|
|
id: body_id from ``add_body()``.
|
|
stream: Optional CUDA stream to synchronise before reading.
|
|
If ``None``, uses the internal stream.
|
|
normalize: If True (default), all fields are divided by the
|
|
accumulated step count (time-normalisation). Sensor velocity
|
|
is always area-normalised internally regardless of this flag.
|
|
|
|
Returns:
|
|
BodyTelemetry with fields ``force``, ``torque``, ``sensor``.
|
|
``sensor`` is a dim-length array of zeros for non-sensor bodies.
|
|
"""
|
|
if self.bodies.obs_pinned is None:
|
|
raise RuntimeError("No obs buffer. Call run() first.")
|
|
if stream is None:
|
|
stream = self.stream
|
|
if stream is not None:
|
|
stream.synchronize()
|
|
return self.bodies.read_body(id, normalize=normalize)
|
|
|
|
@context_bound
|
|
def read_bodies(self, *, stream: cuda.Stream | None = None) -> np.ndarray:
|
|
"""Return all bodies' telemetry as a flat float32 array.
|
|
|
|
Layout follows the packed obs buffer: ``[force_0..n-1, torque_0..n-1, sensor_0..n-1]``
|
|
where each segment is ``max(n_objects, 1) * dim`` floats.
|
|
Total length equals ``obs_layout(dim, n_objects).total_floats``.
|
|
|
|
This is the most efficient readback for DRL control loops
|
|
(single contiguous D2H copy, no Python per-body overhead).
|
|
"""
|
|
if self.bodies.obs_pinned is None:
|
|
raise RuntimeError("No obs buffer. Call run() first.")
|
|
if stream is None:
|
|
stream = self.stream
|
|
if stream is not None:
|
|
stream.synchronize()
|
|
return self.bodies._obs_array()
|
|
|
|
# -- Compilation ---------------------------------------------------------
|
|
@context_bound
|
|
def _recompile(self, expected_n_objects: int | None = None) -> None:
|
|
"""Re-generate config headers and recompile the kernel.
|
|
|
|
If *expected_n_objects* is ``None``, uses ``self.bodies.count``.
|
|
Preserves ``step_count`` and GPU DDF/flag buffers across the
|
|
module reload. Packed action/observation buffers are managed separately
|
|
by the topology synchronization caller.
|
|
|
|
This is safe to call before or after ``initialize()``.
|
|
"""
|
|
if expected_n_objects is None:
|
|
expected_n_objects = self.bodies.count
|
|
self.lbm_cfg.validate_runtime()
|
|
self._assert_allocation_contract()
|
|
arch = self._resolve_compile_arch()
|
|
next_dir = compiler.generate_config(self.lbm_cfg, n_objects=expected_n_objects)
|
|
try:
|
|
next_ptx = compiler.compile_kernel(arch=arch, build_dir=next_dir)
|
|
next_module = compiler.load_module(next_ptx)
|
|
except Exception:
|
|
compiler.cleanup_build_dir(next_dir)
|
|
raise
|
|
old_dir = self._build_dir
|
|
self._build_dir, self._ptx_path, self._module = next_dir, next_ptx, next_module
|
|
self._compiled_n_objects = int(expected_n_objects)
|
|
|
|
# Reconnect field and stepper to new module
|
|
self.field.module = self._module
|
|
self.field.invalidate_module_constant_cache()
|
|
if self.lbm_cfg.viscosity != self._compiled_viscosity:
|
|
self.field._runtime_params["omega"] = self.lbm_cfg.omega
|
|
self._compiled_viscosity = self.lbm_cfg.viscosity
|
|
self.field._upload_params()
|
|
|
|
_prev_step_count = self.stepper._step_count
|
|
self.stepper = LBMStepper(
|
|
self.field, self._module, self.lbm_cfg,
|
|
)
|
|
self.stepper._step_count = _prev_step_count
|
|
|
|
self._assert_object_count_contract(expected_count=expected_n_objects)
|
|
compiler.cleanup_build_dir(old_dir)
|
|
|
|
def recompile(self):
|
|
"""Re-generate config headers and recompile kernel.
|
|
|
|
Call after changing compile-time parameters (collision model, etc.).
|
|
Must not be called after ``initialize()`` -- use ``sync_bodies()``
|
|
for runtime topology changes instead.
|
|
"""
|
|
if self._initialized:
|
|
raise RuntimeError(
|
|
"recompile() must not be called after initialize(); "
|
|
"use sync_bodies() for runtime body changes.")
|
|
self._recompile()
|
|
|
|
# -- Initialization ------------------------------------------------------
|
|
@context_bound
|
|
def initialize(self):
|
|
"""Initialize flow field and sync objects to GPU."""
|
|
self._assert_allocation_contract()
|
|
# Validate the complete geometry before flags, populations or modules
|
|
# change. Compact-list construction publishes no outputs on failure.
|
|
next_flags = self.bodies.build_flags(self.field.build_channel_flags())
|
|
if self.bodies.jets and self.bodies.count == 0:
|
|
raise ValueError("Jets need bodies with cut links")
|
|
self.bodies.build_compact_lists_for(self.bodies.objects, next_flags)
|
|
|
|
# Recompile if objects were added after construction.
|
|
if self._read_generated_n_objects() != self.bodies.count:
|
|
self.recompile()
|
|
|
|
self.field.flag = next_flags
|
|
self.field.upload_flags()
|
|
|
|
self.stepper.initialize()
|
|
if self.bodies.count > 0:
|
|
self.bodies.sync_to_gpu(self.field, rebuild_flags=False)
|
|
else:
|
|
self.bodies.ensure_packed_buffers()
|
|
self._assert_runtime_contracts()
|
|
self._initialized = True
|
|
|
|
# -- Runtime body topology sync -------------------------------------------
|
|
@context_bound
|
|
def sync_bodies(self) -> None:
|
|
"""Apply pending body edits (add/remove) to a running simulation.
|
|
|
|
Supports both ``double_buffer`` and ``esopull`` streaming modes.
|
|
In EsoPull mode, the backing layout is decoded to physical DDF on
|
|
GPU for the patched cells only, patched on host, then encoded back.
|
|
"""
|
|
if not self._initialized:
|
|
raise RuntimeError("Call initialize() before sync_bodies()")
|
|
|
|
self._assert_allocation_contract()
|
|
# 1. Synchronize GPU
|
|
self.stream.synchronize()
|
|
|
|
# 2. Check for pending edits
|
|
if not self.bodies.has_pending_edit():
|
|
return
|
|
|
|
is_esopull = (self.lbm_cfg.streaming == "esopull")
|
|
|
|
# 3. Build sync plan (same for both modes)
|
|
plan = self.bodies.build_sync_plan(self.field)
|
|
|
|
# 4. Obtain physical DDF on host (mode-dependent path)
|
|
if is_esopull:
|
|
patch_set = _build_patch_set(
|
|
plan.added_solid_mask, plan.released_fluid_mask,
|
|
plan.curved_host[0], nx=self.field.nx, ny=self.field.ny)
|
|
scratch_gpu, scratch_host, idx_gpu = \
|
|
self.field.esopull_sync_decode(
|
|
patch_set, self.stepper.step_count, self.stream)
|
|
_scatter_to_field(self.field, scratch_host, patch_set)
|
|
else:
|
|
self.field.download_ddf(
|
|
step_id=self.stepper.step_count, force=True)
|
|
|
|
old_ddf = self.field.ddf.copy()
|
|
old_flags = self.field.flag.copy()
|
|
|
|
# 5. Runtime recompile for new object count
|
|
self._recompile(plan.next_count)
|
|
|
|
# 6. Apply new topology data (flags, compact lists, params)
|
|
self.bodies.apply_sync_plan(self.field, plan)
|
|
|
|
# 7. DDF patch for geometry changes (works on physical DDF, mode-agnostic)
|
|
cl_fluid_idx = plan.curved_host[0]
|
|
patch_ddf_for_body_sync(
|
|
self.field,
|
|
old_ddf, old_flags, plan.next_flags,
|
|
plan.added_solid_mask, plan.released_fluid_mask,
|
|
curved_fluid_indices=cl_fluid_idx,
|
|
)
|
|
|
|
# 8. Upload patched DDF (mode-dependent path)
|
|
if is_esopull:
|
|
_gather_from_field(self.field, scratch_host, patch_set)
|
|
n = len(patch_set)
|
|
self.field.esopull_sync_encode(
|
|
scratch_gpu, scratch_host, idx_gpu,
|
|
n, self.stepper.step_count, self.stream)
|
|
else:
|
|
upload_patched_ddf(
|
|
self.field,
|
|
plan.added_solid_mask, plan.released_fluid_mask,
|
|
curved_fluid_indices=cl_fluid_idx,
|
|
)
|
|
|
|
# 9. Commit formal objects and clear pending
|
|
self.bodies.commit_pending(plan.next_objects, plan.new_sensor_counts)
|
|
|
|
# 10. Build action/obs telemetry for the new object set
|
|
self.bodies._refresh_action_from_objects()
|
|
self.bodies._allocate_packed_telemetry()
|
|
assert self.bodies.obs_pinned is not None
|
|
self.bodies.obs_pinned.fill(0)
|
|
cuda.memcpy_htod(self.bodies.obs_gpu, self.bodies.obs_pinned)
|
|
|
|
self._assert_runtime_contracts()
|
|
|
|
# -- Stepping ------------------------------------------------------------
|
|
@context_bound
|
|
def run(self, steps: int, *,
|
|
stream: cuda.Stream | None = None,
|
|
upload_act: bool = True,
|
|
sync_obs: bool = True,
|
|
zero_obs: bool = True,
|
|
checkpoint_interval: int = 0):
|
|
"""Advance simulation by *steps* time steps.
|
|
|
|
Args:
|
|
steps: Nonnegative integer number of LBM steps.
|
|
stream: CUDA stream for async operations. ``None`` = internal stream.
|
|
External streams must be created under ``with sim.ctx.activate():``.
|
|
upload_act: If True, upload host action array to ``action_gpu``
|
|
before the step group.
|
|
sync_obs: If True, download ``obs_gpu`` to host pinned buffer
|
|
after the step group.
|
|
zero_obs: If True (default), zero all obs segments (force, torque,
|
|
sensor) on GPU and reset the step accumulator before stepping.
|
|
checkpoint_interval: Nonnegative integer; if >0, save every N steps.
|
|
"""
|
|
if not self._initialized:
|
|
raise RuntimeError("Call initialize() first")
|
|
steps = operator.index(steps)
|
|
checkpoint_interval = operator.index(checkpoint_interval)
|
|
if steps < 0 or checkpoint_interval < 0:
|
|
raise ValueError("steps and checkpoint_interval must be nonnegative")
|
|
self._assert_runtime_contracts()
|
|
# Discard any uncommitted body edits before stepping.
|
|
if self.bodies.has_pending_edit():
|
|
self.bodies.clear_pending_edits()
|
|
|
|
# Resolve stream
|
|
if stream is None:
|
|
stream = self.stream
|
|
|
|
# Async upload action
|
|
if upload_act and self.bodies.count > 0:
|
|
self.bodies._upload_action_async(stream)
|
|
|
|
# Zero obs segments before step group
|
|
if zero_obs:
|
|
self.bodies.zero_obs_async(stream)
|
|
|
|
if checkpoint_interval > 0:
|
|
done = 0
|
|
while done < steps:
|
|
batch = min(checkpoint_interval, steps - done)
|
|
self.stepper.step(
|
|
batch,
|
|
action_gpu=self.bodies.action_gpu,
|
|
obs_gpu=self.bodies.obs_gpu,
|
|
stream=stream,
|
|
)
|
|
self.bodies.increment_obs_steps(batch)
|
|
done += batch
|
|
if done < steps or done % checkpoint_interval == 0:
|
|
stream.synchronize()
|
|
self.save_checkpoint()
|
|
else:
|
|
self.stepper.step(
|
|
steps,
|
|
action_gpu=self.bodies.action_gpu,
|
|
obs_gpu=self.bodies.obs_gpu,
|
|
stream=stream,
|
|
)
|
|
|
|
# Count each batch before any checkpoint captures observations.
|
|
if steps > 0:
|
|
self.bodies.increment_obs_steps(steps)
|
|
|
|
# Async download obs
|
|
if sync_obs:
|
|
self.bodies.download_obs_full_async(stream)
|
|
|
|
stream.synchronize()
|
|
|
|
def step(self, n: int = 1):
|
|
"""Advance *n* steps (convenience for interactive use)."""
|
|
self.run(n)
|
|
|
|
# -- Data access ---------------------------------------------------------
|
|
@context_bound
|
|
def get_macroscopic(self) -> Dict[str, np.ndarray]:
|
|
"""Download DDF and return rho, ux, uy [, uz]."""
|
|
return self.field.get_macroscopic(step_id=self.stepper.step_count)
|
|
|
|
@context_bound
|
|
def get_ddf(self) -> np.ndarray:
|
|
self.field.download_ddf(step_id=self.stepper.step_count)
|
|
return self.field.ddf.copy()
|
|
|
|
def get_flags(self) -> np.ndarray:
|
|
return self.field.flag.copy()
|
|
|
|
# -- Runtime parameter updates -------------------------------------------
|
|
@context_bound
|
|
def update_runtime_params(self, **kwargs):
|
|
"""Update __constant__ d_params without recompiling.
|
|
|
|
Accepted: omega, fx, fy, fz, n_objects.
|
|
"""
|
|
self.field.update_params(**kwargs)
|
|
|
|
# -- Snapshots -----------------------------------------------------------
|
|
@context_bound
|
|
def snapshot(self):
|
|
"""Snapshot the flow field; restore keeps the current clock and controls."""
|
|
self.field.snapshot(step_id=self.stepper.step_count)
|
|
|
|
@context_bound
|
|
def restore(self):
|
|
"""Restore the saved flow field at the current streaming parity."""
|
|
self.field.restore(step_id=self.stepper.step_count)
|
|
|
|
# -- Checkpoint / Restore (HDF5) ----------------------------------------
|
|
@context_bound
|
|
def save_checkpoint(self, path: str = None):
|
|
"""Save full simulation state to an HDF5 file."""
|
|
from .common.checkpoint import save_checkpoint
|
|
if self.bodies.jets:
|
|
raise NotImplementedError("Jets are not checkpointed yet")
|
|
return save_checkpoint(self.field, self.stepper, self.lbm_cfg,
|
|
self.bodies, path)
|
|
|
|
@context_bound
|
|
def load_checkpoint(self, path: str):
|
|
"""Restore simulation state from an HDF5 checkpoint."""
|
|
from .common.checkpoint import load_checkpoint
|
|
if not self._initialized:
|
|
self.initialize()
|
|
load_checkpoint(path, self.field, self.stepper, self.lbm_cfg,
|
|
self.bodies)
|
|
self._initialized = True
|
|
|
|
# -- Cleanup -------------------------------------------------------------
|
|
def close(self):
|
|
if getattr(self, "_closed", False):
|
|
return
|
|
ctx = getattr(self, "ctx", None)
|
|
if ctx is not None and ctx.is_active:
|
|
raise RuntimeError("Cannot close Simulation during an active CUDA operation")
|
|
if ctx is not None and not ctx.is_closed:
|
|
with ctx.activate():
|
|
cuda.Context.synchronize()
|
|
# Release owned GPU resources while their context is current.
|
|
self.stepper = None
|
|
self.bodies = None
|
|
self.field = None
|
|
self._module = None
|
|
self._stream = None
|
|
ctx.close()
|
|
compiler.cleanup_build_dir(getattr(self, "_build_dir", None))
|
|
self._build_dir = None
|
|
self._closed = True
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def _resolve_compile_arch(self) -> str:
|
|
"""Return nvcc arch, respecting explicit compute_capability override."""
|
|
if self.lbm_cfg.compute_capability == "auto":
|
|
return self.ctx.sm_arch
|
|
return f"sm_{''.join(self.lbm_cfg.compute_capability.split('.'))}"
|
|
|
|
def _read_generated_n_objects(self) -> int:
|
|
"""Return cached N_OBJS from the last compile (avoids header races)."""
|
|
if hasattr(self, "_compiled_n_objects"):
|
|
return int(self._compiled_n_objects)
|
|
cfg_path = compiler.kernel_path("config/config_objects.h")
|
|
with open(cfg_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line.startswith("#define N_OBJS"):
|
|
return int(line.split()[-1])
|
|
raise RuntimeError(
|
|
f"N_OBJS not found in generated header: {cfg_path}")
|
|
|
|
def _assert_object_count_contract(
|
|
self, expected_count: int | None = None
|
|
) -> None:
|
|
"""Ensure generated compile-time N_OBJS matches runtime object count."""
|
|
runtime_count = (
|
|
self.bodies.count if expected_count is None
|
|
else int(expected_count)
|
|
)
|
|
compiled_count = self._read_generated_n_objects()
|
|
if compiled_count != runtime_count:
|
|
raise RuntimeError(
|
|
"Object count contract mismatch: generated N_OBJS="
|
|
f"{compiled_count}, runtime count={runtime_count}. "
|
|
"Recompile required."
|
|
)
|
|
|
|
def _assert_obs_layout_contract(self) -> None:
|
|
"""Ensure host obs layout matches runtime body count and contract."""
|
|
lay = compiler.obs_layout(self.lbm_cfg.dim, self.bodies.count)
|
|
mgr = self.bodies
|
|
if mgr.obs_n_slots != lay.n_slots:
|
|
raise RuntimeError(
|
|
"obs layout mismatch: n_slots="
|
|
f"{mgr.obs_n_slots}, expected={lay.n_slots}."
|
|
)
|
|
if (
|
|
mgr.slot_stride_floats != lay.slot_stride_floats
|
|
or mgr.torque0_floats != lay.torque0_floats
|
|
or mgr.sensor0_floats != lay.sensor0_floats
|
|
or mgr.obs_total_floats != lay.total_floats
|
|
):
|
|
raise RuntimeError(
|
|
"obs layout mismatch between runtime buffers and expected "
|
|
f"layout for dim={self.lbm_cfg.dim}, "
|
|
f"n_objects={self.bodies.count}."
|
|
)
|
|
|
|
def _assert_runtime_contracts(self) -> None:
|
|
"""Validate object-count and obs layout contracts before stepping."""
|
|
self._assert_allocation_contract()
|
|
self._assert_object_count_contract()
|
|
self._assert_obs_layout_contract()
|
|
|
|
def _assert_allocation_contract(self) -> None:
|
|
if self.lbm_cfg.allocation_signature() != self._allocation_signature:
|
|
raise ValueError(
|
|
"Grid, lattice, streaming or storage configuration changed after "
|
|
"allocation. Construct a new Simulation for structural changes.")
|
|
|
|
|
|
# -- Module-level helpers -------------------------------------------------------
|
|
|
|
def _build_patch_set(added_solid_mask, released_fluid_mask, curved_fluid_idx,
|
|
nx=None, ny=None):
|
|
"""Return sorted uint32 array of cell indices that need DDF patching.
|
|
|
|
Includes a 1-cell halo for the old-fluid donors used by released-cell
|
|
reconstruction. EsoPull same-time encoding itself has disjoint ownership
|
|
and does not require a halo to avoid storage collisions.
|
|
"""
|
|
s = set()
|
|
s.update(np.where(added_solid_mask)[0].tolist())
|
|
s.update(np.where(released_fluid_mask)[0].tolist())
|
|
if curved_fluid_idx is not None:
|
|
s.update(curved_fluid_idx.tolist())
|
|
|
|
if nx is None or ny is None:
|
|
return np.array(sorted(s), dtype=np.uint32)
|
|
|
|
# Expand to include 1-cell halo (D2Q9 Moore neighbourhood)
|
|
expanded = set(s)
|
|
for idx in s:
|
|
x = idx % nx
|
|
y = idx // nx
|
|
for dx in (-1, 0, 1):
|
|
for dy in (-1, 0, 1):
|
|
if dx == 0 and dy == 0:
|
|
continue
|
|
xn, yn = x + dx, y + dy
|
|
if 0 <= xn < nx and 0 <= yn < ny:
|
|
expanded.add(xn + yn * nx)
|
|
return np.array(sorted(expanded), dtype=np.uint32)
|
|
|
|
|
|
def _scatter_to_field(field, scratch_host, patch_indices):
|
|
"""Copy decoded physical DDF from scratch buffer into field.ddf."""
|
|
nq = field.nq
|
|
populations = field.ddf.reshape(nq, -1)
|
|
for j, idx in enumerate(patch_indices):
|
|
populations[:, idx] = scratch_host[j * nq:(j + 1) * nq]
|
|
|
|
|
|
def _gather_from_field(field, scratch_host, patch_indices):
|
|
"""Copy patched physical DDF from field.ddf back to scratch buffer."""
|
|
nq = field.nq
|
|
populations = field.ddf.reshape(nq, -1)
|
|
for j, idx in enumerate(patch_indices):
|
|
scratch_host[j * nq:(j + 1) * nq] = populations[:, idx]
|