CelerisLab/tests/integration/test_sync_bodies_skeleton.py

"""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()