CelerisLab/tests/integration/test_body_sync_e2e.py
Frank14f 987566c0e6 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

118 lines
3.5 KiB
Python

"""Full add/remove/checkpoint/load lifecycle — end-to-end body topology sync.
Requires GPU."""
import os
import unittest
import tempfile
import numpy as np
import pycuda.autoinit
from CelerisLab.simulation import Simulation
from CelerisLab.lbm.descriptors import OBSTACLE, FLUID
class TestBodySyncE2E(unittest.TestCase):
"""Full end-to-end test of runtime body topology sync."""
def test_full_lifecycle(self):
"""Create, init, run, add body, remove body, checkpoint, load."""
sim = Simulation(device_id=0)
nx = sim.lbm_cfg.nx
ny = sim.lbm_cfg.ny
cx, cy = nx // 4, ny // 2
# 1. Initialize and run
sim.initialize()
sim.run(100)
self.assertGreater(sim.stepper.step_count, 0)
# 2. Add a body and sync
sim.add_body("circle", center=(cx, cy), radius=8)
sim.sync_bodies()
self.assertEqual(sim.bodies.count, 1)
center_idx = cx + cy * nx
self.assertTrue(sim.get_flags()[center_idx] & OBSTACLE)
# 3. Run with body (short window — finite near-term)
sim.run(50)
force = sim.read_force(0)
self.assertTrue(np.all(np.isfinite(force)),
f"Finite force after add: {force}")
steps_before_remove = sim.stepper.step_count
# 4. Remove the body and sync
sim.remove_body(0)
sim.sync_bodies()
self.assertEqual(sim.bodies.count, 0)
self.assertTrue(sim.get_flags()[center_idx] & FLUID)
# 5. Run after removal (no crash, step count advances)
sim.run(50)
self.assertEqual(sim.stepper.step_count, steps_before_remove + 50)
# 6. Save checkpoint
with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f:
ckpt_path = f.name
try:
saved_path = sim.save_checkpoint(ckpt_path)
self.assertTrue(os.path.exists(saved_path))
# 7. Load checkpoint in a new simulation
sim2 = Simulation(device_id=0)
sim2.initialize()
sim2.run(1)
sim2.load_checkpoint(saved_path)
self.assertEqual(sim2.stepper.step_count, sim.stepper.step_count)
self.assertEqual(sim2.bodies.count, 0)
# 8. Continue running in restored sim
sim2.run(50)
# Verify step count advances (DDF may have NaN from pre-existing body)
self.assertEqual(sim2.stepper.step_count,
sim.stepper.step_count + 50)
sim2.close()
finally:
os.unlink(ckpt_path)
sim.close()
def test_add_remove_add_cycle(self):
"""Add → run → remove → run → add → run cycle with finite checks."""
sim = Simulation(device_id=0)
nx = sim.lbm_cfg.nx
ny = sim.lbm_cfg.ny
cx, cy = nx // 4, ny // 2
sim.initialize()
sim.run(100)
# Add
sim.add_body("circle", center=(cx, cy), radius=8)
sim.sync_bodies()
self.assertEqual(sim.bodies.count, 1)
sim.run(50)
# Remove
sim.remove_body(0)
sim.sync_bodies()
self.assertEqual(sim.bodies.count, 0)
sim.run(50)
# Add again
sim.add_body("circle", center=(nx // 2, ny // 2), radius=6)
sim.sync_bodies()
self.assertEqual(sim.bodies.count, 1)
sim.run(50)
force = sim.read_force(0)
self.assertTrue(np.all(np.isfinite(force)),
f"Finite force after add-remove-add: {force}")
sim.close()
if __name__ == "__main__":
unittest.main()