Merge reviewed checkpoint additions for release candidate

2026-10-21 15:20:04 +08:00
parent e214480452 b97fc79af9
commit 26e49451ec
1 changed files with 66 additions and 0 deletions
@@ -0,0 +1,66 @@
+"""Checkpoint save -> load round trip for a scene with passive bodies. Requires GPU.
+
+A fresh Simulation with the same scene, restored from the checkpoint must hold the same populations,
+flags, body count and step count, and must stay bitwise identical to the original
+after both continue for the same number of steps. Checkpoints with Jets are
+refused (documented limitation) and are not covered here.
+"""
+
+import os
+import tempfile
+import unittest
+
+import numpy as np
+
+from ._cases import make_sim
+
+# Three separated cylinders (pinball-like); cut-link receivers stay in bulk fluid.
+CIRCLES = (((40, 32), 6), ((56, 22), 6), ((56, 42), 6))
+
+
+def scene():
+    sim = make_sim()
+    sim.initialize()
+    for center, radius in CIRCLES:
+        sim.add_body("circle", center=center, radius=radius)
+    sim.sync_bodies()
+    return sim
+
+
+class CheckpointBodiesGPU(unittest.TestCase):
+    def test_round_trip_is_bitwise_and_continues_identically(self):
+        sim = scene()
+        sim.run(60)
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f:
+            path = f.name
+        restored = None
+        try:
+            saved = sim.save_checkpoint(path)
+            # The loader requires the same body count and geometry: rebuild the scene.
+            restored = scene()
+            restored.load_checkpoint(saved)
+
+            self.assertEqual(restored.bodies.count, len(CIRCLES))
+            self.assertEqual(restored.stepper.step_count, sim.stepper.step_count)
+            np.testing.assert_array_equal(restored.get_flags(), sim.get_flags())
+            a, b = sim.get_ddf(), restored.get_ddf()
+            self.assertEqual(a.dtype, b.dtype)
+            self.assertEqual(a.tobytes(), b.tobytes())
+
+            sim.run(40)
+            restored.run(40)
+            self.assertEqual(restored.stepper.step_count, sim.stepper.step_count)
+            a, b = sim.get_ddf(), restored.get_ddf()
+            self.assertTrue(np.isfinite(a).all())
+            self.assertEqual(a.tobytes(), b.tobytes())
+            for i in range(len(CIRCLES)):
+                self.assertTrue(np.all(np.isfinite(restored.read_force(i))))
+        finally:
+            if restored is not None:
+                restored.close()
+            os.unlink(path)
+            sim.close()
+
+
+if __name__ == "__main__":
+    unittest.main()