Add snapshot fixes from polioland

[jonathanhhb]

Fix snapshot continuity bugs in LaserFrame and add continuity tests
Multi-segment simulations (save snapshot, reload, continue) had three bugs
in the core snapshot infrastructure that caused visible discontinuities:

1. date_of_death offset: absolute-timestep death schedules were not adjusted
   when reloading into a new segment starting at t=0, causing agents to live
   too long (their deaths were scheduled far in the future relative to the
   new timeline). save_snapshot now accepts t= to
   automatically subtracts time steps elapsed from date_of_death.

2. pop_final: results.pop[0] was reset to init_pop on reload, ignoring all
   births and deaths from the first segment and causing a population jump at
   the boundary. save_snapshot now accepts pop_final= (per-node population at
   snapshot time); load_snapshot returns it via pars["pop_final"] so callers
   can restore results.pop[0, :].

3. keep_mask: save_snapshot now accepts an optional keep_mask boolean array
   to squash terminal-state agents (e.g. fully-recovered) before writing,
   reducing snapshot size without requiring callers to mutate the frame first.

All three params are optional and backwards-compatible. t_snap and pop_final
are surfaced in the returned pars dict for model-specific use.

Also adds:

• 8 unit tests in test_laserframe.py covering round-trips, offset math,
  clamping, and absent-field behaviour
• tests/test_snapshot_continuity.py: a model-agnostic continuity test built
  on a minimal deterministic ABM (fixed lifespan + constant birth rate) that
  directly exercises all three bugs and doubles as a runnable visual tool

[Copilot]

Pull request overview

This PR updates LaserFrame snapshot save/load to support multi-segment simulations without discontinuities, and adds regression tests to catch the previously observed boundary issues.

Changes:

• Extend LaserFrame.save_snapshot() to optionally persist t_snap, pop_final, and apply an optional keep_mask.
• Extend LaserFrame.load_snapshot() to surface t_snap/pop_final via pars and offset/clamp date_of_death when t_snap is present.
• Add unit and integration-style tests to validate snapshot round-trips and continuity across a save/reload boundary.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
src/laser/core/laserframe.py	Adds optional snapshot metadata (t_snap, pop_final) and continuity adjustment logic on load; introduces keep_mask handling.
tests/test_laserframe.py	Adds unit tests for t_snap, pop_final, keep_mask, and date_of_death offset/clamping behavior.
tests/test_snapshot_continuity.py	Adds a deterministic staged-vs-flat continuity test to detect boundary discontinuities.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The docstring says the live frame is not mutated, but when keep_mask is provided this method calls self.squash(keep_mask), which mutates the frame (shrinks count and compacts scalar arrays). Please document this side effect explicitly (or clarify that only the date_of_death offset is non-mutating) so callers don’t assume save_snapshot() is safe to call mid-simulation without changing in-memory state.

Suggested change

	to >= 1). The live frame is not mutated.
	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).
	to >= 1). This offsetting is applied only to the data written to disk;
	the in-memory 'date_of_death' property on the live frame is not mutated.
	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 on this LaserFrame before writing, which
	mutates the live frame in-place by shrinking ``_count`` and compacting
	1-D scalar numpy arrays. Callers should not assume that the frame remains
	unchanged when a ``keep_mask`` is supplied.

[Copilot]

save_snapshot() writes the full date_of_death dataset via _save(), then immediately deletes and recreates it when t is provided. For large frames this doubles I/O and memory for that column. Consider avoiding the initial write for date_of_death (e.g., allow _save() to skip/override specific properties) so the adjusted values are written once.

[Copilot]

test_pop_final_continuity() currently doesn’t actually validate the pop_final plumbing: pop_seg2[0] is initialized from frame2.count, which already equals frame1.count at the boundary, so the assertion would still pass even if pop_final were never saved/loaded/used. To exercise the bug described in the PR, simulate a caller that resets the segment-2 population baseline incorrectly (e.g., set pop_seg2[0] to INIT_POP/0 before applying pop_final), and assert that continuity is only restored when pars["pop_final"] is present and applied.

[clorton]

I would really like us to implement a base/core save/load to HDF5 which we then wrap with all the special cases (date_of_death offset, patches restore along with people, etc.).