CelerisLab/tests/integration/test_unified_obs.py

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