feat(body): runtime body add/remove, unified action/obs, FRC_REGION flag

- Add runtime body topology sync (add_body/remove_body + sync_bodies) with recompile, DDF patch (feq + BFS inward fill), and commit. - Unify action/obs flow: set_body/set_force are now host-only; run() auto-uploads action and downloads obs via CUDA stream. - Add read_body(id) -> BodyTelemetry and read_bodies() for DRL loops. - Add FRC_REGION flag (0x0800) for force_region cells. - Extract equilibrium helpers (lbm/equilibrium.py) and DDF patch module (body/ddf_patch.py). - Merge recompile / _runtime_recompile into single _recompile(). - Add n_objects to checkpoint; validate on load. - Add test suite: 40 unit + 19 integration tests (59 total). - Add conftest.py and docs/tests_overview.md for test documentation. - Update README.md and CONFIG.md for new API. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-20 18:17:07 +08:00
parent d5b7e98750
commit 987566c0e6
28 changed files with 2112 additions and 86 deletions
@@ -0,0 +1 @@
+# CelerisLab/tests/unit/__init__.py
@@ -0,0 +1,94 @@
+# CelerisLab/tests/unit/test_body_flags.py
+"""Body type flag masks — OBSTACLE, SENSOR_FLAG, FRC_REGION bits for circle / sensor / force_region.
+
+No GPU required."""
+
+import unittest
+
+import numpy as np
+
+from CelerisLab.body.manager import ObjectManager
+from CelerisLab.body.objects import SimObject
+from CelerisLab.body.geometry.circle import CircleGeometry
+from CelerisLab.lbm.descriptors import (
+    FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION,
+)
+
+
+def _make_obj(cx: float, cy: float, radius: float,
+              is_sensor: bool = False,
+              is_force_region: bool = False) -> SimObject:
+    geom = CircleGeometry(cx, cy, radius)
+    return SimObject(obj_id=-1, geometry=geom,
+                     center=(cx, cy), radius=radius,
+                     is_sensor=is_sensor,
+                     is_force_region=is_force_region)
+
+
+NX, NY = 64, 32
+
+
+class TestBodyFlags(unittest.TestCase):
+    """Verify flag mask bits for each body type."""
+
+    def _obj_flag_mask(self, obj: SimObject) -> np.ndarray:
+        return obj.get_flag_mask(NX, NY)
+
+    def test_circle_has_obstacle_solid_curved(self):
+        mask = self._obj_flag_mask(_make_obj(32, 16, 5))
+        center = 32 + 16 * NX
+        self.assertTrue(mask[center] & OBSTACLE,
+                        "Circle should have OBSTACLE bit")
+        self.assertTrue(mask[center] & SOLID,
+                        "Circle should have SOLID bit")
+        self.assertTrue(mask[center] & BC_CURVED,
+                        "Circle should have BC_CURVED bit")
+        self.assertFalse(mask[center] & FLUID,
+                         "Circle interior should NOT be FLUID")
+
+    def test_sensor_has_sensor_flag(self):
+        mask = self._obj_flag_mask(_make_obj(32, 16, 5, is_sensor=True))
+        center = 32 + 16 * NX
+        self.assertTrue(mask[center] & SENSOR_FLAG,
+                        "Sensor should have SENSOR_FLAG bit")
+        self.assertTrue(mask[center] & FLUID,
+                        "Sensor should be FLUID")
+
+    def test_force_region_has_frc_region_flag(self):
+        mask = self._obj_flag_mask(
+            _make_obj(32, 16, 5, is_force_region=True))
+        center = 32 + 16 * NX
+        self.assertTrue(mask[center] & FRC_REGION,
+                        "Force region should have FRC_REGION bit")
+        self.assertTrue(mask[center] & FLUID,
+                        "Force region should be FLUID")
+
+    def test_circle_has_no_sensor_or_frc_flag(self):
+        mask = self._obj_flag_mask(_make_obj(32, 16, 5))
+        center = 32 + 16 * NX
+        self.assertFalse(mask[center] & SENSOR_FLAG,
+                         "Circle should NOT have SENSOR_FLAG")
+        self.assertFalse(mask[center] & FRC_REGION,
+                         "Circle should NOT have FRC_REGION")
+
+    def test_force_region_has_no_obstacle(self):
+        mask = self._obj_flag_mask(
+            _make_obj(32, 16, 5, is_force_region=True))
+        center = 32 + 16 * NX
+        self.assertFalse(mask[center] & OBSTACLE,
+                         "Force region should NOT have OBSTACLE bit")
+
+    def test_all_body_types_nonzero_masks(self):
+        for obj in [
+            _make_obj(32, 16, 5),
+            _make_obj(32, 16, 5, is_sensor=True),
+            _make_obj(32, 16, 5, is_force_region=True),
+        ]:
+            mask = self._obj_flag_mask(obj)
+            self.assertGreater(np.count_nonzero(mask), 0,
+                               f"{obj.is_sensor=},{obj.is_force_region=}: "
+                               "mask should have non-zero entries")
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,50 @@
+"""D2Q9 equilibrium helpers — compute_feq_d2q9 and compute_macro_from_ddf correctness.
+
+No GPU required."""
+
+import unittest
+
+import numpy as np
+
+from CelerisLab.lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf
+
+
+class TestEquilibrium(unittest.TestCase):
+    """Verify D2Q9 equilibrium and macroscopic helpers."""
+
+    def test_feq_at_rest(self):
+        """Equilibrium at rho=1.0, u=0 should give w_i (weights)."""
+        feq = compute_feq_d2q9(1.0, 0.0, 0.0)
+        w = np.array([4/9, 1/9, 1/9, 1/9, 1/9,
+                      1/36, 1/36, 1/36, 1/36], dtype=np.float32)
+        np.testing.assert_allclose(feq, w, rtol=1e-6)
+
+    def test_feq_sums_to_rho(self):
+        """Sum of feq should equal rho."""
+        rho, ux, uy = 1.2, 0.05, -0.02
+        feq = compute_feq_d2q9(rho, ux, uy)
+        self.assertAlmostEqual(float(np.sum(feq)), rho, places=6)
+
+    def test_macro_preserves_ux_uy(self):
+        """compute_macro_from_ddf(feq) should recover rho, ux, uy."""
+        rho, ux, uy = 1.0, 0.1, 0.0
+        feq = compute_feq_d2q9(rho, ux, uy)
+        rho_out, ux_out, uy_out = compute_macro_from_ddf(feq)
+        self.assertAlmostEqual(rho_out, rho, places=6)
+        self.assertAlmostEqual(ux_out, ux, places=6)
+        self.assertAlmostEqual(uy_out, uy, places=6)
+
+    def test_feq_nonzero_vel(self):
+        """Equilibrium at non-zero velocity should be asymmetric."""
+        feq_x = compute_feq_d2q9(1.0, 0.1, 0.0)
+        feq_y = compute_feq_d2q9(1.0, 0.0, 0.1)
+        # x-directed flow should have f1 > f2 (right > left)
+        self.assertGreater(feq_x[1], feq_x[2],
+                           "Right-moving f1 should exceed left-moving f2")
+        # y-directed flow should have f3 > f4 (up > down)
+        self.assertGreater(feq_y[3], feq_y[4],
+                           "Up-moving f3 should exceed down-moving f4")
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,147 @@
+"""ObjectManager pending edit lifecycle — stage_add, stage_remove, has_pending_edit, clear_pending_edits.
+
+No GPU required."""
+
+import unittest
+
+from CelerisLab.body.manager import ObjectManager
+from CelerisLab.body.objects import SimObject
+from CelerisLab.body.geometry.circle import CircleGeometry
+
+
+def _make_circle_obj(cx: float, cy: float, radius: float,
+                     is_sensor: bool = False) -> SimObject:
+    """Create a minimal SimObject with CircleGeometry."""
+    geom = CircleGeometry(cx, cy, radius)
+    return SimObject(
+        obj_id=-1,
+        geometry=geom,
+        center=(cx, cy),
+        radius=radius,
+        is_sensor=is_sensor,
+    )
+
+
+class TestPendingEditLifecycle(unittest.TestCase):
+    """Test the pending edit state machine on ObjectManager."""
+
+    def setUp(self):
+        # ObjectManager requires nx, ny, nz, nq, cfg.  Use a minimal cfg stub.
+        self.cfg = _StubCfg(dim=2)
+        self.mgr = ObjectManager(nx=64, ny=32, nz=1, nq=9, cfg=self.cfg)
+
+    # -- stage_add -----------------------------------------------------------
+
+    def test_stage_add_sets_edit_active(self):
+        obj = _make_circle_obj(30, 16, 5)
+        self.assertFalse(self.mgr.has_pending_edit())
+        self.mgr.stage_add(obj)
+        self.assertTrue(self.mgr.has_pending_edit())
+
+    def test_stage_add_does_not_change_formal_count(self):
+        obj = _make_circle_obj(30, 16, 5)
+        self.assertEqual(self.mgr.count, 0)
+        self.mgr.stage_add(obj)
+        self.assertEqual(self.mgr.count, 0)
+
+    def test_stage_add_multiple(self):
+        for i in range(3):
+            self.mgr.stage_add(_make_circle_obj(10 + i * 10, 16, 3))
+        self.assertTrue(self.mgr.has_pending_edit())
+
+    # -- stage_remove --------------------------------------------------------
+
+    def test_stage_remove_valid_id(self):
+        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.mgr.stage_remove(body_id)
+        self.assertTrue(self.mgr.has_pending_edit())
+
+    def test_stage_remove_invalid_id_raises(self):
+        with self.assertRaises(IndexError):
+            self.mgr.stage_remove(999)
+
+    def test_stage_remove_does_not_change_formal_count(self):
+        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.assertEqual(self.mgr.count, 1)
+        self.mgr.stage_remove(body_id)
+        self.assertEqual(self.mgr.count, 1)
+
+    # -- has_pending_edit ----------------------------------------------------
+
+    def test_has_pending_edit_false_initially(self):
+        self.assertFalse(self.mgr.has_pending_edit())
+
+    def test_has_pending_edit_false_after_clear(self):
+        self.mgr.stage_add(_make_circle_obj(30, 16, 5))
+        self.mgr.clear_pending_edits()
+        self.assertFalse(self.mgr.has_pending_edit())
+
+    def test_has_pending_edit_false_after_empty_stage(self):
+        # If we add and then remove the same pending add, edit is still
+        # "active" (edit_active=True) but has no pending content.
+        # has_pending_edit should return False.
+        obj = _make_circle_obj(30, 16, 5)
+        self.mgr.stage_add(obj)
+        # Manually clear pending_add to simulate an empty edit window
+        self.mgr._pending_add.clear()
+        self.assertFalse(self.mgr.has_pending_edit())
+
+    # -- clear_pending_edits -------------------------------------------------
+
+    def test_clear_resets_all_pending(self):
+        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
+        self.mgr.stage_remove(body_id)
+        self.assertTrue(self.mgr.has_pending_edit())
+
+        self.mgr.clear_pending_edits()
+        self.assertFalse(self.mgr.has_pending_edit())
+        self.assertEqual(len(self.mgr._pending_add), 0)
+        self.assertEqual(len(self.mgr._pending_remove), 0)
+        self.assertFalse(self.mgr._edit_active)
+
+    def test_clear_preserves_formal_registry(self):
+        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.mgr.stage_remove(body_id)
+        self.mgr.clear_pending_edits()
+        # Formal object should still be there
+        self.assertEqual(self.mgr.count, 1)
+        self.assertEqual(self.mgr.get(body_id).obj_id, body_id)
+
+    # -- Combination: add + remove -------------------------------------------
+
+    def test_stage_add_and_remove_together(self):
+        body_id = self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
+        self.mgr.stage_remove(body_id)
+        self.assertTrue(self.mgr.has_pending_edit())
+        # Formal count unchanged
+        self.assertEqual(self.mgr.count, 1)
+
+    # -- Formal add (pre-initialize path) ------------------------------------
+
+    def test_formal_add_still_works(self):
+        """The existing add() path must remain functional."""
+        obj = _make_circle_obj(30, 16, 5)
+        body_id = self.mgr.add(obj)
+        self.assertEqual(body_id, 0)
+        self.assertEqual(self.mgr.count, 1)
+
+    def test_formal_add_and_pending_coexist(self):
+        """Formal add + pending stage should not interfere."""
+        self.mgr.add(_make_circle_obj(30, 16, 5))
+        self.mgr.stage_add(_make_circle_obj(40, 16, 3))
+        self.assertEqual(self.mgr.count, 1)
+        self.assertTrue(self.mgr.has_pending_edit())
+
+
+class _StubCfg:
+    """Minimal LBMConfig-like stub for ObjectManager construction."""
+
+    def __init__(self, dim: int = 2):
+        self.dim = dim
+        self.is_d3q19 = (dim == 3)
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,234 @@
+"""BodySyncPlan construction — build_flags_for, build_compact_lists_for, build_next_objects, build_sync_plan, commit_pending.
+
+No GPU required."""
+
+import unittest
+
+import numpy as np
+
+from CelerisLab.body.manager import ObjectManager
+from CelerisLab.body.objects import SimObject
+from CelerisLab.body.geometry.circle import CircleGeometry
+from CelerisLab.body.sync_plan import BodySyncPlan
+from CelerisLab.lbm.descriptors import FLUID, SOLID, OBSTACLE, BC_CURVED
+
+
+def _make_circle_obj(cx: float, cy: float, radius: float,
+                     is_sensor: bool = False) -> SimObject:
+    geom = CircleGeometry(cx, cy, radius)
+    return SimObject(
+        obj_id=-1,
+        geometry=geom,
+        center=(cx, cy),
+        radius=radius,
+        is_sensor=is_sensor,
+    )
+
+
+class _StubCfg:
+    def __init__(self, dim=2):
+        self.dim = dim
+        self.is_d3q19 = (dim == 3)
+
+
+class _StubField:
+    """Minimal LBMField stub for build_sync_plan testing."""
+
+    def __init__(self, nx: int, ny: int):
+        self.nx = nx
+        self.ny = ny
+        # Build a simple channel flag array (fluid everywhere except top/bottom walls).
+        n = nx * ny
+        self.flag = np.ones(n, dtype=np.uint16) * FLUID
+        self.flag[:nx] = SOLID | 0x0010  # bottom wall
+        self.flag[(ny - 1) * nx:ny * nx] = SOLID | 0x0010  # top wall
+        # Save a clean copy for build_channel_flags
+        self._channel_flags = self.flag.copy()
+
+    def build_channel_flags(self) -> np.ndarray:
+        """Return a clean channel base (no object overlays)."""
+        return self._channel_flags.copy()
+
+
+NX, NY = 64, 32
+
+
+class TestBuildFlagsFor(unittest.TestCase):
+
+    def setUp(self):
+        self.cfg = _StubCfg()
+        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
+        self.field = _StubField(NX, NY)
+
+    def test_empty_objects_returns_base(self):
+        base = self.field.build_channel_flags()
+        result = ObjectManager.build_flags_for(
+            [], base, nx=NX, ny=NY, nz=1)
+        np.testing.assert_array_equal(result, base)
+
+    def test_one_circle_produces_solid_region(self):
+        base = self.field.build_channel_flags()
+        obj = _make_circle_obj(32, 16, 5)
+        result = ObjectManager.build_flags_for(
+            [obj], base, nx=NX, ny=NY, nz=1)
+        # Center cell should be solid with OBSTACLE and BC_CURVED bits
+        center_idx = 32 + 16 * NX
+        self.assertTrue(result[center_idx] & SOLID)
+        self.assertTrue(result[center_idx] & OBSTACLE)
+
+    def test_instance_method_delegates(self):
+        base = self.field.build_channel_flags()
+        obj = _make_circle_obj(32, 16, 5)
+        self.mgr.add(obj)
+        result = self.mgr.build_flags(base)
+        center_idx = 32 + 16 * NX
+        self.assertTrue(result[center_idx] & OBSTACLE)
+
+
+class TestBuildCompactListsFor(unittest.TestCase):
+
+    def setUp(self):
+        self.cfg = _StubCfg()
+        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
+
+    def test_circle_produces_curved_links(self):
+        obj = _make_circle_obj(32, 16, 5)
+        obj.obj_id = 0
+        result = self.mgr.build_compact_lists_for([obj])
+        cl_fluid_idx = result[0]
+        self.assertGreater(len(cl_fluid_idx), 0)
+
+    def test_sensor_produces_sensor_cells(self):
+        obj = _make_circle_obj(32, 16, 5, is_sensor=True)
+        obj.obj_id = 0
+        result = self.mgr.build_compact_lists_for([obj])
+        sensor_cells = result[8]
+        self.assertGreater(len(sensor_cells), 0)
+
+    def test_instance_method_delegates(self):
+        obj = _make_circle_obj(32, 16, 5)
+        self.mgr.add(obj)
+        result = self.mgr.build_compact_lists()
+        cl_fluid_idx = result[0]
+        self.assertGreater(len(cl_fluid_idx), 0)
+
+
+class TestBuildNextObjects(unittest.TestCase):
+
+    def setUp(self):
+        self.cfg = _StubCfg()
+        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
+
+    def test_no_pending_returns_formal_objects(self):
+        obj = _make_circle_obj(32, 16, 5)
+        self.mgr.add(obj)
+        result = self.mgr.build_next_objects()
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].obj_id, 0)
+
+    def test_removal_excludes_object(self):
+        obj0 = _make_circle_obj(20, 16, 3)
+        obj1 = _make_circle_obj(40, 16, 3)
+        id0 = self.mgr.add(obj0)
+        self.mgr.add(obj1)
+        self.mgr.stage_remove(id0)
+        result = self.mgr.build_next_objects()
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].obj_id, 0)
+        # The remaining object should be the second one (center at 40)
+        self.assertAlmostEqual(result[0].center[0], 40.0)
+
+    def test_add_appends_new_object(self):
+        obj0 = _make_circle_obj(20, 16, 3)
+        self.mgr.add(obj0)
+        new_obj = _make_circle_obj(40, 16, 3)
+        self.mgr.stage_add(new_obj)
+        result = self.mgr.build_next_objects()
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0].obj_id, 0)
+        self.assertEqual(result[1].obj_id, 1)
+
+    def test_ids_are_consecutive(self):
+        obj0 = _make_circle_obj(20, 16, 3)
+        obj1 = _make_circle_obj(30, 16, 3)
+        obj2 = _make_circle_obj(40, 16, 3)
+        id0 = self.mgr.add(obj0)
+        self.mgr.add(obj1)
+        self.mgr.add(obj2)
+        # Remove middle object
+        self.mgr.stage_remove(id0 + 1)
+        result = self.mgr.build_next_objects()
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0].obj_id, 0)
+        self.assertEqual(result[1].obj_id, 1)
+
+    def test_formal_registry_unchanged(self):
+        obj = _make_circle_obj(32, 16, 5)
+        id0 = self.mgr.add(obj)
+        self.mgr.stage_remove(id0)
+        self.mgr.build_next_objects()
+        # Formal registry should be untouched
+        self.assertEqual(self.mgr.count, 1)
+
+
+class TestBuildSyncPlan(unittest.TestCase):
+
+    def setUp(self):
+        self.cfg = _StubCfg()
+        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
+        self.field = _StubField(NX, NY)
+
+    def test_add_body_produces_added_solid_mask(self):
+        new_obj = _make_circle_obj(32, 16, 5)
+        self.mgr.stage_add(new_obj)
+        plan = self.mgr.build_sync_plan(self.field)
+        self.assertIsInstance(plan, BodySyncPlan)
+        self.assertEqual(plan.next_count, 1)
+        # Center should be in added_solid_mask (was fluid, becomes solid)
+        center_idx = 32 + 16 * NX
+        self.assertTrue(plan.added_solid_mask[center_idx])
+
+    def test_remove_body_produces_released_fluid_mask(self):
+        obj = _make_circle_obj(32, 16, 5)
+        id0 = self.mgr.add(obj)
+        # Build current flags so the field "knows" about this body
+        base = self.field.build_channel_flags()
+        self.field.flag = self.mgr.build_flags(base)
+        self.mgr.stage_remove(id0)
+        plan = self.mgr.build_sync_plan(self.field)
+        center_idx = 32 + 16 * NX
+        self.assertTrue(plan.released_fluid_mask[center_idx])
+
+    def test_no_change_masks_are_empty(self):
+        plan = self.mgr.build_sync_plan(self.field)
+        self.assertFalse(np.any(plan.added_solid_mask))
+        self.assertFalse(np.any(plan.released_fluid_mask))
+
+
+class TestCommitPending(unittest.TestCase):
+
+    def setUp(self):
+        self.cfg = _StubCfg()
+        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
+
+    def test_commit_replaces_registry(self):
+        obj0 = _make_circle_obj(20, 16, 3)
+        obj1 = _make_circle_obj(40, 16, 3)
+        id0 = self.mgr.add(obj0)
+        self.mgr.add(obj1)
+        self.mgr.stage_remove(id0)
+        next_objs = self.mgr.build_next_objects()
+        self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32))
+        self.assertEqual(self.mgr.count, 1)
+        self.assertEqual(self.mgr.get(0).obj_id, 0)
+        self.assertFalse(self.mgr.has_pending_edit())
+
+    def test_commit_clears_pending(self):
+        self.mgr.stage_add(_make_circle_obj(32, 16, 5))
+        next_objs = self.mgr.build_next_objects()
+        self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32))
+        self.assertFalse(self.mgr.has_pending_edit())
+
+
+if __name__ == "__main__":
+    unittest.main()