"""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()