ClaudeandClaude Opus 5.5
d318911221
The package already exposed it as CelerisLab.cuda.compiler; the module now
has that name and every tip caller, doc and current-tree test imports it
directly (no alias).
Historical-revision tooling is not edited: the Taylor-Green harness
(run_case.py, hash-pinned check_projected_pairs.py, assess receipts) only
runs against a135252, whose package has cuda/compiler_v2.py, and the J
campaign runners and dated docs describe past trees. Its two current-tree
tests now import that package from git (historical_package) instead of the
tip, and the kernel-root test first checks the accepted case.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
864 lines
35 KiB
Python
864 lines
35 KiB
Python
# CelerisLab/body/manager.py
|
|
"""
|
|
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`
|
|
|
|
Design::
|
|
Packed ``obs`` layout matches ``generate_config(..., n_objects=count)`` and
|
|
``config_obs.h`` macros. Host sizing uses :func:`~CelerisLab.cuda.compiler.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 .objects import SimObject
|
|
from .registry import BodyRegistry
|
|
from .preprocess.flag_overlay import merge_flag_masks, validate_cut_link_ownership
|
|
from ..cuda.compiler import obs_layout
|
|
from ..cuda.obs_layout import ObsLayout
|
|
from ..lbm.descriptors import (
|
|
FLUID, SOLID,
|
|
D2Q9_EX,
|
|
D2Q9_EY,
|
|
D3Q19_EX,
|
|
D3Q19_EY,
|
|
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.
|
|
|
|
Frozen telemetry layout::
|
|
``obs`` packs force, torque, then sensor segments.
|
|
Segment offsets and sizes come from :func:`obs_layout(dim, count)`.
|
|
Device macros ``OBS_*`` in ``config_obs.h`` must match ``obs_layout(dim, count)``.
|
|
"""
|
|
|
|
def __init__(self, nx: int, ny: int, nz: int, nq: int, cfg=None):
|
|
self.nx = nx
|
|
self.ny = ny
|
|
self.nz = nz
|
|
self.nq = nq
|
|
self.cfg = cfg
|
|
self._registry = BodyRegistry()
|
|
self.jets: list = [] # scene-level Jets; selection is spatial, not per body
|
|
|
|
self.action_gpu: Optional[cuda.DeviceAllocation] = None
|
|
self.obs_gpu: Optional[cuda.DeviceAllocation] = None
|
|
self._action_nbytes: int = 0
|
|
self._obs_alloc_nbytes: int = 0
|
|
|
|
self.action = np.zeros(0, dtype=np.float32)
|
|
self._action_pinned: Optional[np.ndarray] = None
|
|
self.obs_pinned: Optional[np.ndarray] = None
|
|
self.obs_force_view: Optional[np.ndarray] = None
|
|
self.obs_torque_view: Optional[np.ndarray] = None
|
|
self.obs_sensor_view: Optional[np.ndarray] = None
|
|
self.sensor_cell_counts = np.zeros(0, dtype=np.int32)
|
|
self.slot_stride_floats: int = 0
|
|
self.torque_stride_floats: int = 0
|
|
self.torque_components: int = 0
|
|
self.torque0_floats: int = 0
|
|
self.sensor0_floats: int = 0
|
|
self.obs_total_floats: int = 0
|
|
self.obs_nbytes: int = 0
|
|
|
|
# -- Accumulated step count (for obs time-normalization) --------------
|
|
self._obs_accum_steps: 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()
|
|
|
|
# -- Registry delegation -------------------------------------------------
|
|
def add(self, obj: SimObject) -> int:
|
|
"""Register an object and return its id."""
|
|
obj.validate_surface_actuation(self.cfg)
|
|
body_id = self._registry.add(obj)
|
|
self._resize_buffers()
|
|
return body_id
|
|
|
|
def remove(self, obj_id: int):
|
|
"""Remove immediately and renumber remaining IDs, as runtime sync does.
|
|
|
|
Previously returned numeric IDs are invalidated by this operation.
|
|
"""
|
|
self._registry.remove(obj_id)
|
|
remaining = self.objects
|
|
self._registry.clear()
|
|
for obj in remaining:
|
|
self._registry.add(obj)
|
|
self._resize_buffers()
|
|
self._refresh_action_from_objects()
|
|
|
|
def get(self, obj_id: int) -> SimObject:
|
|
return self._registry.get(obj_id)
|
|
|
|
@property
|
|
def objects(self) -> List[SimObject]:
|
|
return self._registry.objects
|
|
|
|
@property
|
|
def count(self) -> int:
|
|
return self._registry.count
|
|
|
|
# -- Buffer sizing -------------------------------------------------------
|
|
def _resize_buffers(self):
|
|
n = self.count
|
|
dim = self.cfg.dim if self.cfg else 2
|
|
action_slot_floats = 3 * dim
|
|
new_action = np.zeros(max(n * action_slot_floats, 1), dtype=np.float32)
|
|
if self.action.size > 0:
|
|
copy_n = min(self.action.size, new_action.size)
|
|
new_action[:copy_n] = self.action[:copy_n]
|
|
self.action = new_action
|
|
|
|
new_sensor_counts = np.zeros(max(n, 1), dtype=np.int32)
|
|
if self.sensor_cell_counts.size > 0:
|
|
copy_n = min(self.sensor_cell_counts.size, new_sensor_counts.size)
|
|
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.
|
|
"""
|
|
obj.validate_surface_actuation(self.cfg)
|
|
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,
|
|
surface_actuation=obj.surface_actuation)
|
|
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)
|
|
hits = []
|
|
(
|
|
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, hits_out=hits,
|
|
)
|
|
|
|
# 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,
|
|
curved_hits=hits[0],
|
|
)
|
|
|
|
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)
|
|
field.curved.assign_surface(plan.next_objects, self.cfg, plan.next_flags,
|
|
jets=self.jets, hits=plan.curved_hits)
|
|
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.
|
|
|
|
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(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, *, hits_out=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, hits_out=hits_out,
|
|
)
|
|
|
|
def build_compact_lists_for(
|
|
self,
|
|
objects: List[SimObject],
|
|
domain_flags: np.ndarray | None = None,
|
|
*,
|
|
sensor_cell_counts_out: np.ndarray | None = None,
|
|
hits_out: list | 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 receiver/donor classification.
|
|
sensor_cell_counts_out: Int array to fill with per-object sensor
|
|
cell counts after the full geometry has passed validation.
|
|
hits_out: List that receives one FP64 (n, 2) array of true hit
|
|
points row-aligned with the cut-link columns (for Jets).
|
|
|
|
Returns:
|
|
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
|
|
"""
|
|
cl_idx, cl_dir, cl_q, cl_bid = [], [], [], []
|
|
cl_rx, cl_ry, cl_rz, cl_fallback = [], [], [], []
|
|
cl_hit = []
|
|
s_cells, s_ids = [], []
|
|
fr_cells, fr_ids = [], []
|
|
if domain_flags is None:
|
|
domain_flags = self.build_flags_for(
|
|
objects, np.full(self.nx * self.ny * self.nz, FLUID, dtype=np.uint16),
|
|
nx=self.nx, ny=self.ny, nz=self.nz,
|
|
)
|
|
active_fluid = ((domain_flags & FLUID) != 0) & ((domain_flags & SOLID) == 0)
|
|
force_owners: set[int] = set()
|
|
|
|
sensor_counts = (np.zeros_like(sensor_cell_counts_out)
|
|
if sensor_cell_counts_out is not None else None)
|
|
|
|
ez = None
|
|
if self.cfg and self.cfg.is_d3q19:
|
|
ex, ey, ez = D3Q19_EX, D3Q19_EY, D3Q19_EZ
|
|
else:
|
|
ex, ey = D2Q9_EX, D2Q9_EY
|
|
|
|
for obj in objects:
|
|
obj.validate_surface_actuation(self.cfg)
|
|
if (not obj.is_sensor) and (not obj.is_force_region) and hasattr(obj, 'get_curved_list'):
|
|
(
|
|
fluid_idx, dirs, q_vals, body_ids,
|
|
rx_vals, ry_vals, rz_vals, fallback_vals
|
|
) = obj.get_curved_list(
|
|
self.nx, self.ny, self.nq,
|
|
nz=self.nz, ex=ex, ey=ey, ez=ez,
|
|
domain_flags=domain_flags, hits_out=cl_hit,
|
|
)
|
|
if len(fluid_idx) > 0:
|
|
cl_idx.append(fluid_idx)
|
|
cl_dir.append(dirs)
|
|
cl_q.append(q_vals)
|
|
cl_bid.append(body_ids)
|
|
cl_rx.append(rx_vals)
|
|
cl_ry.append(ry_vals)
|
|
cl_rz.append(rz_vals)
|
|
cl_fallback.append(fallback_vals)
|
|
if obj.is_sensor and hasattr(obj, 'get_sensor_list'):
|
|
cells, ids = obj.get_sensor_list(self.nx, self.ny, self.nz)
|
|
keep = active_fluid[cells]
|
|
cells, ids = cells[keep], ids[keep]
|
|
if len(cells) > 0:
|
|
s_cells.append(cells)
|
|
s_ids.append(ids)
|
|
if (sensor_counts is not None
|
|
and 0 <= obj.obj_id < sensor_counts.size):
|
|
sensor_counts[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)
|
|
keep = active_fluid[cells]
|
|
cells, ids = cells[keep], ids[keep]
|
|
if len(cells) > 0:
|
|
if force_owners.intersection(map(int, cells)):
|
|
raise ValueError("Overlapping force_region fluid cells are unsupported")
|
|
force_owners.update(map(int, cells))
|
|
fr_cells.append(cells)
|
|
fr_ids.append(ids)
|
|
|
|
r_idx = np.concatenate(cl_idx) if cl_idx else np.zeros(0, dtype=np.uint32)
|
|
r_dir = np.concatenate(cl_dir) if cl_dir else np.zeros(0, dtype=np.uint8)
|
|
r_q = np.concatenate(cl_q) if cl_q else np.zeros(0, dtype=np.float32)
|
|
r_bid = np.concatenate(cl_bid) if cl_bid else np.zeros(0, dtype=np.int32)
|
|
r_rx = np.concatenate(cl_rx) if cl_rx else np.zeros(0, dtype=np.float32)
|
|
r_ry = np.concatenate(cl_ry) if cl_ry else np.zeros(0, dtype=np.float32)
|
|
r_rz = np.concatenate(cl_rz) if cl_rz else np.zeros(0, dtype=np.float32)
|
|
r_fallback = (
|
|
np.concatenate(cl_fallback) if cl_fallback
|
|
else np.zeros(0, dtype=np.uint8)
|
|
)
|
|
sensor_cells = np.concatenate(s_cells) if s_cells else np.zeros(0, dtype=np.uint32)
|
|
sensor_obj_id = np.concatenate(s_ids) if s_ids else np.zeros(0, dtype=np.int32)
|
|
r_fr_cells = np.concatenate(fr_cells) if fr_cells else np.zeros(0, dtype=np.uint32)
|
|
r_fr_ids = np.concatenate(fr_ids) if fr_ids else np.zeros(0, dtype=np.int32)
|
|
|
|
r_hit = np.concatenate(cl_hit) if cl_hit else np.zeros((0, 2))
|
|
# Passive walls need the same ownership guarantees as driven walls.
|
|
# Overlapping solids are not a supported union-boundary operation.
|
|
validate_cut_link_ownership(domain_flags, r_idx, r_dir, self.nq)
|
|
if self.jets or any(obj.surface_actuation is not None for obj in objects):
|
|
# Reject unresolved apertures/ownership before sync or restart mutates state.
|
|
from ..lbm.curved_links import CurvedLinkSoA
|
|
links = CurvedLinkSoA()
|
|
links.assign_host(r_idx, r_dir, r_q, r_bid, r_rx, r_ry, r_rz, r_fallback)
|
|
links.assign_surface(objects, self.cfg, domain_flags, jets=self.jets, hits=r_hit)
|
|
|
|
if sensor_cell_counts_out is not None:
|
|
sensor_cell_counts_out[:] = sensor_counts
|
|
if hits_out is not None:
|
|
hits_out.append(r_hit)
|
|
|
|
return (
|
|
r_idx, r_dir, r_q, r_bid,
|
|
r_rx, r_ry, r_rz, r_fallback,
|
|
sensor_cells, sensor_obj_id,
|
|
r_fr_cells, r_fr_ids,
|
|
)
|
|
|
|
# -- GPU sync ------------------------------------------------------------
|
|
def sync_to_gpu(self, field, *, rebuild_flags: bool = True):
|
|
"""Upload compact lists, action buffer, and packed ``obs`` GPU buffer.
|
|
|
|
Args:
|
|
rebuild_flags: Rebuild ``field.flag`` from the clean channel base and
|
|
upload it before constructing compact lists. Pass ``False`` when
|
|
the caller has already prepared and uploaded the final flag field
|
|
for this geometry, such as during initialization.
|
|
"""
|
|
next_flags = self.build_flags(field.build_channel_flags()) if rebuild_flags else field.flag
|
|
hits = []
|
|
|
|
(
|
|
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(domain_flags=next_flags, hits_out=hits)
|
|
|
|
self._telemetry_field = field
|
|
if rebuild_flags:
|
|
field.flag = next_flags
|
|
field.upload_flags()
|
|
|
|
field.curved.assign_host(
|
|
cl_fluid_idx, cl_dir, cl_q, cl_body_id,
|
|
cl_rx, cl_ry, cl_rz, cl_fallback_class,
|
|
)
|
|
field.curved.assign_surface(self.objects, self.cfg, field.flag,
|
|
jets=self.jets, hits=hits[0])
|
|
field.sensors.assign_host(sensor_cells, sensor_obj_id)
|
|
field.force_regions.assign_host(fr_cells, fr_obj_id)
|
|
|
|
field.upload_compact_lists()
|
|
field.update_params(n_objects=self.count)
|
|
|
|
self._refresh_action_from_objects()
|
|
self._allocate_packed_telemetry()
|
|
assert self.obs_pinned is not None
|
|
self.obs_pinned.fill(0)
|
|
cuda.memcpy_htod(self.obs_gpu, self.obs_pinned)
|
|
|
|
def _allocate_packed_telemetry(self) -> None:
|
|
"""Resize ``action_gpu`` / ``obs_gpu`` / ``obs_pinned`` when layout changes."""
|
|
dim = self.cfg.dim if self.cfg else 2
|
|
lay = obs_layout(dim, self.count)
|
|
self._apply_obs_layout(lay)
|
|
|
|
# Buffer re-allocation implies a fresh start -- reset step counter.
|
|
self._obs_accum_steps = 0
|
|
|
|
action_nbytes = int(self.action.nbytes)
|
|
if self.action_gpu is None or self._action_nbytes != action_nbytes:
|
|
if self.action_gpu is not None:
|
|
self.action_gpu.free()
|
|
self.action_gpu = None
|
|
self.action_gpu = cuda.mem_alloc(action_nbytes)
|
|
self._action_nbytes = action_nbytes
|
|
# 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:
|
|
self.obs_gpu.free()
|
|
self.obs_gpu = None
|
|
self.obs_gpu = cuda.mem_alloc(self.obs_nbytes)
|
|
self._obs_alloc_nbytes = self.obs_nbytes
|
|
self.obs_pinned = cuda.pagelocked_empty(lay.total_floats, np.float32)
|
|
self.obs_pinned.fill(0)
|
|
n_slots = max(self.count, 1)
|
|
self.obs_force_view = self.obs_pinned[
|
|
lay.force0_floats: lay.force0_floats + lay.force_floats
|
|
].reshape(
|
|
n_slots, dim)
|
|
self.obs_torque_view = self.obs_pinned[
|
|
lay.torque0_floats: lay.torque0_floats + lay.torque_floats
|
|
].reshape(n_slots, lay.torque_components)
|
|
self.obs_sensor_view = self.obs_pinned[
|
|
lay.sensor0_floats: lay.sensor0_floats + lay.slot_stride_floats
|
|
].reshape(n_slots, dim)
|
|
|
|
def _apply_obs_layout(self, lay: ObsLayout) -> None:
|
|
self.slot_stride_floats = lay.slot_stride_floats
|
|
self.torque_stride_floats = lay.torque_floats
|
|
self.torque_components = lay.torque_components
|
|
self.torque0_floats = lay.torque0_floats
|
|
self.sensor0_floats = lay.sensor0_floats
|
|
self.obs_total_floats = lay.total_floats
|
|
self.obs_nbytes = lay.total_floats * 4
|
|
|
|
def ensure_packed_buffers(self) -> None:
|
|
"""Allocate ``action_gpu`` and packed ``obs`` when there are zero objects.
|
|
|
|
Call after ``generate_config(..., n_objects=0)`` so ``config_obs.h`` matches
|
|
:func:`obs_layout(dim, 0)`.
|
|
"""
|
|
if self.cfg is None:
|
|
return
|
|
self._allocate_packed_telemetry()
|
|
assert self.obs_pinned is not None
|
|
self.obs_pinned.fill(0)
|
|
cuda.memcpy_htod(self.obs_gpu, self.obs_pinned)
|
|
|
|
def zero_force_segment_async(self, stream: cuda.Stream):
|
|
"""Zero body telemetry (force + torque) of ``obs_gpu``.
|
|
|
|
Deprecated: prefer :meth:`zero_obs_async` which zeros all three
|
|
obs segments (force + torque + sensor) in a single memset.
|
|
"""
|
|
n_floats = self.sensor0_floats
|
|
cuda.memset_d32_async(self.obs_gpu, 0, n_floats, stream)
|
|
|
|
def zero_sensor_segment_async(self, stream: cuda.Stream):
|
|
"""Zero the sensor segment (second stride-sized block of floats).
|
|
|
|
Deprecated: prefer :meth:`zero_obs_async` which zeros all three
|
|
obs segments (force + torque + sensor) in a single memset.
|
|
"""
|
|
offset_bytes = self.sensor0_floats * 4
|
|
ptr = int(self.obs_gpu) + offset_bytes
|
|
cuda.memset_d32_async(ptr, 0, self.slot_stride_floats, stream)
|
|
|
|
def zero_obs_async(self, stream: cuda.Stream):
|
|
"""Zero ALL obs segments (force + torque + sensor) on GPU.
|
|
|
|
One ``memset`` covers the entire ``obs_gpu`` buffer.
|
|
Resets the step accumulator so that ``read_*(normalize=True)``
|
|
returns raw values (dividing by zero is avoided -- see read methods).
|
|
"""
|
|
n_floats = self.obs_total_floats
|
|
cuda.memset_d32_async(self.obs_gpu, 0, n_floats, stream)
|
|
self._obs_accum_steps = 0
|
|
|
|
def increment_obs_steps(self, n: int):
|
|
"""Accumulate *n* LBM steps for time-normalization."""
|
|
self._obs_accum_steps += n
|
|
|
|
def download_obs_full_async(self, stream: cuda.Stream):
|
|
"""Enqueue full DTOH copy ``obs_gpu`` -> ``obs_pinned``."""
|
|
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.
|
|
|
|
Single gate for every action path (setters or direct ``action`` edits):
|
|
surface-actuated and Jet-driven bodies must keep zero motion slots.
|
|
"""
|
|
slot = 3*(self.cfg.dim if self.cfg else 2)
|
|
for obj in self.objects:
|
|
obj.validate_surface_actuation(self.cfg)
|
|
base = obj.obj_id*slot
|
|
if obj.surface_actuation is not None and self.action[base+slot-1] != 0:
|
|
raise ValueError("Surface actuation cannot share the rotation action")
|
|
self.validate_jet_motion(obj.obj_id, {"vx": self.action[base], "vy": self.action[base+1],
|
|
"omega": self.action[base+slot-1]})
|
|
if self.action_gpu is None or self.action.size == 0:
|
|
return
|
|
if self._action_pinned is None or self._action_pinned.size != self.action.size:
|
|
self._action_pinned = cuda.pagelocked_empty(self.action.size, np.float32)
|
|
self._action_pinned[:] = self.action
|
|
cuda.memcpy_htod_async(self.action_gpu, self._action_pinned, stream)
|
|
|
|
@property
|
|
def obs_n_slots(self) -> int:
|
|
"""At-least-one slot count ``max(count, 1)`` matching ``config_obs.h``."""
|
|
return max(self.count, 1)
|
|
|
|
@property
|
|
def obs_slot_stride_floats(self) -> int:
|
|
"""Floats per telemetry segment (``obs_n_slots * DIM``)."""
|
|
return self.slot_stride_floats
|
|
|
|
@property
|
|
def obs_sensor_base_floats(self) -> int:
|
|
"""Float index where the sensor segment begins."""
|
|
return self.sensor0_floats
|
|
|
|
@property
|
|
def obs_torque_base_floats(self) -> int:
|
|
"""Float index where the torque segment begins."""
|
|
return self.torque0_floats
|
|
|
|
def read_force(self, body_id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return the DIM-vector force on body ``body_id`` from ``obs_pinned``.
|
|
|
|
When *normalize* is ``True`` (default), divides by the number of
|
|
accumulated LBM steps since the last zero so the result is the
|
|
**average force per step**. When ``False``, returns the raw
|
|
GPU-accumulated sum (no time-normalisation).
|
|
|
|
Caller must have synchronised the CUDA stream before reading.
|
|
"""
|
|
self._validate_body_id(body_id)
|
|
assert self.cfg is not None
|
|
assert self.obs_pinned is not None
|
|
d = self.cfg.dim
|
|
i0 = body_id * d
|
|
values = np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)
|
|
if normalize and self._obs_accum_steps > 0:
|
|
values /= np.float32(self._obs_accum_steps)
|
|
return values
|
|
|
|
def read_torque(self, body_id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return torque vector for ``body_id`` from ``obs_pinned``.
|
|
|
|
See :meth:`read_force` for the *normalize* semantics.
|
|
"""
|
|
self._validate_body_id(body_id)
|
|
assert self.obs_pinned is not None
|
|
i0 = self.torque0_floats + body_id * self.torque_components
|
|
values = np.array(
|
|
self.obs_pinned[i0:i0 + self.torque_components], dtype=np.float32)
|
|
if normalize and self._obs_accum_steps > 0:
|
|
values /= np.float32(self._obs_accum_steps)
|
|
return values
|
|
|
|
def read_sensor(self, body_id: int, *, normalize: bool = True) -> np.ndarray:
|
|
"""Return sensor accumulation for ``body_id``.
|
|
|
|
**Area-normalisation is always applied internally** -- the raw GPU
|
|
sum is divided by the number of sensor cells in the body footprint.
|
|
When *normalize* is ``True`` (default), the result is further divided
|
|
by the number of accumulated LBM steps, giving a **per-step
|
|
area-averaged velocity**. Set ``normalize=False`` to obtain the
|
|
area-averaged value without time-normalisation.
|
|
|
|
Caller must have synchronised the CUDA stream before reading.
|
|
"""
|
|
self._validate_body_id(body_id)
|
|
assert self.obs_pinned is not None
|
|
assert self.cfg is not None
|
|
d = self.cfg.dim
|
|
i0 = self.sensor0_floats + body_id * d
|
|
values = np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)
|
|
# Always area-normalise
|
|
count = int(self.sensor_cell_counts[body_id]) if body_id < self.sensor_cell_counts.size else 0
|
|
if count > 0:
|
|
values /= np.float32(count)
|
|
# Optionally time-normalise
|
|
if normalize and self._obs_accum_steps > 0:
|
|
values /= np.float32(self._obs_accum_steps)
|
|
return values
|
|
|
|
def read_body(self, body_id: int, *, normalize: bool = True) -> BodyTelemetry:
|
|
"""Return unified telemetry for one body from the pinned obs buffer.
|
|
|
|
See :meth:`read_force`, :meth:`read_torque`, and :meth:`read_sensor`
|
|
for the *normalize* semantics applied to each field.
|
|
|
|
The caller must ensure ``run(sync_obs=True)`` or an explicit
|
|
``download_obs_full_async + synchronize`` has completed.
|
|
"""
|
|
force = self.read_force(body_id, normalize=normalize)
|
|
torque = self.read_torque(body_id, normalize=normalize)
|
|
sensor = self.read_sensor(body_id, normalize=normalize)
|
|
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 validate_jet_motion(self, body_id: int, state: dict) -> None:
|
|
"""Reject motion of a body whose synced cut links are Jet-driven.
|
|
|
|
Before the first sync no links exist; sync-time checks cover that case.
|
|
"""
|
|
field = self._telemetry_field
|
|
cl = field.curved if field is not None else None
|
|
if (cl is not None and cl.jet_id is not None
|
|
and np.any(cl.body_id[cl.jet_id >= 0] == body_id)
|
|
and any(state.get(k, 0.) != 0 for k in ("vx", "vy", "omega"))):
|
|
raise ValueError("Jet-driven bodies must stay stationary")
|
|
|
|
def set_body_state(self, body_id: int, omega: float = 0.0) -> None:
|
|
"""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)
|
|
obj = self.get(body_id)
|
|
proposed = dict(vars(obj.state), omega=float(omega))
|
|
obj.validate_surface_actuation(self.cfg, proposed)
|
|
self.validate_jet_motion(body_id, proposed)
|
|
dim = self.cfg.dim if self.cfg else 2
|
|
slot = 3 * dim
|
|
base = body_id * slot
|
|
self.get(body_id).state.omega = float(omega)
|
|
self.action[base + slot - 1] = np.float32(omega)
|
|
|
|
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 (host action 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
|
|
obj = self.get(body_id)
|
|
if not obj.is_force_region:
|
|
raise ValueError("set_force_state requires a force_region object")
|
|
# vx/vy are the existing force-region action aliases. Keep their
|
|
# authoritative host state in sync so edits to other bodies cannot
|
|
# discard this control during a full action refresh or topology sync.
|
|
obj.state.vx = float(fx)
|
|
obj.state.vy = float(fy)
|
|
self.action[base + 0] = np.float32(fx)
|
|
self.action[base + 1] = np.float32(fy)
|
|
|
|
def _validate_body_id(self, body_id: int) -> None:
|
|
if body_id < 0 or body_id >= self.count:
|
|
raise IndexError(
|
|
f"body_id {body_id} out of range for {self.count} objects")
|
|
|
|
def _refresh_action_from_objects(self) -> None:
|
|
"""Populate action slots from object state using the runtime contract.
|
|
|
|
Slot layout per object (DIM=2 in this example, slots derived from 3*DIM):
|
|
[0] = vx (= fx for force_region objects; 0 for bodies/sensors)
|
|
[1] = vy (= fy for force_region objects; 0 for bodies/sensors)
|
|
[5] = omega (= body rotation for circle objects; 0 for sensors/force_region)
|
|
|
|
Slot 0/1 are consumed by ForceRegionKernel for force_region objects.
|
|
Slot -1 is consumed by CurvedBoundaryKernel for body rotation (action_omega).
|
|
Sensors do not consume action slots.
|
|
"""
|
|
if self.count == 0:
|
|
return
|
|
dim = self.cfg.dim if self.cfg else 2
|
|
slot = 3 * dim
|
|
self.action.fill(0)
|
|
for obj in self._registry.objects:
|
|
base = obj.obj_id * slot
|
|
if base + slot > self.action.size:
|
|
continue
|
|
self.action[base + 0] = np.float32(getattr(obj.state, "vx", 0.0))
|
|
self.action[base + 1] = np.float32(getattr(obj.state, "vy", 0.0))
|
|
self.action[base + slot - 1] = np.float32(
|
|
getattr(obj.state, "omega", 0.0))
|
|
|
|
def fallback_link_count(self) -> int:
|
|
"""Return number of curved links marked for fallback bounce-back."""
|
|
if self._telemetry_field is None:
|
|
return 0
|
|
curved = self._telemetry_field.curved
|
|
if curved.count == 0:
|
|
return 0
|
|
return int(np.count_nonzero(curved.fallback_class[:curved.count] != 0))
|
|
|
|
def low_q_link_count(self, threshold: float = 0.5) -> int:
|
|
"""Return count of links where q is below ``threshold``."""
|
|
if self._telemetry_field is None:
|
|
return 0
|
|
curved = self._telemetry_field.curved
|
|
if curved.count == 0:
|
|
return 0
|
|
return int(np.count_nonzero(
|
|
curved.q[:curved.count] < np.float32(threshold)))
|
|
|
|
# -- Placeholders for future use -----------------------------------------
|
|
def update_states(self, dt: float):
|
|
"""Integrate object motion (placeholder)."""
|
|
pass
|
|
|
|
def detect_collisions(self):
|
|
"""Multi-body collision detection (placeholder)."""
|
|
pass
|
|
|
|
def exchange_forces(self, obs: np.ndarray):
|
|
"""Distribute fluid forces to objects (placeholder)."""
|
|
pass
|