Add snapshot fixes from polioland

.github/workflows/github-actions.yml

@@ -190,15 +190,14 @@ jobs:
         CIBW_ARCHS: '${{ matrix.cibw_arch }}'
         CIBW_BUILD: '${{ matrix.cibw_build }}'
         CIBW_BUILD_VERBOSITY: '3'
-        CIBW_TEST_REQUIRES: >
-          tox
-          tox-direct
+        CIBW_BEFORE_TEST: "python -m ensurepip --upgrade"
+        CIBW_TEST_REQUIRES: tox>=4
         CIBW_TEST_COMMAND: >
           cd {project} &&
-          tox --skip-pkg-install --direct-yolo -e ${{ matrix.tox_env }} -v
+          tox -e ${{ matrix.tox_env }} -v
         CIBW_TEST_COMMAND_WINDOWS: >
           cd /d {project} &&
-          tox --skip-pkg-install --direct-yolo -e ${{ matrix.tox_env }} -v
+          tox -e ${{ matrix.tox_env }} -v
     - name: regular build and test
       env:
         TOXPYTHON: '${{ matrix.toxpython }}'

src/laser/core/laserframe.py

@@ -282,30 +282,62 @@ def squash(self, indices, verbose: bool = False) -> None:
 
         return
 
-    def save_snapshot(self, path, results_r=None, pars=None):
+    def save_snapshot(self, path, results_r=None, pars=None, t=None, pop_final=None, keep_mask=None):
         """
         Save this LaserFrame and optional extras to an HDF5 snapshot file.
 
         Parameters:
             path (Path): Destination file path
             results_r (np.ndarray): Optional 2D numpy array of recovered counts
             pars (PropertySet or dict): Optional PropertySet or dict of parameters
+            t (int, optional): Current simulation timestep. If provided and the frame
+                has a 'date_of_death' property, that field is written to the snapshot
+                already offset to the new segment's timeline (values minus t, clamped
+                to >= 1). The offset is applied only to the data written to disk;
+                the in-memory frame is not modified.
+            pop_final (np.ndarray, optional): 1-D array of final population counts per node
+                at the snapshot boundary. Returned via pars["pop_final"] on load so the
+                caller can restore results.pop[0, :] for a continuous population time-series.
+            keep_mask (np.ndarray, optional): Boolean array of shape (count,). If provided,
+                squash(keep_mask) is called before writing to compact terminal-state agents
+                out of the snapshot (e.g. fully-recovered agents that no longer affect dynamics).
+                NOTE: this mutates the frame — count is reduced and scalar arrays are
+                compacted in place. Do not pass keep_mask if the frame is still needed
+                in its current state after this call.
         """
         from laser.core.propertyset import PropertySet  # to avoid circular import
 
+        if keep_mask is not None:
+            self.squash(keep_mask)
+
+        # Build overrides for _save: date_of_death is written already offset to the
+        # new segment's timeline so load_snapshot receives clean, ready-to-use values.
+        # Computing this before opening the file keeps the HDF5 write to a single pass.
+        overrides = None
+        if t is not None and "date_of_death" in self._properties:
+            dod = self._properties["date_of_death"]
+            overrides = {"date_of_death": np.maximum(dod[: self._count].astype(np.int64) - t, 1).astype(dod.dtype)}
+
         with h5py.File(path, "w") as f:
-            self._save(f, "people")
+            self._save(f, "people", overrides=overrides)
 
             if results_r is not None:
                 f.create_dataset("recovered", data=results_r)
 
+            if pop_final is not None:
+                f.create_dataset("pop_final", data=pop_final)
+
             if pars is not None and isinstance(pars, (dict, PropertySet)):
                 data = pars.to_dict() if isinstance(pars, PropertySet) else pars
                 self._save_dict(data, f.create_group("pars"))
 
-    def _save(self, parent_group, name):
+    def _save(self, parent_group, name, overrides=None):
         """
         Internal method to save this LaserFrame under the given group name.
+
+        overrides (dict, optional): Maps property names to pre-computed arrays that
+            should be written in place of the live property data (e.g. an already-offset
+            date_of_death). Each value must have length == self._count.
         """
         group = parent_group.create_group(name)
         group.attrs["count"] = self._count
@@ -314,7 +346,8 @@ def _save(self, parent_group, name):
         for name, data in self._properties.items():
             # Currently only saving scalar properties (implied by loading logic)
             if data.shape == (self._capacity,):
-                group.create_dataset(name, data=data[0 : self._count])
+                value = overrides[name] if (overrides and name in overrides) else data[0 : self._count]
+                group.create_dataset(name, data=value)
 
         return
 
@@ -352,7 +385,11 @@ def load_snapshot(cls, path, cbr, nt):
                 shape (time, nodes), or None if not present in the snapshot.
 
             pars (dict): Dictionary of model parameters stored in the snapshot,
-                or empty if none are found.
+                or empty if none are found. Also contains any of the following
+                keys written by save_snapshot:
+                  - "pop_final" (np.ndarray): Final per-node population counts at
+                    the snapshot boundary; use to restore results.pop[0, :] for
+                    a continuous population time-series across the boundary.
 
         Raises:
             ValueError: If only one of cbr or nt is provided.
