feat(body): runtime body add/remove, unified action/obs, FRC_REGION flag
- Add runtime body topology sync (add_body/remove_body + sync_bodies) with recompile, DDF patch (feq + BFS inward fill), and commit. - Unify action/obs flow: set_body/set_force are now host-only; run() auto-uploads action and downloads obs via CUDA stream. - Add read_body(id) -> BodyTelemetry and read_bodies() for DRL loops. - Add FRC_REGION flag (0x0800) for force_region cells. - Extract equilibrium helpers (lbm/equilibrium.py) and DDF patch module (body/ddf_patch.py). - Merge recompile / _runtime_recompile into single _recompile(). - Add n_objects to checkpoint; validate on load. - Add test suite: 40 unit + 19 integration tests (59 total). - Add conftest.py and docs/tests_overview.md for test documentation. - Update README.md and CONFIG.md for new API. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
1 parent
d5b7e98750
commit
987566c0e6
28 files changed
+2112
-86
No files matched your search
@@ -5,6 +5,7 @@ Body / object management for immersed and rigid objects.
|
||||
from .objects import SimObject, ObjectState, ObjectControl
|
||||
from .manager import ObjectManager
|
||||
from .registry import BodyRegistry
|
||||
from .sync_plan import BodySyncPlan
|
||||
from .geometry.base import CutLink, SensorCell, Geometry
|
||||
from .geometry.circle import CircleGeometry
|
||||
from .coupling.soa_packer import pack_cut_links_to_soa, pack_sensor_to_soa
|
||||
@@ -13,6 +14,7 @@ __all__ = [
|
||||
"SimObject", "ObjectState", "ObjectControl",
|
||||
"ObjectManager",
|
||||
"BodyRegistry",
|
||||
"BodySyncPlan",
|
||||
"CutLink", "SensorCell", "Geometry",
|
||||
"CircleGeometry",
|
||||
"pack_cut_links_to_soa", "pack_sensor_to_soa",
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
# CelerisLab/body/ddf_patch.py
|
||||
"""
|
||||
Host-side DDF patching for runtime body topology sync. (D2Q9 only)
|
||||
|
||||
When bodies are added or removed at runtime, the host DDF must be
|
||||
patched so that newly-solid cells contain equilibrium (not stale
|
||||
fluid values) and newly-fluid cells initialise from their neighbours.
|
||||
|
||||
Called by ``Simulation.sync_bodies()`` after topology rebuild and
|
||||
before GPU upload.
|
||||
|
||||
All equilibrium calculations use ``lbm/equilibrium.py`` which is D2Q9-only.
|
||||
D3Q19 DDF patching is not yet implemented on the host side.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import pycuda.driver as cuda
|
||||
|
||||
from ..lbm.descriptors import FLUID
|
||||
from ..lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf
|
||||
|
||||
|
||||
def patch_ddf_for_body_sync(
|
||||
field,
|
||||
old_ddf: np.ndarray,
|
||||
old_flags: np.ndarray,
|
||||
new_flags: np.ndarray,
|
||||
added_solid_mask: np.ndarray,
|
||||
released_fluid_mask: np.ndarray,
|
||||
curved_fluid_indices: np.ndarray | None = None,
|
||||
) -> None:
|
||||
"""Patch host-side DDF after a body topology change.
|
||||
|
||||
Operates on physical DDF (``field.ddf``). After patching, the caller
|
||||
should call ``upload_patched_ddf()`` to push only the changed cells
|
||||
to both GPU buffers.
|
||||
|
||||
Three patch types:
|
||||
|
||||
* ``old_fluid -> new_solid``: Write equilibrium at rest (rho=RHO, u=0)
|
||||
so stale fluid values do not linger inside new solid cells.
|
||||
|
||||
* ``old_solid -> new_fluid``: BFS inward fill from surrounding old-fluid
|
||||
cells. Layer-0 cells average their old-fluid neighbours; deeper layers
|
||||
average the previously filled layer. Fallback: equilibrium at inlet
|
||||
velocity.
|
||||
|
||||
* **Boundary fluid cells** (new curved boundary neighbours): If new
|
||||
solid cells appear, the fluid cells just outside them (on curved
|
||||
boundary links) are also set to equilibrium at their local velocity
|
||||
(or rest if NaN), preventing abrupt-boundary instability.
|
||||
|
||||
Args:
|
||||
field: LBMField instance (``field.ddf`` is modified in-place).
|
||||
old_ddf: Physical DDF snapshot before topology change (n*NQ, float32).
|
||||
old_flags: Flag array before topology change.
|
||||
new_flags: Flag array after topology change.
|
||||
added_solid_mask: Bool array -- True where fluid became solid.
|
||||
released_fluid_mask: Bool array -- True where obstacle became fluid.
|
||||
curved_fluid_indices: Optional array of fluid cell indices on curved
|
||||
boundary links. When *added_solid_mask* has any True entries,
|
||||
these cells are set to equilibrium at local velocity (or rest).
|
||||
"""
|
||||
_patch_added_solid(field, added_solid_mask)
|
||||
_patch_boundary_fluid(field, curved_fluid_indices, added_solid_mask)
|
||||
_patch_released_fluid(field, old_ddf, old_flags, released_fluid_mask)
|
||||
|
||||
|
||||
def upload_patched_ddf(
|
||||
field,
|
||||
added_solid_mask: np.ndarray,
|
||||
released_fluid_mask: np.ndarray,
|
||||
curved_fluid_indices: np.ndarray | None = None,
|
||||
) -> None:
|
||||
"""Upload only the DDF cells that changed during patching.
|
||||
|
||||
Uses per-cell ``cuda.memcpy_htod`` for both ``ddf_gpu`` and
|
||||
``temp_gpu``. Leaves all other cells (especially boundary cells
|
||||
with NaN) untouched on the GPU, preserving their streaming-
|
||||
compatible values from before the sync.
|
||||
"""
|
||||
changed_indices: set[int] = set()
|
||||
changed_indices.update(np.where(added_solid_mask)[0].tolist())
|
||||
changed_indices.update(np.where(released_fluid_mask)[0].tolist())
|
||||
if curved_fluid_indices is not None:
|
||||
changed_indices.update(curved_fluid_indices.tolist())
|
||||
if not changed_indices:
|
||||
field.invalidate_host_ddf_cache()
|
||||
field._host_ddf_step = None
|
||||
return
|
||||
|
||||
nq = field.nq
|
||||
cell_bytes = nq * 4
|
||||
ddf_ptr = int(field.ddf_gpu)
|
||||
temp_ptr = int(field.temp_gpu)
|
||||
|
||||
for idx in sorted(changed_indices):
|
||||
offset = idx * cell_bytes
|
||||
cell_data = field.ddf[idx * nq:(idx + 1) * nq].tobytes()
|
||||
cuda.memcpy_htod(ddf_ptr + offset, cell_data)
|
||||
cuda.memcpy_htod(temp_ptr + offset, cell_data)
|
||||
|
||||
field.invalidate_host_ddf_cache()
|
||||
field._host_ddf_step = None
|
||||
|
||||
|
||||
# -- Internal helpers ---------------------------------------------------------
|
||||
|
||||
def _patch_added_solid(field, added_solid_mask: np.ndarray) -> None:
|
||||
"""Write resting equilibrium for cells that became solid."""
|
||||
if not np.any(added_solid_mask):
|
||||
return
|
||||
rho0 = field.cfg.rho
|
||||
feq = compute_feq_d2q9(rho0, 0.0, 0.0)
|
||||
indices = np.where(added_solid_mask)[0]
|
||||
nq = field.nq
|
||||
for idx in indices:
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq
|
||||
|
||||
|
||||
def _patch_boundary_fluid(
|
||||
field,
|
||||
curved_fluid_indices: np.ndarray | None,
|
||||
added_solid_mask: np.ndarray,
|
||||
) -> None:
|
||||
"""Patch fluid cells on new curved boundary links to local equilibrium."""
|
||||
if curved_fluid_indices is None or len(curved_fluid_indices) == 0:
|
||||
return
|
||||
if not np.any(added_solid_mask):
|
||||
return
|
||||
nq = field.nq
|
||||
feq_fallback = compute_feq_d2q9(field.cfg.rho, 0.0, 0.0)
|
||||
for idx in curved_fluid_indices:
|
||||
f = field.ddf[idx * nq:(idx + 1) * nq]
|
||||
if np.any(np.isnan(f)):
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq_fallback
|
||||
continue
|
||||
rho, ux, uy = compute_macro_from_ddf(f)
|
||||
feq = compute_feq_d2q9(float(rho), float(ux), float(uy))
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq
|
||||
|
||||
|
||||
def _patch_released_fluid(
|
||||
field,
|
||||
old_ddf: np.ndarray,
|
||||
old_flags: np.ndarray,
|
||||
released_fluid_mask: np.ndarray,
|
||||
) -> None:
|
||||
"""BFS inward fill for cells that became fluid after body removal."""
|
||||
if not np.any(released_fluid_mask):
|
||||
return
|
||||
|
||||
nx, ny = field.nx, field.ny
|
||||
nq = field.nq
|
||||
old_is_fluid = (old_flags & FLUID) != 0
|
||||
released = np.where(released_fluid_mask)[0]
|
||||
|
||||
filled = np.zeros(field.n, dtype=bool)
|
||||
|
||||
offsets = [
|
||||
(-1, -1), (-1, 0), (-1, 1),
|
||||
(0, -1), (0, 1),
|
||||
(1, -1), (1, 0), (1, 1),
|
||||
]
|
||||
|
||||
# Precompute macroscopic from old_ddf for old-fluid cells
|
||||
# Skip cells with NaN (pre-existing body curvature may have
|
||||
# contaminated boundary-fluid DDF values).
|
||||
old_macro: dict[int, tuple[float, float, float]] = {}
|
||||
fluid_indices = np.where(old_is_fluid)[0]
|
||||
for idx in fluid_indices:
|
||||
f = old_ddf[idx * nq:(idx + 1) * nq]
|
||||
if np.any(np.isnan(f)):
|
||||
continue
|
||||
rho, ux, uy = compute_macro_from_ddf(f)
|
||||
old_macro[idx] = (rho, ux, uy)
|
||||
|
||||
def _neighbours(idx: int):
|
||||
x = idx % nx
|
||||
y = idx // nx
|
||||
for dx, dy in offsets:
|
||||
xn, yn = x + dx, y + dy
|
||||
if 0 <= xn < nx and 0 <= yn < ny:
|
||||
yield xn + yn * nx
|
||||
|
||||
prev_layer_filled: dict[int, tuple[float, float, float]] = {}
|
||||
remaining = set(released.tolist())
|
||||
|
||||
for idx in released:
|
||||
nbrs = list(_neighbours(idx))
|
||||
fluid_nbrs = [n for n in nbrs if n in old_macro]
|
||||
if not fluid_nbrs:
|
||||
continue
|
||||
rho_avg = np.mean([old_macro[n][0] for n in fluid_nbrs])
|
||||
ux_avg = np.mean([old_macro[n][1] for n in fluid_nbrs])
|
||||
uy_avg = np.mean([old_macro[n][2] for n in fluid_nbrs])
|
||||
feq = compute_feq_d2q9(float(rho_avg), float(ux_avg), float(uy_avg))
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq
|
||||
filled[idx] = True
|
||||
prev_layer_filled[idx] = (
|
||||
float(rho_avg), float(ux_avg), float(uy_avg))
|
||||
remaining.discard(idx)
|
||||
|
||||
# BFS layers
|
||||
max_layers = max(nx, ny)
|
||||
for _ in range(max_layers):
|
||||
if not remaining:
|
||||
break
|
||||
current_layer: dict[int, tuple[float, float, float]] = {}
|
||||
for idx in list(remaining):
|
||||
nbrs = list(_neighbours(idx))
|
||||
filled_nbrs = [n for n in nbrs if n in prev_layer_filled]
|
||||
if not filled_nbrs:
|
||||
continue
|
||||
rho_avg = np.mean([prev_layer_filled[n][0] for n in filled_nbrs])
|
||||
ux_avg = np.mean([prev_layer_filled[n][1] for n in filled_nbrs])
|
||||
uy_avg = np.mean([prev_layer_filled[n][2] for n in filled_nbrs])
|
||||
feq = compute_feq_d2q9(
|
||||
float(rho_avg), float(ux_avg), float(uy_avg))
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq
|
||||
filled[idx] = True
|
||||
current_layer[idx] = (
|
||||
float(rho_avg), float(ux_avg), float(uy_avg))
|
||||
remaining.discard(idx)
|
||||
prev_layer_filled = current_layer
|
||||
|
||||
# Fallback for any remaining unfilled cells
|
||||
u_inlet = field.cfg.velocity
|
||||
for idx in remaining:
|
||||
y = idx // nx
|
||||
ux_fallback = (
|
||||
u_inlet * 4.0 * y * (ny - 1 - y) / ((ny - 1) ** 2)
|
||||
if ny > 2 else u_inlet)
|
||||
feq = compute_feq_d2q9(field.cfg.rho, float(ux_fallback), 0.0)
|
||||
field.ddf[idx * nq:(idx + 1) * nq] = feq
|
||||
@@ -135,3 +135,19 @@ class Geometry(ABC):
|
||||
uint16 array of shape (nx*ny,) with flag bits set for sensor cells.
|
||||
"""
|
||||
...
|
||||
|
||||
def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray:
|
||||
"""Return (nx*ny,) uint16 flag mask for force-region cells.
|
||||
|
||||
This variant marks force-region cells (``FLUID|FRC_REGION``).
|
||||
Default implementation calls ``build_sensor_flag_mask`` for
|
||||
subclasses that treat sensor and force-region identically.
|
||||
|
||||
Args:
|
||||
nx, ny: Grid dimensions.
|
||||
|
||||
Returns:
|
||||
uint16 array of shape (nx*ny,) with flag bits set for
|
||||
force-region cells.
|
||||
"""
|
||||
return self.build_sensor_flag_mask(nx, ny)
|
||||
@@ -14,7 +14,7 @@ import numpy as np
|
||||
|
||||
from .base import CutLink, Geometry, SensorCell
|
||||
from ...lbm.descriptors import (
|
||||
FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG,
|
||||
FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION,
|
||||
D2Q9_EX, D2Q9_EY,
|
||||
)
|
||||
|
||||
@@ -175,6 +175,10 @@ class CircleGeometry(Geometry):
|
||||
"""Return (nx*ny,) uint16 flag mask with FLUID|SENSOR_FLAG."""
|
||||
return self._build_flag_mask_inner(nx, ny, FLUID | SENSOR_FLAG)
|
||||
|
||||
def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray:
|
||||
"""Return (nx*ny,) uint16 flag mask with FLUID|FRC_REGION."""
|
||||
return self._build_flag_mask_inner(nx, ny, FLUID | FRC_REGION)
|
||||
|
||||
def _build_flag_mask_inner(self, nx: int, ny: int, flag_val: int) -> np.ndarray:
|
||||
n = nx * ny
|
||||
mask = np.zeros(n, dtype=np.uint16)
|
||||
|
||||
+268
-21
@@ -4,6 +4,7 @@ ObjectManager -- batch management of SimObjects.
|
||||
|
||||
Responsibilities:
|
||||
- Add / remove / query objects (delegated to ``BodyRegistry``)
|
||||
- Pending edit state for runtime body topology changes
|
||||
- Build merged flag masks and compact cut-link / sensor lists
|
||||
- Allocate packed telemetry ``obs_gpu`` + pagelocked mirror ``obs_pinned``
|
||||
- Sync geometry to :class:`~CelerisLab.lbm.field.LBMField`
|
||||
@@ -11,13 +12,18 @@ Responsibilities:
|
||||
Design::
|
||||
Packed ``obs`` layout matches ``generate_config(..., n_objects=count)`` and
|
||||
``config_obs.h`` macros. Host sizing uses :func:`~CelerisLab.cuda.compiler_v2.obs_layout`.
|
||||
|
||||
Pending edits are lightweight: only the "intent" (add / remove) is stored.
|
||||
All derived quantities (flags, compact lists, GPU buffers) are built at
|
||||
``sync_bodies()`` time from the formal registry + pending edits.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, NamedTuple, Optional, TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import pycuda.driver as cuda
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from .objects import SimObject
|
||||
from .registry import BodyRegistry
|
||||
@@ -32,6 +38,20 @@ from ..lbm.descriptors import (
|
||||
D3Q19_EZ,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .sync_plan import BodySyncPlan
|
||||
|
||||
|
||||
class BodyTelemetry(NamedTuple):
|
||||
"""Unified telemetry for one body, read from the pinned obs buffer.
|
||||
|
||||
After ``run(sync_obs=True)``, all fields are populated from GPU results.
|
||||
``sensor`` is a dim-length array — zeros for non-sensor bodies.
|
||||
"""
|
||||
force: np.ndarray
|
||||
torque: np.ndarray
|
||||
sensor: np.ndarray
|
||||
|
||||
|
||||
class ObjectManager:
|
||||
"""Central registry for all simulation objects.
|
||||
@@ -70,6 +90,12 @@ class ObjectManager:
|
||||
self.obs_nbytes: int = 0
|
||||
|
||||
self._telemetry_field: Optional[object] = None
|
||||
|
||||
# -- Pending edit state (runtime body topology sync) -------------------
|
||||
self._edit_active: bool = False
|
||||
self._pending_add: List[SimObject] = []
|
||||
self._pending_remove: set[int] = set()
|
||||
|
||||
# Ensure host action buffer is non-empty so CUDA alloc never sees 0 bytes
|
||||
# when bodies.count == 0 (matches _resize_buffers sizing for n=0).
|
||||
self._resize_buffers()
|
||||
@@ -113,24 +139,217 @@ class ObjectManager:
|
||||
new_sensor_counts[:copy_n] = self.sensor_cell_counts[:copy_n]
|
||||
self.sensor_cell_counts = new_sensor_counts
|
||||
|
||||
# -- Pending edit state (runtime body topology sync) ---------------------
|
||||
def stage_add(self, obj: SimObject) -> None:
|
||||
"""Stage an object for addition at the next ``sync_bodies()``.
|
||||
|
||||
The object is not added to the formal registry until ``commit_pending()``
|
||||
is called. This implicitly activates the edit window.
|
||||
"""
|
||||
self._edit_active = True
|
||||
self._pending_add.append(obj)
|
||||
|
||||
def stage_remove(self, obj_id: int) -> None:
|
||||
"""Stage removal of a formal object at the next ``sync_bodies()``.
|
||||
|
||||
Args:
|
||||
obj_id: Id of a currently registered formal object.
|
||||
|
||||
Raises:
|
||||
IndexError: If *obj_id* is not in the formal registry.
|
||||
"""
|
||||
self._validate_body_id(obj_id)
|
||||
self._edit_active = True
|
||||
self._pending_remove.add(obj_id)
|
||||
|
||||
def has_pending_edit(self) -> bool:
|
||||
"""Return True if there are uncommitted staged edits."""
|
||||
return self._edit_active and bool(self._pending_add or self._pending_remove)
|
||||
|
||||
def clear_pending_edits(self) -> None:
|
||||
"""Discard all staged edits without applying them."""
|
||||
self._pending_add.clear()
|
||||
self._pending_remove.clear()
|
||||
self._edit_active = False
|
||||
|
||||
# -- Sync plan construction (runtime body topology sync) -----------------
|
||||
def build_next_objects(self) -> List[SimObject]:
|
||||
"""Merge formal objects with pending edits into a new object list.
|
||||
|
||||
Formal objects whose id is in ``_pending_remove`` are excluded.
|
||||
Objects in ``_pending_add`` are appended. All returned objects
|
||||
receive consecutive ``obj_id = 0..n-1``.
|
||||
|
||||
Returns:
|
||||
New list of SimObject copies with reassigned ids.
|
||||
"""
|
||||
def _copy(obj, oid):
|
||||
new = SimObject(obj_id=oid, geometry=obj.geometry,
|
||||
center=obj.center, radius=obj.radius,
|
||||
is_sensor=obj.is_sensor,
|
||||
is_force_region=obj.is_force_region)
|
||||
new.state = obj.state
|
||||
new.control = obj.control
|
||||
return new
|
||||
|
||||
result: List[SimObject] = []
|
||||
next_id = 0
|
||||
for obj in self._registry.objects:
|
||||
if obj.obj_id in self._pending_remove:
|
||||
continue
|
||||
result.append(_copy(obj, next_id))
|
||||
next_id += 1
|
||||
for obj in self._pending_add:
|
||||
result.append(_copy(obj, next_id))
|
||||
next_id += 1
|
||||
return result
|
||||
|
||||
def build_sync_plan(self, field) -> 'BodySyncPlan':
|
||||
"""Construct a full topology transition plan from pending edits.
|
||||
|
||||
Args:
|
||||
field: ``LBMField`` instance (used for old flags and channel base).
|
||||
|
||||
Returns:
|
||||
:class:`BodySyncPlan` describing the transition.
|
||||
"""
|
||||
from .sync_plan import BodySyncPlan
|
||||
from ..lbm.descriptors import FLUID, OBSTACLE
|
||||
|
||||
old_flags = field.flag.copy()
|
||||
|
||||
next_objects = self.build_next_objects()
|
||||
next_count = len(next_objects)
|
||||
|
||||
# Build new flags from clean channel base + new object set
|
||||
base_flags = field.build_channel_flags()
|
||||
next_flags = self.build_flags_for(
|
||||
next_objects, base_flags,
|
||||
nx=self.nx, ny=self.ny, nz=self.nz,
|
||||
)
|
||||
|
||||
# Build compact lists for new object set
|
||||
new_sensor_counts = np.zeros(max(next_count, 1), dtype=np.int32)
|
||||
(
|
||||
cl_fluid_idx, cl_dir, cl_q, cl_body_id,
|
||||
cl_rx, cl_ry, cl_rz, cl_fallback_class,
|
||||
sensor_cells, sensor_obj_id,
|
||||
fr_cells, fr_obj_id,
|
||||
) = self.build_compact_lists_for(
|
||||
next_objects, domain_flags=next_flags,
|
||||
sensor_cell_counts_out=new_sensor_counts,
|
||||
)
|
||||
|
||||
# Compute change masks
|
||||
old_is_solid = (old_flags & 0x0002) != 0
|
||||
new_is_solid = (next_flags & 0x0002) != 0
|
||||
old_is_obstacle = (old_flags & OBSTACLE) != 0
|
||||
new_is_fluid = (next_flags & FLUID) != 0
|
||||
|
||||
added_solid_mask = (~old_is_solid) & new_is_solid
|
||||
released_fluid_mask = old_is_obstacle & new_is_fluid
|
||||
|
||||
return BodySyncPlan(
|
||||
next_objects=next_objects,
|
||||
next_count=next_count,
|
||||
next_flags=next_flags,
|
||||
curved_host=(cl_fluid_idx, cl_dir, cl_q, cl_body_id,
|
||||
cl_rx, cl_ry, cl_rz, cl_fallback_class),
|
||||
sensor_host=(sensor_cells, sensor_obj_id),
|
||||
force_region_host=(fr_cells, fr_obj_id),
|
||||
added_solid_mask=added_solid_mask,
|
||||
released_fluid_mask=released_fluid_mask,
|
||||
new_sensor_counts=new_sensor_counts,
|
||||
)
|
||||
|
||||
def commit_pending(self, next_objects: List[SimObject],
|
||||
sensor_counts: np.ndarray) -> None:
|
||||
"""Replace the formal registry with *next_objects* and clear pending edits.
|
||||
|
||||
Args:
|
||||
next_objects: New formal object list with ``obj_id = 0..n-1``.
|
||||
sensor_counts: New per-object sensor cell counts.
|
||||
"""
|
||||
self._registry.clear()
|
||||
for obj in next_objects:
|
||||
self._registry.add(obj)
|
||||
self.sensor_cell_counts = sensor_counts.copy()
|
||||
self._resize_buffers()
|
||||
self.clear_pending_edits()
|
||||
|
||||
def apply_sync_plan(self, field, plan) -> None:
|
||||
"""Apply topology data from a ``BodySyncPlan`` to the field.
|
||||
|
||||
Sets ``field.flag``, assigns compact list host arrays, uploads
|
||||
compact lists, and refreshes action buffer from object state.
|
||||
|
||||
Args:
|
||||
field: LBMField instance.
|
||||
plan: BodySyncPlan from ``build_sync_plan()``.
|
||||
"""
|
||||
field.flag = plan.next_flags
|
||||
field.upload_flags()
|
||||
|
||||
cl_host = plan.curved_host
|
||||
field.curved.assign_host(*cl_host)
|
||||
s_cells, s_ids = plan.sensor_host
|
||||
field.sensors.assign_host(s_cells, s_ids)
|
||||
fr_cells, fr_ids = plan.force_region_host
|
||||
field.force_regions.assign_host(fr_cells, fr_ids)
|
||||
|
||||
field.upload_compact_lists()
|
||||
field.update_params(n_objects=plan.next_count)
|
||||
|
||||
# -- Flag composition ----------------------------------------------------
|
||||
def build_flags(self, base_flags: np.ndarray) -> np.ndarray:
|
||||
"""Merge formal object flag masks onto a clean domain base."""
|
||||
return self.build_flags_for(
|
||||
self._registry.objects, base_flags,
|
||||
nx=self.nx, ny=self.ny, nz=self.nz,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def build_flags_for(objects: List[SimObject],
|
||||
base_flags: np.ndarray,
|
||||
*, nx: int, ny: int, nz: int) -> np.ndarray:
|
||||
"""Merge object flag masks onto a clean domain base.
|
||||
|
||||
Callers should pass a fresh channel-layout flag array. This keeps object
|
||||
overlays stateless and prevents stale obstacle bits from surviving after
|
||||
geometry edits or future body motion.
|
||||
Args:
|
||||
objects: List of SimObject instances.
|
||||
base_flags: Fresh channel-layout flag array.
|
||||
nx, ny, nz: Grid dimensions.
|
||||
|
||||
Returns:
|
||||
New uint16 array with object overlays applied.
|
||||
"""
|
||||
masks = [
|
||||
obj.get_flag_mask(self.nx, self.ny, self.nz)
|
||||
for obj in self._registry.objects
|
||||
]
|
||||
masks = [obj.get_flag_mask(nx, ny, nz) for obj in objects]
|
||||
return merge_flag_masks(base_flags, masks)
|
||||
|
||||
# -- Compact list building -----------------------------------------------
|
||||
def build_compact_lists(self, domain_flags: np.ndarray | None = None):
|
||||
"""Build cut-link SoA columns, sensor lists, and force-region lists
|
||||
from the formal registry objects.
|
||||
"""
|
||||
return self.build_compact_lists_for(
|
||||
self._registry.objects, domain_flags,
|
||||
sensor_cell_counts_out=self.sensor_cell_counts,
|
||||
)
|
||||
|
||||
def build_compact_lists_for(
|
||||
self,
|
||||
objects: List[SimObject],
|
||||
domain_flags: np.ndarray | None = None,
|
||||
*,
|
||||
sensor_cell_counts_out: np.ndarray | None = None,
|
||||
):
|
||||
"""Build cut-link SoA columns, sensor lists, and force-region lists.
|
||||
|
||||
Args:
|
||||
objects: List of SimObject instances to process.
|
||||
domain_flags: Optional uint16 flags for donor-cell classification.
|
||||
sensor_cell_counts_out: Int array to fill with per-object sensor
|
||||
cell counts. If provided, ``fill(0)`` is called first.
|
||||
|
||||
Returns:
|
||||
cl_fluid_idx, cl_dir, cl_q, cl_body_id, cl_rx, cl_ry, cl_rz,
|
||||
cl_fallback_class,
|
||||
@@ -141,7 +360,9 @@ class ObjectManager:
|
||||
cl_rx, cl_ry, cl_rz, cl_fallback = [], [], [], []
|
||||
s_cells, s_ids = [], []
|
||||
fr_cells, fr_ids = [], []
|
||||
self.sensor_cell_counts.fill(0)
|
||||
|
||||
if sensor_cell_counts_out is not None:
|
||||
sensor_cell_counts_out.fill(0)
|
||||
|
||||
ez = None
|
||||
if self.cfg and self.cfg.is_d3q19:
|
||||
@@ -149,8 +370,7 @@ class ObjectManager:
|
||||
else:
|
||||
ex, ey = D2Q9_EX, D2Q9_EY
|
||||
|
||||
for obj in self._registry.objects:
|
||||
# Curved list: only real obstacle bodies (not sensors, not force_regions).
|
||||
for obj in objects:
|
||||
if (not obj.is_sensor) and (not obj.is_force_region) and hasattr(obj, 'get_curved_list'):
|
||||
(
|
||||
fluid_idx, dirs, q_vals, body_ids,
|
||||
@@ -174,8 +394,9 @@ class ObjectManager:
|
||||
if len(cells) > 0:
|
||||
s_cells.append(cells)
|
||||
s_ids.append(ids)
|
||||
if 0 <= obj.obj_id < self.sensor_cell_counts.size:
|
||||
self.sensor_cell_counts[obj.obj_id] = int(len(cells))
|
||||
if (sensor_cell_counts_out is not None
|
||||
and 0 <= obj.obj_id < sensor_cell_counts_out.size):
|
||||
sensor_cell_counts_out[obj.obj_id] = int(len(cells))
|
||||
if obj.is_force_region and hasattr(obj, 'get_sensor_list'):
|
||||
cells, ids = obj.get_sensor_list(self.nx, self.ny, self.nz)
|
||||
if len(cells) > 0:
|
||||
@@ -257,7 +478,8 @@ class ObjectManager:
|
||||
self.action_gpu = None
|
||||
self.action_gpu = cuda.mem_alloc(action_nbytes)
|
||||
self._action_nbytes = action_nbytes
|
||||
cuda.memcpy_htod(self.action_gpu, self.action)
|
||||
# Initialise with host action (zeros until first run() upload).
|
||||
cuda.memcpy_htod(self.action_gpu, self.action)
|
||||
|
||||
if self.obs_gpu is None or self._obs_alloc_nbytes != self.obs_nbytes:
|
||||
if self.obs_gpu is not None:
|
||||
@@ -321,6 +543,12 @@ class ObjectManager:
|
||||
assert self.obs_pinned is not None
|
||||
cuda.memcpy_dtoh_async(self.obs_pinned, self.obs_gpu, stream)
|
||||
|
||||
def _upload_action_async(self, stream: cuda.Stream) -> None:
|
||||
"""Async H2D of host action array to action_gpu."""
|
||||
if self.action_gpu is None or self.action.size == 0:
|
||||
return
|
||||
cuda.memcpy_htod_async(self.action_gpu, self.action, stream)
|
||||
|
||||
@property
|
||||
def obs_n_slots(self) -> int:
|
||||
"""At-least-one slot count ``max(count, 1)`` matching ``config_obs.h``."""
|
||||
@@ -380,20 +608,41 @@ class ObjectManager:
|
||||
return values
|
||||
return values / np.float32(count)
|
||||
|
||||
def read_body(self, body_id: int) -> BodyTelemetry:
|
||||
"""Return unified telemetry for one body from the pinned obs buffer.
|
||||
|
||||
The caller must ensure ``run(sync_obs=True)`` or an explicit
|
||||
``download_obs_full_async + synchronize`` has completed.
|
||||
"""
|
||||
force = self.read_force(body_id)
|
||||
torque = self.read_torque(body_id)
|
||||
sensor = self.read_sensor(body_id, normalize=True)
|
||||
return BodyTelemetry(force=force, torque=torque, sensor=sensor)
|
||||
|
||||
def _obs_array(self) -> np.ndarray:
|
||||
"""Return the full pinned obs buffer as a flat float32 array.
|
||||
|
||||
Shape: ``(total_floats,)`` with layout ``[force|torque|sensor]``.
|
||||
Efficient for DRL batch reads.
|
||||
"""
|
||||
assert self.obs_pinned is not None
|
||||
return self.obs_pinned.copy()
|
||||
|
||||
def set_body_state(self, body_id: int, omega: float = 0.0) -> None:
|
||||
"""Set runtime body state used by kernels (currently only omega)."""
|
||||
"""Set runtime body state (host action array only, no H2D).
|
||||
|
||||
The GPU action buffer is updated at the next ``run()`` call.
|
||||
"""
|
||||
self._validate_body_id(body_id)
|
||||
dim = self.cfg.dim if self.cfg else 2
|
||||
slot = 3 * dim
|
||||
base = body_id * slot
|
||||
self.action[base + slot - 1] = np.float32(omega)
|
||||
if self.action_gpu is not None:
|
||||
cuda.memcpy_htod(self.action_gpu, self.action)
|
||||
|
||||
def set_force_state(self, body_id: int, fx: float = 0.0, fy: float = 0.0) -> None:
|
||||
"""Set force density on a force_region object (action slot 0/1).
|
||||
"""Set force density on a force_region object (host action only, no H2D).
|
||||
|
||||
Only touches the force component slots; does not affect omega.
|
||||
The GPU action buffer is updated at the next ``run()`` call.
|
||||
"""
|
||||
self._validate_body_id(body_id)
|
||||
dim = self.cfg.dim if self.cfg else 2
|
||||
@@ -401,8 +650,6 @@ class ObjectManager:
|
||||
base = body_id * slot
|
||||
self.action[base + 0] = np.float32(fx)
|
||||
self.action[base + 1] = np.float32(fy)
|
||||
if self.action_gpu is not None:
|
||||
cuda.memcpy_htod(self.action_gpu, self.action)
|
||||
|
||||
def _validate_body_id(self, body_id: int) -> None:
|
||||
if body_id < 0 or body_id >= self.count:
|
||||
|
||||
@@ -118,11 +118,11 @@ class SimObject:
|
||||
"""Return (nx*ny,) uint16 array with flag bits set for this object.
|
||||
|
||||
Sensors produce ``FLUID|SENSOR_FLAG``; bodies produce
|
||||
``SOLID|OBSTACLE|BC_CURVED``. Force-region objects produce zero
|
||||
(no flag overlay — forcing is applied via the compact-list kernel).
|
||||
``SOLID|OBSTACLE|BC_CURVED``; force-region objects produce
|
||||
``FLUID|FRC_REGION``.
|
||||
"""
|
||||
if self._is_force_region:
|
||||
return np.zeros(nx * max(ny, 1) * max(nz, 1), dtype=np.uint16)
|
||||
return self.geometry.build_force_region_flag_mask(nx, ny)
|
||||
if self._is_sensor:
|
||||
return self.geometry.build_sensor_flag_mask(nx, ny)
|
||||
return self.geometry.build_flag_mask(nx, ny)
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
# CelerisLab/body/sync_plan.py
|
||||
"""
|
||||
BodySyncPlan -- topology change descriptor for runtime body sync. (D2Q9 only)
|
||||
|
||||
Design:
|
||||
A ``BodySyncPlan`` is constructed by ``ObjectManager.build_sync_plan()``
|
||||
and consumed by ``Simulation.sync_bodies()``. It describes the full
|
||||
topology transition from the current formal object set to the next one,
|
||||
including flags, compact lists, and cell-level change masks.
|
||||
|
||||
The plan does NOT hold DDF data -- that must be read from ``LBMField``
|
||||
at sync time. Keeping the plan free of large array copies makes it
|
||||
cheap to construct and discard.
|
||||
|
||||
Change masks use D2Q9 neighbour offsets. D3Q19 support would require
|
||||
extending ``released_fluid_mask`` neighbour enumeration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .objects import SimObject
|
||||
|
||||
|
||||
@dataclass
|
||||
class BodySyncPlan:
|
||||
"""Topology transition plan produced by ``ObjectManager.build_sync_plan()``.
|
||||
|
||||
Attributes:
|
||||
next_objects: New formal object list with ``obj_id = 0..n-1``.
|
||||
next_count: Length of *next_objects*.
|
||||
next_flags: Merged uint16 flag array for the new topology.
|
||||
curved_host: 8-tuple of SoA arrays for CurvedLinkSoA.assign_host().
|
||||
sensor_host: 2-tuple (cells, obj_id) for SensorSoA.assign_host().
|
||||
force_region_host: 2-tuple (cells, obj_id) for ForceRegionSoA.assign_host().
|
||||
added_solid_mask: Bool array -- True where ``old_fluid -> new_solid``.
|
||||
released_fluid_mask: Bool array -- True where ``old_obstacle -> new_fluid``.
|
||||
new_sensor_counts: Int array of sensor cell counts per object.
|
||||
"""
|
||||
next_objects: List[SimObject]
|
||||
next_count: int
|
||||
next_flags: np.ndarray
|
||||
curved_host: tuple
|
||||
sensor_host: tuple
|
||||
force_region_host: tuple
|
||||
added_solid_mask: np.ndarray
|
||||
released_fluid_mask: np.ndarray
|
||||
new_sensor_counts: np.ndarray
|
||||
@@ -70,6 +70,9 @@ def save_checkpoint(field, stepper, lbm_cfg, bodies, path=None):
|
||||
hf.attrs["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%S")
|
||||
hf.attrs["step_count"] = stepper.step_count
|
||||
|
||||
# Object count metadata for runtime sync compatibility
|
||||
hf.attrs["n_objects"] = bodies.count
|
||||
|
||||
# Config as JSON string (human-readable in HDF5 viewers)
|
||||
cfg = lbm_cfg
|
||||
config_dict = {
|
||||
@@ -156,6 +159,17 @@ def load_checkpoint(path, field, stepper, lbm_cfg, bodies):
|
||||
"checkpoint file",
|
||||
)
|
||||
|
||||
# Object count validation: checkpoint's n_objects must match
|
||||
# the compiled N_OBJS (which equals current bodies.count).
|
||||
ckpt_n_objects = int(hf.attrs.get("n_objects", 0))
|
||||
if ckpt_n_objects != bodies.count:
|
||||
raise ValueError(
|
||||
f"Object count mismatch: checkpoint has {ckpt_n_objects} "
|
||||
f"objects, but the simulation is compiled with "
|
||||
f"N_OBJS={bodies.count}. Use sync_bodies() to adjust "
|
||||
"the object count before loading, or rebuild the "
|
||||
"simulation with the correct number of objects.")
|
||||
|
||||
# Restore GPU arrays
|
||||
ddf_buf = hf["ddf"][:]
|
||||
tmp_buf = hf["temp"][:]
|
||||
|
||||
@@ -28,7 +28,7 @@ Python `config.py` 只负责读取和校验,不是配置位置。
|
||||
| 字段 | 类型 | 默认 | 允许值 | 说明 |
|
||||
|------|------|------|--------|------|
|
||||
| `collision` | string | `"SRT"` | `SRT`, `TRT`, `MRT` | 碰撞算子 |
|
||||
| `streaming` | string | `"double_buffer"` | `double_buffer`, `esopull` | 流传输方式 |
|
||||
| `streaming` | string | `"double_buffer"` | `double_buffer`, `esopull` | 流传输方式。运行时 body 拓扑同步 (``sync_bodies()``) 仅支持 ``double_buffer`` |
|
||||
| `store_precision` | string | `"FP32"` | `FP32`, `FP16S`, `FP16C` | GPU 存储精度。当前运行时已实现 `FP32` 与 `FP16S`,`FP16C` 仍为保留选项 |
|
||||
| `ddf_shifting` | bool | false | | 存储 f−w 而非 f,提升 FP16 精度 |
|
||||
| `les.enabled` | bool | false | | LES Smagorinsky 子格模型 |
|
||||
@@ -89,4 +89,15 @@ step/one_step_*.cu → kernel 编排
|
||||
### 稳定性默认窗口(Stability defaults)
|
||||
|
||||
- 默认 `method.omega_guard = {min: 0.01, max: 1.99}`,用于约束碰撞频率与 LES 有效粘度路径。
|
||||
- 高 Re 调参建议将 `omega_guard.max` 保持在 `1.90-1.99` 区间;过高上界会增大接近 `omega -> 2` 的不稳定风险。
|
||||
- 高 Re 调参建议将 `omega_guard.max` 保持在 `1.90-1.99` 区间;过高上界会增大接近 `omega -> 2` 的不稳定风险。
|
||||
|
||||
### 运行时 body 拓扑同步
|
||||
|
||||
- 运行时增删 body(`add_body()` / `remove_body()` + `sync_bodies()`)仅支持 `streaming: "double_buffer"`。
|
||||
- `esopull` 模式下调用 `sync_bodies()` 会抛出 `NotImplementedError`。
|
||||
|
||||
### 力区域标记(Force region flag)
|
||||
|
||||
D2Q9/D3Q19 的 flag 扩展位 `0x0800` 定义为 `FRC_REGION`,用于标记 force_region 类型物体的格子。
|
||||
- ForceRegionKernel 不依赖 flag 位执行,flag 只用于后处理可视化。
|
||||
- `merge_flag_masks` 会正确合并 force_region 的 `FLUID | FRC_REGION` 标记。
|
||||
@@ -31,6 +31,7 @@ BC_MOVING = 0x0060
|
||||
# Extension [15:8]
|
||||
SENSOR_FLAG = 0x0100
|
||||
OBSTACLE = 0x0200
|
||||
FRC_REGION = 0x0800
|
||||
|
||||
# ===================================================================
|
||||
# D2Q9 lattice velocity vectors
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
# CelerisLab/lbm/equilibrium.py
|
||||
"""
|
||||
Host-side equilibrium and macroscopic helpers for D2Q9. (D2Q9 only)
|
||||
|
||||
These functions are shared by LBMField (macroscopic computation) and
|
||||
body DDF patching. They operate on host NumPy arrays only and do not
|
||||
require a CUDA module.
|
||||
|
||||
For the authoritative GPU-side equilibrium, see ``operators/equilibrium.cuh``.
|
||||
|
||||
Callers with ``dim != 2`` must not call these functions. The D3Q19
|
||||
equilibrium path lives in the CUDA kernel layer and is not mirrored on
|
||||
the host at present.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
# D2Q9 velocity vectors (must match descriptors.cuh / descriptors.py)
|
||||
_CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.float32)
|
||||
_CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.float32)
|
||||
_W = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36],
|
||||
dtype=np.float32)
|
||||
|
||||
|
||||
def compute_feq_d2q9(rho: float, ux: float, uy: float) -> np.ndarray:
|
||||
"""Compute D2Q9 equilibrium distribution ``f_eq``.
|
||||
|
||||
Args:
|
||||
rho: Macroscopic density.
|
||||
ux: x-velocity (lattice units).
|
||||
uy: y-velocity (lattice units).
|
||||
|
||||
Returns:
|
||||
9-element float32 array ``[f0 .. f8]``.
|
||||
"""
|
||||
u_sq = ux * ux + uy * uy
|
||||
feq = np.empty(9, dtype=np.float32)
|
||||
for i in range(9):
|
||||
cu = _CX[i] * ux + _CY[i] * uy
|
||||
feq[i] = rho * _W[i] * (
|
||||
1.0 + 3.0 * cu + 4.5 * cu * cu - 1.5 * u_sq)
|
||||
return feq
|
||||
|
||||
|
||||
def compute_macro_from_ddf(f: np.ndarray) -> tuple[float, float, float]:
|
||||
"""Compute rho, ux, uy from a 9-element D2Q9 DDF array.
|
||||
|
||||
Args:
|
||||
f: 9-element float32 array of distribution functions.
|
||||
|
||||
Returns:
|
||||
``(rho, ux, uy)`` tuple.
|
||||
"""
|
||||
rho = float(np.sum(f))
|
||||
rho_safe = rho if abs(rho) > 1e-12 else 1.0
|
||||
ux = float(np.sum(f * _CX)) / rho_safe
|
||||
uy = float(np.sum(f * _CY)) / rho_safe
|
||||
return rho, ux, uy
|
||||
@@ -402,6 +402,11 @@ class LBMField:
|
||||
"uy": self._macro_uy.reshape(self.ny, self.nx),
|
||||
}
|
||||
|
||||
# -- DDF patch for runtime body sync (delegated to body/ddf_patch.py) -----
|
||||
# patch_ddf_for_body_sync and upload_patched_ddf are imported and called
|
||||
# directly from Simulation.sync_bodies(). They are no longer methods on
|
||||
# LBMField.
|
||||
|
||||
def snapshot(self):
|
||||
self.download_ddf(force=True)
|
||||
self._ddf_snap = self.ddf.copy()
|
||||
|
||||
@@ -39,7 +39,8 @@
|
||||
#define FLAG_SENSOR ((uint16_t)0x0100)
|
||||
#define FLAG_OBSTACLE ((uint16_t)0x0200)
|
||||
#define FLAG_HALO ((uint16_t)0x0400)
|
||||
// 0x0800..0x8000 reserved for multi-GPU / AMR / multi-phase
|
||||
#define FLAG_FRC_REGION ((uint16_t)0x0800)
|
||||
// 0x1000..0x8000 reserved for multi-GPU / AMR / multi-phase
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Masks
|
||||
@@ -79,5 +80,6 @@ __device__ __forceinline__ bool has_bc(uint16_t fl) { return (fl & MASK_BC_T
|
||||
__device__ __forceinline__ bool is_sensor(uint16_t fl) { return (fl & FLAG_SENSOR) != 0; }
|
||||
__device__ __forceinline__ bool is_obstacle(uint16_t fl) { return (fl & FLAG_OBSTACLE) != 0; }
|
||||
__device__ __forceinline__ bool is_halo(uint16_t fl) { return (fl & FLAG_HALO) != 0; }
|
||||
__device__ __forceinline__ bool is_force_region(uint16_t fl) { return (fl & FLAG_FRC_REGION) != 0; }
|
||||
|
||||
#endif // CELERIS_CORE_FLAGS_CUH
|
||||
@@ -54,6 +54,8 @@ __device__ __forceinline__ uint16_t finalize_domain_flag(
|
||||
const uint16_t base = channel_flag_from_coords(x, y);
|
||||
if (is_sensor(fl) && base == FLAG_FLUID)
|
||||
return (uint16_t)(base | FLAG_SENSOR);
|
||||
if (is_force_region(fl) && base == FLAG_FLUID)
|
||||
return (uint16_t)(base | FLAG_FRC_REGION);
|
||||
return base;
|
||||
}
|
||||
|
||||
|
||||
+216
-26
@@ -29,6 +29,7 @@ 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:
|
||||
@@ -87,6 +88,10 @@ class Simulation:
|
||||
radius: float = 0.0, **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"``).
|
||||
@@ -95,7 +100,7 @@ class Simulation:
|
||||
**kwargs: passed to the geometry constructor.
|
||||
|
||||
Returns:
|
||||
int body_id.
|
||||
int body_id (formal id before initialize; -1 if staged).
|
||||
"""
|
||||
type_lower = str(type).strip().lower()
|
||||
|
||||
@@ -125,8 +130,27 @@ class Simulation:
|
||||
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 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()``.
|
||||
|
||||
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:
|
||||
@@ -174,7 +198,9 @@ class Simulation:
|
||||
# -- Runtime control -----------------------------------------------------
|
||||
def set_body(self, id: int, omega: float = None,
|
||||
vx: float = None, vy: float = None) -> None:
|
||||
"""Set runtime body state. Updates are implicitly uploaded to GPU.
|
||||
"""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().
|
||||
@@ -195,11 +221,11 @@ class Simulation:
|
||||
changed = True
|
||||
if changed:
|
||||
self.bodies._refresh_action_from_objects()
|
||||
if self.bodies.action_gpu is not None:
|
||||
cuda.memcpy_htod(self.bodies.action_gpu, self.bodies.action)
|
||||
|
||||
def set_force(self, id: int, fx: float = 0.0, fy: float = 0.0) -> None:
|
||||
"""Set force density on a force_region object (implicit GPU upload).
|
||||
"""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', ...).
|
||||
@@ -210,55 +236,113 @@ class Simulation:
|
||||
|
||||
# -- Telemetry readback --------------------------------------------------
|
||||
def read_force(self, id: int) -> np.ndarray:
|
||||
"""Download and return the force vector on body *id* (async stream, synced)."""
|
||||
self._stream_obs_download()
|
||||
"""Return the force vector on body *id* from the pinned obs buffer."""
|
||||
if self.bodies.obs_pinned is None:
|
||||
raise RuntimeError("No obs buffer. Call run() first.")
|
||||
return self.bodies.read_force(id)
|
||||
|
||||
def read_torque(self, id: int) -> np.ndarray:
|
||||
"""Download and return the torque on body *id* (async stream, synced)."""
|
||||
self._stream_obs_download()
|
||||
"""Return the torque on body *id* from the pinned obs buffer."""
|
||||
if self.bodies.obs_pinned is None:
|
||||
raise RuntimeError("No obs buffer. Call run() first.")
|
||||
return self.bodies.read_torque(id)
|
||||
|
||||
def read_sensor(self, id: int, *, normalize: bool = True) -> np.ndarray:
|
||||
"""Download and return the sensor reading for body *id* (async stream, synced).
|
||||
"""Return the sensor reading for body *id* from the pinned obs buffer.
|
||||
|
||||
Args:
|
||||
normalize: If True, return area-averaged velocity. If False,
|
||||
return the raw sum accumulated by the GPU SensorKernel.
|
||||
"""
|
||||
self._stream_obs_download()
|
||||
if self.bodies.obs_pinned is None:
|
||||
raise RuntimeError("No obs buffer. Call run() first.")
|
||||
return self.bodies.read_sensor(id, normalize=normalize)
|
||||
|
||||
def _stream_obs_download(self) -> None:
|
||||
"""Async obs download on internal stream, then synchronize."""
|
||||
if self.bodies.obs_pinned is not None:
|
||||
self.bodies.download_obs_full_async(self.stream)
|
||||
self.stream.synchronize()
|
||||
def read_body(self, id: int, *, stream: cuda.Stream | None = None):
|
||||
"""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.
|
||||
|
||||
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)
|
||||
|
||||
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 ---------------------------------------------------------
|
||||
def recompile(self):
|
||||
"""Re-generate config headers and recompile kernel.
|
||||
def _recompile(self, expected_n_objects: int | None = None) -> None:
|
||||
"""Re-generate config headers and recompile the kernel.
|
||||
|
||||
Call after changing compile-time parameters (collision model, etc.).
|
||||
If *expected_n_objects* is ``None``, uses ``self.bodies.count``.
|
||||
Preserves ``step_count`` and GPU DDF/flag buffers across the
|
||||
module reload. Does NOT preserve ``action_gpu`` / ``obs_gpu``
|
||||
(those are recreated by ``sync_to_gpu`` or ``run()``).
|
||||
|
||||
This is safe to call before or after ``initialize()``.
|
||||
"""
|
||||
if self._initialized:
|
||||
raise RuntimeError(
|
||||
"recompile() must not be called after initialize(); "
|
||||
"rebuild Simulation instead.")
|
||||
if expected_n_objects is None:
|
||||
expected_n_objects = self.bodies.count
|
||||
arch = self._resolve_compile_arch()
|
||||
compiler.generate_config(self.lbm_cfg, n_objects=self.bodies.count)
|
||||
compiler.generate_config(self.lbm_cfg, n_objects=expected_n_objects)
|
||||
# Remove stale PTX to prevent PyCUDA module handle conflicts.
|
||||
import os as _os
|
||||
if _os.path.exists(self._ptx_path):
|
||||
_os.remove(self._ptx_path)
|
||||
self._ptx_path = compiler.compile_kernel(arch=arch)
|
||||
self._module = compiler.load_module(self._ptx_path)
|
||||
|
||||
# Reconnect field and stepper to new module
|
||||
self.field.module = self._module
|
||||
self.field.invalidate_module_constant_cache()
|
||||
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=self.bodies.count)
|
||||
|
||||
self._assert_object_count_contract(expected_count=expected_n_objects)
|
||||
|
||||
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 ------------------------------------------------------
|
||||
def initialize(self):
|
||||
@@ -281,15 +365,113 @@ class Simulation:
|
||||
self._assert_runtime_contracts()
|
||||
self._initialized = True
|
||||
|
||||
# -- Runtime body topology sync -------------------------------------------
|
||||
def sync_bodies(self) -> None:
|
||||
"""Apply pending body edits (add/remove) to a running simulation.
|
||||
|
||||
This is the main entry point for runtime body topology changes.
|
||||
It downloads the current DDF, rebuilds topology, recompiles the
|
||||
kernel for the new object count, patches the DDF for geometry
|
||||
changes, and re-uploads everything to GPU.
|
||||
|
||||
Currently only supports ``double_buffer`` streaming mode.
|
||||
``esopull`` mode raises ``NotImplementedError``.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called before ``initialize()``.
|
||||
NotImplementedError: If streaming mode is ``esopull``.
|
||||
"""
|
||||
if not self._initialized:
|
||||
raise RuntimeError("Call initialize() before sync_bodies()")
|
||||
|
||||
# 1. Synchronize GPU
|
||||
self.stream.synchronize()
|
||||
|
||||
# 2. Check for pending edits
|
||||
if not self.bodies.has_pending_edit():
|
||||
return
|
||||
|
||||
# 3. Check streaming mode
|
||||
if self.lbm_cfg.streaming == "esopull":
|
||||
raise NotImplementedError(
|
||||
"Runtime body sync is not yet supported for esopull "
|
||||
"streaming mode. Use double_buffer instead.")
|
||||
|
||||
# 4. Download current DDF to host (snapshot for patching)
|
||||
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. Build sync plan
|
||||
plan = self.bodies.build_sync_plan(self.field)
|
||||
|
||||
# 6. Runtime recompile for new object count
|
||||
self._recompile(plan.next_count)
|
||||
|
||||
# 7. Apply new topology data (flags, compact lists, params)
|
||||
self.bodies.apply_sync_plan(self.field, plan)
|
||||
|
||||
# 8. DDF patch for geometry changes
|
||||
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,
|
||||
)
|
||||
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 ------------------------------------------------------------
|
||||
def run(self, steps: int, checkpoint_interval: int = 0):
|
||||
def run(self, steps: int, *,
|
||||
stream: cuda.Stream | None = None,
|
||||
upload_act: bool = True,
|
||||
sync_obs: bool = True,
|
||||
checkpoint_interval: int = 0):
|
||||
"""Advance simulation by *steps* time steps.
|
||||
|
||||
Args:
|
||||
steps: Number of LBM steps.
|
||||
stream: CUDA stream for async operations. ``None`` = internal stream.
|
||||
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.
|
||||
checkpoint_interval: If >0, save checkpoint every N steps.
|
||||
"""
|
||||
if not self._initialized:
|
||||
raise RuntimeError("Call initialize() first")
|
||||
# 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 force segment before step group
|
||||
self.bodies.zero_force_segment_async(stream)
|
||||
|
||||
self._assert_runtime_contracts()
|
||||
if checkpoint_interval > 0:
|
||||
done = 0
|
||||
@@ -299,6 +481,7 @@ class Simulation:
|
||||
batch,
|
||||
action_gpu=self.bodies.action_gpu,
|
||||
obs_gpu=self.bodies.obs_gpu,
|
||||
stream=stream,
|
||||
)
|
||||
done += batch
|
||||
if done < steps or done % checkpoint_interval == 0:
|
||||
@@ -308,8 +491,15 @@ class Simulation:
|
||||
steps,
|
||||
action_gpu=self.bodies.action_gpu,
|
||||
obs_gpu=self.bodies.obs_gpu,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
Reference in new issue
Block a user