04c2bc75ea
- 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>
250 lines
9.1 KiB
Python
250 lines
9.1 KiB
Python
"""Unified action/obs flow — host-only set_body, auto transfer, stream API, DRL loop pattern.
|
|
|
|
Requires GPU."""
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
import pycuda.autoinit
|
|
|
|
from CelerisLab.simulation import Simulation
|
|
|
|
|
|
class TestUnifiedObs(unittest.TestCase):
|
|
"""Test unified action/obs flow."""
|
|
|
|
def test_set_body_then_run_read_body(self):
|
|
"""set_body (host-only), run, read_body returns finite force."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
sim.run(50)
|
|
|
|
# set_body should not trigger H2D (no error expected)
|
|
sim.set_body(0, omega=0.001)
|
|
|
|
# run will auto-upload action
|
|
sim.run(50)
|
|
data = sim.read_body(0)
|
|
self.assertTrue(np.all(np.isfinite(data.force)),
|
|
f"Force finite: {data.force}")
|
|
sim.close()
|
|
|
|
def test_skip_transfer(self):
|
|
"""run(upload_act=False, sync_obs=False) should not crash."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
sim.run(50, upload_act=False, sync_obs=False)
|
|
# After no-sync run, step count should still advance
|
|
self.assertEqual(sim.stepper.step_count, 50)
|
|
sim.close()
|
|
|
|
def test_external_stream(self):
|
|
"""Providing an external CUDA stream should not crash."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
s = cuda.Stream()
|
|
sim.run(50, stream=s)
|
|
force = sim.read_force(0)
|
|
self.assertTrue(np.all(np.isfinite(force)),
|
|
f"Force finite with external stream: {force}")
|
|
sim.close()
|
|
|
|
def test_read_body_before_run_returns_zeros(self):
|
|
"""read_body before any run() should return zero force (buffer is
|
|
initialized to zero during sync_to_gpu)."""
|
|
sim = Simulation(device_id=0)
|
|
sim.add_body("circle", center=(128, 128), radius=8)
|
|
sim.initialize()
|
|
force = sim.read_force(0)
|
|
np.testing.assert_array_equal(force, np.zeros(2, dtype=np.float32))
|
|
sim.close()
|
|
|
|
def test_drl_pattern(self):
|
|
"""DRL-style loop: run → read → set → run → read."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
for i in range(3):
|
|
sim.run(50)
|
|
data = sim.read_body(0)
|
|
self.assertTrue(np.all(np.isfinite(data.force)))
|
|
sim.set_body(0, omega=0.001 * i)
|
|
self.assertEqual(sim.stepper.step_count, 150)
|
|
sim.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|
|
class TestObsZeroingAndNormalize(unittest.TestCase):
|
|
"""Test obs zeroing, step accumulation, and time-normalize."""
|
|
|
|
def test_zero_obs_true_resets_force(self):
|
|
"""run(zero_obs=True) resets the step counter; force per-step
|
|
should be the same order of magnitude across blocks."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
# Warmup to reach a more developed flow state
|
|
sim.run(200, zero_obs=True)
|
|
|
|
sim.run(50, zero_obs=True)
|
|
self.assertEqual(sim.bodies._obs_accum_steps, 50,
|
|
"Step counter should be 50 after one run(zero_obs=True)")
|
|
|
|
sim.run(50, zero_obs=True)
|
|
self.assertEqual(sim.bodies._obs_accum_steps, 50,
|
|
"Step counter should be reset to 50 after zero_obs=True")
|
|
sim.close()
|
|
|
|
def test_zero_obs_false_accumulates_force(self):
|
|
"""run(zero_obs=False) twice → step counter accumulates."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=False)
|
|
self.assertEqual(sim.bodies._obs_accum_steps, 50)
|
|
|
|
sim.run(50, zero_obs=False)
|
|
self.assertEqual(sim.bodies._obs_accum_steps, 100,
|
|
"Step counter should accumulate across zero_obs=False calls")
|
|
|
|
# Normalized value = raw / accumulated steps
|
|
raw = sim.read_force(0, normalize=False)
|
|
avg = sim.read_force(0, normalize=True)
|
|
np.testing.assert_allclose(avg, raw / 100.0, rtol=1e-6)
|
|
sim.close()
|
|
|
|
def test_zero_obs_true_resets_sensor(self):
|
|
"""Sensor values should not spill across run() calls with zero_obs."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("sensor", center=(nx // 2, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=True)
|
|
s1 = sim.read_sensor(0, normalize=False)
|
|
|
|
sim.run(50, zero_obs=True)
|
|
s2 = sim.read_sensor(0, normalize=False)
|
|
|
|
# Each block should start fresh — magnitudes should be similar
|
|
mag1 = np.sqrt(np.sum(s1**2))
|
|
mag2 = np.sqrt(np.sum(s2**2))
|
|
self.assertGreater(mag1, 0.0)
|
|
self.assertGreater(mag2, 0.0)
|
|
sim.close()
|
|
|
|
def test_normalize_divides_by_steps(self):
|
|
"""read_force(normalize=True) should give per-step force."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=True)
|
|
raw = sim.read_force(0, normalize=False)
|
|
avg = sim.read_force(0, normalize=True)
|
|
|
|
# avg should be roughly raw / 50
|
|
expected = raw / np.float32(50)
|
|
np.testing.assert_allclose(avg, expected, rtol=1e-6)
|
|
sim.close()
|
|
|
|
def test_read_sensor_normalize_false(self):
|
|
"""read_sensor(normalize=False) returns area-averaged but not
|
|
time-averaged value."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("sensor", center=(nx // 2, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=True)
|
|
raw = sim.read_sensor(0, normalize=False)
|
|
tim_avg = sim.read_sensor(0, normalize=True)
|
|
|
|
# raw should be sensor sum/cell_count (area-average only)
|
|
# tim_avg should be raw / 50
|
|
expected = raw / np.float32(50)
|
|
np.testing.assert_allclose(tim_avg, expected, rtol=1e-6)
|
|
sim.close()
|
|
|
|
def test_read_body_normalize(self):
|
|
"""read_body(normalize=True) divides all fields by step count."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
|
|
sim.add_body("sensor", center=(nx // 2, ny // 2), radius=6)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=True)
|
|
data = sim.read_body(0, normalize=True)
|
|
raw_f = sim.read_force(0, normalize=False)
|
|
raw_t = sim.read_torque(0, normalize=False)
|
|
|
|
# Normalized values = raw / 50
|
|
np.testing.assert_allclose(data.force, raw_f / 50.0, rtol=1e-6)
|
|
np.testing.assert_allclose(data.torque, raw_t / 50.0, rtol=1e-6)
|
|
sim.close()
|
|
|
|
def test_normalize_returns_zero_before_run(self):
|
|
"""read(..., normalize=True) before any run() returns zeros."""
|
|
sim = Simulation(device_id=0)
|
|
sim.add_body("circle", center=(128, 128), radius=8)
|
|
sim.initialize()
|
|
|
|
force = sim.read_force(0, normalize=True)
|
|
torque = sim.read_torque(0, normalize=True)
|
|
sensor = sim.read_sensor(0, normalize=True)
|
|
np.testing.assert_array_equal(force, np.zeros(2, dtype=np.float32))
|
|
np.testing.assert_array_equal(torque, np.zeros(1, dtype=np.float32))
|
|
np.testing.assert_array_equal(sensor, np.zeros(2, dtype=np.float32))
|
|
sim.close()
|
|
|
|
def test_sensor_area_always_normalized(self):
|
|
"""Sensor always does area-normalisation internally.
|
|
normalize=False should NOT equal GPU raw (should be smaller by cell_count)."""
|
|
sim = Simulation(device_id=0)
|
|
nx = sim.lbm_cfg.nx
|
|
ny = sim.lbm_cfg.ny
|
|
sim.add_body("sensor", center=(nx // 2, ny // 2), radius=8)
|
|
sim.initialize()
|
|
|
|
sim.run(50, zero_obs=True)
|
|
|
|
# read_sensor with normalize=False returns area-averaged value.
|
|
# This value should be non-zero if there's flow.
|
|
sensor_val = sim.read_sensor(0, normalize=False)
|
|
self.assertTrue(np.all(np.isfinite(sensor_val)),
|
|
f"Sensor should be finite: {sensor_val}")
|
|
|
|
# Area-only normalization: if cell_count > 1, the value should be less
|
|
# than the raw GPU accumulator magnitude in most cases.
|
|
cells_arr, _ = sim.bodies.get(0).get_sensor_list(nx, ny)
|
|
n_cells = len(cells_arr)
|
|
self.assertGreater(n_cells, 0)
|
|
sim.close()
|