@@ -363,6 +400,9 @@ def load_snapshot(cls, path, cbr, nt):
             - Snapshots must contain a per-agent 'node_id' property.
             - The recovered array is assumed to be in (time, node) layout.
             - The capacity estimate includes both current and recovered agents at t=0.
+            - If save_snapshot was called with t=, the 'date_of_death' field in the
+              snapshot is already offset to the new segment's timeline and requires
+              no further adjustment on load.
         """
         with h5py.File(path, "r") as f:
             group = f["people"]
@@ -379,6 +419,9 @@ def load_snapshot(cls, path, cbr, nt):
             else:
                 pars = {}
 
+            if "pop_final" in f:
+                pars["pop_final"] = f["pop_final"][()]
+
             # Validate that cbr and nt are both provided or both None
             if (cbr is None) != (nt is None):
                 raise ValueError("cbr and nt must both be provided or both be None. " "Cannot calculate capacity with only one parameter.")

tests/test_laserframe.py

@@ -669,3 +669,115 @@ def test_cannot_reassign_property(self):
             lf.age = np.arange(10)
 
         return
+
+    def test_date_of_death_offset_at_save_time(self):
+        """save_snapshot with t= writes offset date_of_death values; live frame is unchanged."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            count = 5
+            t = 10
+            frame = LaserFrame(capacity=count, initial_count=count)
+            frame.add_scalar_property("date_of_death", dtype=np.int32, default=0)
+            original = np.array([12, 15, 20, 25, 30], dtype=np.int32)
+            frame.date_of_death[:] = original
+            frame.save_snapshot(path, t=t)
+
+            # Snapshot values are already offset
+            loaded, _, _ = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+            expected = np.maximum(original - t, 1)
+            assert np.array_equal(loaded.date_of_death, expected)
+
+            # Live frame is not mutated
+            assert np.array_equal(frame.date_of_death, original)
+        finally:
+            Path(path).unlink()
+
+    def test_date_of_death_clamped_to_one_at_save_time(self):
+        """date_of_death values that would go <= 0 after offset are clamped to 1 in the snapshot."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            count = 3
+            t = 50
+            frame = LaserFrame(capacity=count, initial_count=count)
+            frame.add_scalar_property("date_of_death", dtype=np.int32, default=0)
+            frame.date_of_death[:] = np.array([48, 50, 51], dtype=np.int32)
+            frame.save_snapshot(path, t=t)
+
+            loaded, _, _ = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+
+            assert np.all(loaded.date_of_death >= 1)
+            assert loaded.date_of_death[2] == 1  # 51 - 50 = 1
+            assert loaded.date_of_death[0] == 1  # 48 - 50 = -2 → clamped to 1
+        finally:
+            Path(path).unlink()
+
+    def test_date_of_death_not_offset_without_t(self):
+        """date_of_death in the snapshot is unchanged when save_snapshot is called without t."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            count = 3
+            frame = LaserFrame(capacity=count, initial_count=count)
+            frame.add_scalar_property("date_of_death", dtype=np.int32, default=0)
+            original = np.array([20, 30, 40], dtype=np.int32)
+            frame.date_of_death[:] = original
+            frame.save_snapshot(path)  # no t=
+
+            loaded, _, _ = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+            assert np.array_equal(loaded.date_of_death, original)
+        finally:
+            Path(path).unlink()
+
+    def test_pop_final_round_trip(self):
+        """pop_final saved with save_snapshot is returned in pars dict by load_snapshot."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            frame = LaserFrame(capacity=100, initial_count=10)
+            frame.add_scalar_property("node_id", dtype=np.int32, default=0)
+            pop_final = np.array([1234, 5678], dtype=np.int32)
+            frame.save_snapshot(path, pop_final=pop_final)
+
+            _, _, pars = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+            assert "pop_final" in pars
+            assert np.array_equal(pars["pop_final"], pop_final)
+        finally:
+            Path(path).unlink()
+
+    def test_pop_final_absent_when_not_saved(self):
+        """When pop_final is not passed to save_snapshot, it is absent from loaded pars."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            frame = LaserFrame(capacity=100, initial_count=10)
+            frame.add_scalar_property("node_id", dtype=np.int32, default=0)
+            frame.save_snapshot(path)
+
+            _, _, pars = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+            assert "pop_final" not in pars
+        finally:
+            Path(path).unlink()
+
+    def test_keep_mask_squash_before_save(self):
+        """keep_mask causes squash before writing; only kept agents appear in snapshot."""
+        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
+            path = tmp.name
+        try:
+            count = 10
+            frame = LaserFrame(capacity=count, initial_count=count)
+            frame.add_scalar_property("age", dtype=np.int32, default=0)
+            frame.age[:] = np.arange(count, dtype=np.int32)
+
+            # Keep only agents with age >= 5
+            keep = frame.age >= 5
+            kept_ages = np.arange(5, count, dtype=np.int32)
+
+            frame.save_snapshot(path, keep_mask=keep)
+
+            loaded, _, _ = LaserFrame.load_snapshot(path, cbr=None, nt=None)
+            assert loaded.count == 5
+            assert np.array_equal(loaded.age, kept_ages)
+        finally:
+            Path(path).unlink()