Fix saved simulation reloads and live-viewer cancellation

SyMBac/physics/joints.py

@@ -1,6 +1,4 @@
 import pymunk
-from SyMBac.physics.config import CellConfig
-from SyMBac.physics.segments import CellSegment
 
 class CellJoint(pymunk.PivotJoint):
     """Represents a custom PivotJoint connecting two CellSegments.
@@ -16,7 +14,7 @@ class CellJoint(pymunk.PivotJoint):
         max_force: The maximum force this pivot joint can exert,
             derived from `config.PIVOT_JOINT_STIFFNESS`.
     """
-    def __init__(self, segment_a: CellSegment, segment_b: CellSegment, config: CellConfig) -> None:
+    def __init__(self, segment_a, segment_b, config, anchor_b=None) -> None:
         """Initializes a CellJoint instance.
 
         This constructor sets up a `pymunk.PivotJoint` between two
@@ -28,6 +26,10 @@ def __init__(self, segment_a: CellSegment, segment_b: CellSegment, config: CellC
             segment_b: The second CellSegment to connect.
             config: The CellConfig object providing joint parameters.
         """
+        if anchor_b is not None:
+            super().__init__(segment_a, segment_b, config, anchor_b)
+            return
+
         self.joint_distance: float = config.SEGMENT_RADIUS / config.GRANULARITY
         anchor_on_prev = (self.joint_distance / 2, 0)  # Local coordinates relative to the body for the pivot joint
         anchor_on_curr = (-self.joint_distance / 2, 0)
@@ -43,7 +45,11 @@ class CellRotaryLimitJoint(pymunk.RotaryLimitJoint):
     """
     A custom RotaryLimitJoint that uses the configuration from CellConfig.
     """
-    def __init__(self, segment_a: CellSegment, segment_b: CellSegment, config: CellConfig) -> None:
+    def __init__(self, segment_a, segment_b, config, max_angle=None) -> None:
+        if max_angle is not None:
+            super().__init__(segment_a, segment_b, config, max_angle)
+            return
+
         assert config.STIFFNESS is not None
         assert config.MAX_BEND_ANGLE is not None
         assert config.ROTARY_LIMIT_JOINT
@@ -58,7 +64,11 @@ class CellDampedRotarySpring(pymunk.DampedRotarySpring):
     """
     A custom DampedRotarySpring that uses the configuration from CellConfig.
     """
-    def __init__(self, segment_a: CellSegment, segment_b: CellSegment, config: CellConfig) -> None:
+    def __init__(self, segment_a, segment_b, config, stiffness=None, damping=None) -> None:
+        if damping is not None:
+            super().__init__(segment_a, segment_b, config, stiffness, damping)
+            return
+
         assert config.DAMPED_ROTARY_SPRING
         assert config.ROTARY_SPRING_STIFFNESS is not None
         assert config.ROTARY_SPRING_DAMPING is not None

SyMBac/simulation.py

@@ -773,6 +773,7 @@ def step_and_capture_frame(frame_idx):
             }
             state_lock = threading.Lock()
             stop_event = threading.Event()
