"""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 from CelerisLab.simulation import Simulation, _scatter_to_field, _gather_from_field from ._cases import make_sim # 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 make_sim() 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.""" return make_sim(nx=nx, ny=ny, streaming="esopull") 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) with sim.ctx.activate(): 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()