Frank14f/CelerisLab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(obs): unified zero_obs control and time-normalised readback

Browse Source 
- Replace split zero_force_segment / zero_sensor_segment with unified
  zero_obs_async() — a single memset covers all three obs segments
  (force, torque, sensor), resetting the step accumulator.
- Add _obs_accum_steps counter so read_*(normalize=True) returns the
  physically meaningful per-step average for all telemetry fields.
- Sensor now always applies area-normalisation internally; the normalize
  parameter only controls the additional time-normalisation step.
- run() gains zero_obs=True parameter (default) to control reset-on-step.
- 7 new integration tests covering accumulation, zeroing, and normalise.
- Fix bug in test_sensor_accuracy.py (undefined loop variable i).
- Bump version to 0.4.0 for the API change.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Frank14f
2026-06-21 00:50:20 +08:00
co-authored by Cursor
parent 7a609b2c76
commit 04c2bc75ea
10 changed files with 322 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CelerisLab/__init__.py
+1 -1
View File
@@ -14,7 +14,7 @@ Usage::
force = sim.read_force(0)
"""
__version__ = "0.3.0"
__version__ = "0.4.0"
from . import common, cuda, lbm, body, config
src/CelerisLab/body/manager.py
+69 -19
View File
@@ -89,6 +89,9 @@ class ObjectManager:
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) -------------------
@@ -471,6 +474,9 @@ class ObjectManager:
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:
@@ -524,20 +530,39 @@ class ObjectManager:
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``."""
"""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).
This always issues a ``memset`` on the sensor sub-range. Call it only when
the sensor kernel runs (e.g. ``field.n_sensor > 0``); the runner decides.
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
@@ -569,9 +594,14 @@ class ObjectManager:
"""Float index where the torque segment begins."""
return self.torque0_floats
def read_force(self, body_id: int) -> np.ndarray:
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)
@@ -579,21 +609,36 @@ class ObjectManager:
assert self.obs_pinned is not None
d = self.cfg.dim
i0 = body_id * d
return np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)
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) -> np.ndarray:
"""Return torque vector for ``body_id`` from ``obs_pinned``."""
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
return np.array(
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``.
By default this returns the area-averaged value over the sensor footprint.
Set ``normalize=False`` to get the raw sum accumulated by ``SensorKernel``.
**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
@@ -601,22 +646,27 @@ class ObjectManager:
d = self.cfg.dim
i0 = self.sensor0_floats + body_id * d
values = np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)
if not normalize:
return values
# Always area-normalise
count = int(self.sensor_cell_counts[body_id]) if body_id < self.sensor_cell_counts.size else 0
if count <= 0:
return values
return values / np.float32(count)
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) -> BodyTelemetry:
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)
torque = self.read_torque(body_id)
sensor = self.read_sensor(body_id, normalize=True)
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:
src/CelerisLab/simulation.py
+44 -14
View File
@@ -235,36 +235,58 @@ class Simulation:
self.bodies.set_force_state(body_id=id, fx=float(fx), fy=float(fy))
# -- Telemetry readback --------------------------------------------------
def read_force(self, id: int) -> np.ndarray:
"""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_force(self, id: int, *, normalize: bool = True) -> np.ndarray:
"""Return the force vector on body *id* from the pinned obs buffer.
def read_torque(self, id: int) -> 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 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_torque(id)
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, return area-averaged velocity. If False,
return the raw sum accumulated by the GPU SensorKernel.
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)
def read_body(self, id: int, *, stream: cuda.Stream | None = None):
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``.
@@ -276,7 +298,7 @@ class Simulation:
stream = self.stream
if stream is not None:
stream.synchronize()
return self.bodies.read_body(id)
return self.bodies.read_body(id, normalize=normalize)
def read_bodies(self, *, stream: cuda.Stream | None = None) -> np.ndarray:
"""Return all bodies' telemetry as a flat float32 array.
@@ -443,6 +465,7 @@ class Simulation:
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.
@@ -453,6 +476,8 @@ class Simulation:
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: If >0, save checkpoint every N steps.
"""
if not self._initialized:
@@ -469,8 +494,9 @@ class Simulation:
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)
# Zero obs segments before step group
if zero_obs:
self.bodies.zero_obs_async(stream)
self._assert_runtime_contracts()
if checkpoint_interval > 0:
@@ -494,6 +520,10 @@ class Simulation:
stream=stream,
)
# Accumulate step count for obs time-normalisation
if steps > 0:
self.bodies.increment_obs_steps(steps)
# Async download obs
if sync_obs:
self.bodies.download_obs_full_async(stream)