+            cancelled = False
 
             def publish_live_frame(frame_idx):
                 image = render_live_frame_image(
@@ -810,15 +811,15 @@ def stop_loop():
                     progress_bar.update(state["frame_idx"] - state["progress_idx"])
                     state["progress_idx"] = state["frame_idx"]
                 progress_bar.close()
-                if worker_thread.is_alive():
-                    worker_thread.join(timeout=1.0)
+                worker_thread.join()
                 viewer.close()
 
             last_drawn_frame_idx = None
             try:
                 while True:
                     viewer.process_events()
                     if viewer.closed:
+                        cancelled = not state["done"]
                         stop_loop()
                         break
 
@@ -848,6 +849,8 @@ def stop_loop():
 
             if state["worker_error"] is not None:
                 raise RuntimeError("Simulation worker failed during show_window=True execution.") from state["worker_error"]
+            if cancelled:
+                return
         else:
             for frame_idx in tqdm(range(total_frames), desc='Running simulation'):
                 step_and_capture_frame(frame_idx)

tests/test_simulation_persistence.py

@@ -2,6 +2,7 @@
 import os
 import pickle
 import sys
+import threading
 import types
 
 import numpy as np
@@ -16,6 +17,7 @@
 from SyMBac.simulation import Simulation
 import SyMBac.simulation as simulation_module
 import SyMBac.cell_snapshot as cell_snapshot_module
+import SyMBac.live_viewer as live_viewer_module
 import SyMBac.physics.simulator as physics_simulator_module
 from SyMBac.physics.microfluidic_geometry import (
     Bounds2D,
@@ -24,6 +26,7 @@
     SegmentPrimitive,
     TrenchGeometrySpec,
 )
+from SyMBac.physics.joints import CellDampedRotarySpring, CellJoint, CellRotaryLimitJoint
 from pymunk.vec2d import Vec2d
 
 
@@ -269,6 +272,106 @@ def tracking_replace(src, dst):
     assert list(save_dir.glob(".tmp_*.p")) == []
 
 
+def test_real_simulation_save_reload(tmp_path):
+    kwargs = _simulation_kwargs(tmp_path)
+    kwargs["sim_length"] = 1
+    kwargs["cell_config_overrides"] = {
+        "DAMPED_ROTARY_SPRING": True,
+        "ROTARY_SPRING_STIFFNESS": 100.0,
+        "ROTARY_SPRING_DAMPING": 10.0,
+    }
+
+    simulation = Simulation(**kwargs)
+    simulation.run_simulation(show_window=False)
+
+    reload_kwargs = dict(kwargs)
+    reload_kwargs["save_dir"] = str(tmp_path / "reload_target")
+    reload_kwargs["load_sim_dir"] = kwargs["save_dir"]
+    reloaded = Simulation(**reload_kwargs)
+
+    assert len(reloaded.cell_timeseries) == 1
+    for constraint_type in (CellJoint, CellRotaryLimitJoint, CellDampedRotarySpring):
+        saved_count = sum(
+            isinstance(constraint, constraint_type)
+            for constraint in simulation.space.constraints
+        )
+        reloaded_count = sum(
+            isinstance(constraint, constraint_type)
+            for constraint in reloaded.space.constraints
+        )
+        assert saved_count > 0
+        assert reloaded_count == saved_count
+
+
+def test_live_viewer_cancellation_joins_worker_without_saving_partial_artifacts(
+    tmp_path, monkeypatch
+):
+    real_thread = threading.Thread
+    worker_in_step = threading.Event()
+    release_worker = threading.Event()
+    worker_threads = []
+    join_timeouts = []
+    join_started_while_alive = []
+    step_count = 0
+
+    class _BlockingSimulator(_DummySimulator):
+        def step(self):
+            nonlocal step_count
+            step_count += 1
+            if step_count == 4:
+                worker_in_step.set()
+                if not release_worker.wait(timeout=5):
+                    raise TimeoutError("Test worker was not released")
+
+    class _ClosingViewer:
+        def __init__(self, **kwargs):
+            pass
+
+        def process_events(self):
+            if not worker_in_step.wait(timeout=5):
+                raise TimeoutError("Simulation worker did not reach the blocking frame")
+
+        @property
+        def closed(self):
+            return True
+
+        def close(self):
+            return None
+
+    class _TrackingThread(real_thread):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            worker_threads.append(self)
+
+        def join(self, timeout=None):
+            join_started_while_alive.append(self.is_alive())
+            join_timeouts.append(timeout)
+            release_worker.set()
+            super().join(timeout)
+
+    monkeypatch.setattr(physics_simulator_module, "Simulator", _BlockingSimulator)
+    monkeypatch.setattr(cell_snapshot_module, "CellSnapshot", _DummySnapshot)
+    monkeypatch.setattr(live_viewer_module, "LiveSimulationViewer", _ClosingViewer)
+    monkeypatch.setattr(threading, "Thread", _TrackingThread)
+
+    kwargs = _simulation_kwargs(tmp_path)
+    kwargs["sim_length"] = 3
+    simulation = Simulation(**kwargs)
+    try:
+        simulation.run_simulation(show_window=True)
+    finally:
+        release_worker.set()
+        for worker_thread in worker_threads:
+            real_thread.join(worker_thread, timeout=5)
+
+    assert step_count == 4
+    assert join_started_while_alive == [True]
+    assert join_timeouts == [None]
+    assert all(not worker_thread.is_alive() for worker_thread in worker_threads)
+    assert not (tmp_path / "sim" / "cell_timeseries.p").exists()
+    assert not (tmp_path / "sim" / "space_timeseries.p").exists()
+
+
 def test_load_sim_dir_missing_artifacts_error_names_expected_files(tmp_path):
     load_dir = tmp_path / "load_artifacts"
     load_dir.mkdir()