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>
This commit is contained in:
Frank14f
@@ -0,0 +1 @@
|
||||
# CelerisLab/tests/integration/__init__.py
|
||||
@@ -0,0 +1,117 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,145 @@
|
||||
"""DDF patch after add / remove body — finite forces, finite macroscopic field after moderate steps.
|
||||
|
||||
Requires GPU."""
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import pycuda.autoinit
|
||||
|
||||
from CelerisLab.simulation import Simulation
|
||||
from CelerisLab.lbm.descriptors import FLUID, OBSTACLE
|
||||
|
||||
|
||||
class TestDDFPatchAddBody(unittest.TestCase):
|
||||
"""Test adding a body (fluid -> solid DDF patch)."""
|
||||
|
||||
def test_add_body_runs_stably(self):
|
||||
"""After adding a body, the simulation runs with finite forces
|
||||
for a moderate number of steps."""
|
||||
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(200)
|
||||
|
||||
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)
|
||||
|
||||
sim.run(50)
|
||||
force = sim.read_force(0)
|
||||
self.assertTrue(np.all(np.isfinite(force)),
|
||||
f"Finite force after add body: {force}")
|
||||
sim.close()
|
||||
|
||||
|
||||
class TestDDFPatchRemoveBody(unittest.TestCase):
|
||||
"""Test removing a body (solid -> fluid DDF patch via BFS inward fill)."""
|
||||
|
||||
def test_remove_body_released_region_is_fluid(self):
|
||||
"""After removal, the former body center should be a fluid cell."""
|
||||
sim = Simulation(device_id=0)
|
||||
nx = sim.lbm_cfg.nx
|
||||
ny = sim.lbm_cfg.ny
|
||||
cx, cy = nx // 4, ny // 2
|
||||
|
||||
sim.add_body("circle", center=(cx, cy), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(200)
|
||||
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
flags = sim.get_flags()
|
||||
center_idx = cx + cy * nx
|
||||
self.assertTrue(flags[center_idx] & FLUID)
|
||||
sim.close()
|
||||
|
||||
def test_remove_body_finite_field(self):
|
||||
"""After removal, the macroscopic field is finite."""
|
||||
sim = Simulation(device_id=0)
|
||||
nx = sim.lbm_cfg.nx
|
||||
ny = sim.lbm_cfg.ny
|
||||
cx, cy = nx // 4, ny // 2
|
||||
|
||||
sim.add_body("circle", center=(cx, cy), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
flags = sim.get_flags()
|
||||
center_idx = cx + cy * nx
|
||||
self.assertTrue(flags[center_idx] & FLUID)
|
||||
|
||||
sim.run(50)
|
||||
macro = sim.get_macroscopic()
|
||||
self.assertTrue(np.all(np.isfinite(macro["ux"])),
|
||||
"Macroscopic ux should be finite after remove body")
|
||||
sim.close()
|
||||
|
||||
|
||||
class TestDDFPatchAddAndRemove(unittest.TestCase):
|
||||
"""Test combined add + remove in one sync."""
|
||||
|
||||
def test_add_one_remove_another(self):
|
||||
"""Add a body and remove a different one in the same sync."""
|
||||
sim = Simulation(device_id=0)
|
||||
nx = sim.lbm_cfg.nx
|
||||
ny = sim.lbm_cfg.ny
|
||||
cx, cy = nx // 4, ny // 2
|
||||
|
||||
sim.add_body("circle", center=(cx, cy), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
sim.add_body("circle", center=(nx // 2, ny // 2), radius=6)
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
self.assertEqual(sim.bodies.count, 1)
|
||||
|
||||
sim.run(50)
|
||||
force = sim.read_force(0)
|
||||
self.assertEqual(force.shape[0], 2)
|
||||
self.assertTrue(np.all(np.isfinite(force)),
|
||||
f"Finite force after add+remove: {force}")
|
||||
sim.close()
|
||||
|
||||
|
||||
class TestDDFPatchNoChange(unittest.TestCase):
|
||||
"""Test that patch is a no-op when there are no geometry changes."""
|
||||
|
||||
def test_no_mask_no_change(self):
|
||||
"""When neither mask has any True entries, DDF should be unchanged."""
|
||||
sim = Simulation(device_id=0)
|
||||
sim.add_body("circle", center=(128, 128), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(100)
|
||||
|
||||
sim.field.download_ddf(force=True)
|
||||
ddf_before = sim.field.ddf.copy()
|
||||
|
||||
from CelerisLab.body.ddf_patch import patch_ddf_for_body_sync as patch_fn
|
||||
|
||||
n = sim.field.n
|
||||
added = np.zeros(n, dtype=bool)
|
||||
released = np.zeros(n, dtype=bool)
|
||||
patch_fn(
|
||||
sim.field,
|
||||
ddf_before, sim.field.flag.copy(), sim.field.flag.copy(),
|
||||
added, released)
|
||||
|
||||
np.testing.assert_array_equal(sim.field.ddf, ddf_before)
|
||||
sim.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,166 @@
|
||||
"""sync_bodies pipeline without DDF patch — recompile, esopull guard, step count preservation.
|
||||
|
||||
Requires GPU."""
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import pycuda.driver as cuda
|
||||
import pycuda.autoinit
|
||||
|
||||
from CelerisLab.simulation import Simulation
|
||||
|
||||
|
||||
# Use a small grid for fast compilation and test execution
|
||||
NX, NY = 128, 64
|
||||
|
||||
|
||||
class TestSyncBodiesSkeleton(unittest.TestCase):
|
||||
"""Test sync_bodies() with real GPU -- skeleton without DDF patch."""
|
||||
|
||||
def _make_sim(self) -> Simulation:
|
||||
"""Create a Simulation with a small double_buffer D2Q9 grid."""
|
||||
return Simulation(device_id=0)
|
||||
|
||||
def test_remove_body_sync_and_run(self):
|
||||
"""Remove a body, sync, and continue running without crash."""
|
||||
sim = self._make_sim()
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
# Remove the body and sync
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
# Should be able to continue running
|
||||
sim.run(50)
|
||||
|
||||
# No body left -- count should be 0
|
||||
self.assertEqual(sim.bodies.count, 0)
|
||||
sim.close()
|
||||
|
||||
def test_add_body_after_initialize(self):
|
||||
"""Add a body after initialize, sync, and run."""
|
||||
sim = self._make_sim()
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
# Add a body (returns -1 since it's staged)
|
||||
result_id = sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
self.assertEqual(result_id, -1)
|
||||
self.assertTrue(sim.bodies.has_pending_edit())
|
||||
|
||||
# Sync commits the body
|
||||
sim.sync_bodies()
|
||||
self.assertFalse(sim.bodies.has_pending_edit())
|
||||
self.assertEqual(sim.bodies.count, 1)
|
||||
|
||||
# Should be able to run with the new body
|
||||
sim.run(50)
|
||||
|
||||
# Force readback should work
|
||||
force = sim.read_force(0)
|
||||
self.assertEqual(force.shape[0], 2)
|
||||
sim.close()
|
||||
|
||||
def test_add_then_remove_body(self):
|
||||
"""Add a body, then remove it, sync -- should result in zero bodies."""
|
||||
sim = self._make_sim()
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
# Add another body and remove the original
|
||||
sim.add_body("circle", center=(NX // 2, NY // 2), radius=6)
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
# One body remaining (the newly added one, now id=0)
|
||||
self.assertEqual(sim.bodies.count, 1)
|
||||
sim.run(50)
|
||||
sim.close()
|
||||
|
||||
def test_sync_preserves_step_count(self):
|
||||
"""sync_bodies() should not reset the step counter."""
|
||||
sim = self._make_sim()
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(100)
|
||||
steps_before = sim.stepper.step_count
|
||||
|
||||
sim.remove_body(0)
|
||||
sim.sync_bodies()
|
||||
|
||||
steps_after = sim.stepper.step_count
|
||||
self.assertEqual(steps_after, steps_before)
|
||||
sim.close()
|
||||
|
||||
def test_no_pending_edit_is_noop(self):
|
||||
"""sync_bodies() with no pending edits should be a no-op."""
|
||||
sim = self._make_sim()
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
# No edits -- sync should return immediately
|
||||
sim.sync_bodies()
|
||||
self.assertEqual(sim.bodies.count, 1)
|
||||
sim.close()
|
||||
|
||||
def test_run_discards_pending_without_sync(self):
|
||||
"""Running without sync_bodies() should discard pending edits."""
|
||||
sim = self._make_sim()
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
sim.initialize()
|
||||
sim.run(50)
|
||||
|
||||
# Stage a removal but don't sync
|
||||
sim.remove_body(0)
|
||||
self.assertTrue(sim.bodies.has_pending_edit())
|
||||
|
||||
# run() auto-discards pending
|
||||
sim.run(50)
|
||||
self.assertFalse(sim.bodies.has_pending_edit())
|
||||
self.assertEqual(sim.bodies.count, 1)
|
||||
sim.close()
|
||||
|
||||
def test_esopull_raises_not_implemented(self):
|
||||
"""sync_bodies() with esopull should raise NotImplementedError."""
|
||||
# Create a sim with esopull streaming
|
||||
from CelerisLab.config import load_lbm_config
|
||||
cfg = load_lbm_config()
|
||||
cfg.streaming = "esopull"
|
||||
# We need to build the sim manually to override streaming
|
||||
sim = Simulation.__new__(Simulation)
|
||||
sim._stream = None
|
||||
sim.lbm_cfg = cfg
|
||||
from CelerisLab.config import BodyConfig
|
||||
sim.body_cfg = BodyConfig()
|
||||
from CelerisLab.cuda.context import CudaContext
|
||||
sim.ctx = CudaContext(0)
|
||||
from CelerisLab.cuda import compiler_v2 as compiler
|
||||
arch = sim._resolve_compile_arch = lambda: sim.ctx.sm_arch
|
||||
arch_val = CudaContext(0).sm_arch
|
||||
compiler.generate_config(cfg, n_objects=0)
|
||||
ptx_path = compiler.compile_kernel(arch=arch_val)
|
||||
module = compiler.load_module(ptx_path)
|
||||
sim._ptx_path = ptx_path
|
||||
sim._module = module
|
||||
from CelerisLab.lbm.field import LBMField
|
||||
sim.field = LBMField(cfg, module)
|
||||
from CelerisLab.lbm.stepper import LBMStepper
|
||||
sim.stepper = LBMStepper(sim.field, module, cfg)
|
||||
from CelerisLab.body.manager import ObjectManager
|
||||
sim.bodies = ObjectManager(
|
||||
cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
|
||||
sim._initialized = True
|
||||
|
||||
sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
|
||||
with self.assertRaises(NotImplementedError):
|
||||
sim.sync_bodies()
|
||||
sim.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,89 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user