"""sync_bodies pipeline — recompile, esopull runtime body sync, step count preservation. Requires GPU.""" import unittest import tempfile import os import json import numpy as np import pycuda.driver as cuda import pycuda.autoinit from CelerisLab.simulation import Simulation, _scatter_to_field, _gather_from_field # Use a small grid for fast compilation and test execution NX, NY = 128, 64 class TestSyncBodiesSkeleton(unittest.TestCase): """Test sync_bodies() with real GPU -- double_buffer path.""" 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() class TestSyncBodiesEsoPull(unittest.TestCase): """Test sync_bodies() with esopull streaming.""" def _make_esopull_sim(self, nx=128, ny=64) -> Simulation: """Create a Simulation with a small esopull D2Q9 grid.""" cfg = { "grid": {"lattice_model": "D2Q9", "nx": nx, "ny": ny, "nz": 1}, "physics": {"data_type": "FP32", "viscosity": 0.05, "velocity": 0.03, "rho": 1.0}, "method": { "collision": "SRT", "streaming": "esopull", "store_precision": "FP32", "ddf_shifting": False, "les": {"enabled": False, "cs": 0.16, "closed_form": True}, "trt": {"magic_param": 0.1875}, "inlet": {"profile": "parabolic", "scheme": "zou_he_local"}, "outlet": {"mode": "neq_extrap", "backflow_clamp": True, "blend_alpha": 0.7}, "y_wall_bc": "bounce_back", "omega_guard": {"min": 0.01, "max": 1.99}, }, "cuda": {"threads_per_block": 256, "compute_capability": "auto"}, } tmpd = tempfile.mkdtemp(prefix="esopull_sync_") lbm_path = os.path.join(tmpd, "config_lbm.json") with open(lbm_path, "w") as f: json.dump(cfg, f) return Simulation(lbm_config_path=lbm_path) def test_esopull_add_body(self): """EsoPull: add body, sync, run, read finite force.""" sim = self._make_esopull_sim() sim.initialize() sim.run(100) sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) sim.sync_bodies() self.assertEqual(sim.bodies.count, 1) sim.run(100) force = sim.read_force(0) self.assertTrue(np.all(np.isfinite(force)), f"EsoPull sync: finite force expected, got {force}") sim.close() def test_esopull_remove_body(self): """EsoPull: add body → run → remove → sync → run.""" sim = self._make_esopull_sim() sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) sim.initialize() sim.run(100) sim.remove_body(0) sim.sync_bodies() self.assertEqual(sim.bodies.count, 0) sim.run(100) macro = sim.get_macroscopic() self.assertTrue(np.all(np.isfinite(macro["ux"])), "Macroscopic ux should be finite after esopull sync") sim.close() def test_esopull_add_remove_add(self): """EsoPull: add → remove → add cycle with finite checks.""" sim = self._make_esopull_sim() sim.initialize() sim.run(100) # Add sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) sim.sync_bodies() sim.run(100) # Remove sim.remove_body(0) sim.sync_bodies() self.assertEqual(sim.bodies.count, 0) sim.run(100) # Add again sim.add_body("circle", center=(NX // 2, NY // 2), radius=6) sim.sync_bodies() self.assertEqual(sim.bodies.count, 1) sim.run(100) force = sim.read_force(0) self.assertTrue(np.all(np.isfinite(force)), f"Finite force after esopull add-remove-add: {force}") sim.close() def test_esopull_roundtrip_fullgrid(self): """EsoPull: decode/encode roundtrip on full grid (no patch).""" sim = self._make_esopull_sim() sim.initialize() sim.run(100) nx, ny, nq = sim.field.nx, sim.field.ny, sim.field.nq all_cells = np.arange(nx * ny, dtype=np.uint32) step = sim.stepper.step_count stream = sim.stream s_gpu, s_host, idx_gpu = sim.field.esopull_sync_decode( all_cells, step, stream) _scatter_to_field(sim.field, s_host, all_cells) _gather_from_field(sim.field, s_host, all_cells) sim.field.esopull_sync_encode( s_gpu, s_host, idx_gpu, len(all_cells), step, stream) stream.synchronize() sim.run(100) macro = sim.get_macroscopic() self.assertTrue(np.all(np.isfinite(macro["ux"])), "ux finite after full-grid roundtrip") sim.close() if __name__ == "__main__": unittest.main()