diff --git a/docs/api/core.rst b/docs/api/core.rst index 12dcf732..891e2a12 100644 --- a/docs/api/core.rst +++ b/docs/api/core.rst @@ -30,6 +30,14 @@ Particle state and results :undoc-members: :show-inheritance: +Checkpointing +------------- + +.. automodule:: sequince.checkpoint + :members: + :undoc-members: + :show-inheritance: + Posterior utilities ------------------- diff --git a/docs/checkpointing.rst b/docs/checkpointing.rst new file mode 100644 index 00000000..8ead034a --- /dev/null +++ b/docs/checkpointing.rst @@ -0,0 +1,223 @@ +Checkpointing +============= + +Checkpointing records the state of a Sequential Monte Carlo run so +that it can continue in a later process. A checkpoint is captured after a +complete iteration: step processing has finished, weights have been normalised, +the new particle state has been added to the result, and callbacks have +finished. Consequently, resuming starts at the next iteration. + +Configuring a checkpoint writer +------------------------------- + +Pass one or more checkpoint writers to +:class:`~sequince.smc.SequentialMonteCarlo`. The ``every`` argument is measured +in completed iterations, so ``every=10`` writes after iterations 10, 20, and so +on. + +.. code-block:: python + + from sequince import SequentialMonteCarlo + from sequince.checkpoint import HDF5CheckpointWriter + + writer = HDF5CheckpointWriter( + "run.h5", + every=10, + save_rng=True, + include_history=True, + ) + smc = SequentialMonteCarlo( + steps=steps, + callbacks=callbacks, + checkpoint_writers=[writer], + ) + result = smc.run( + initial_state, + observations=observations, + rng=123, + shared_context=shared_context, + ) + +The available writers are: + +``HDF5CheckpointWriter`` + Stores arrays and supported values as native HDF5 datasets, and mappings and + sequences as groups. This format is suitable when the checkpoint should be + inspectable or when its arrays may be useful independently. Install the + optional dependency with ``pip install "sequince[hdf5]"``. + +``PickleCheckpointWriter`` + Stores the checkpoint as one opaque pickle payload. This is the simpler + format when inspecting the file outside Python is not important. + +Both writers replace the latest runtime checkpoint. The HDF5 writer appends new +entries to the result history rather than rewriting existing entries on every +iteration. The pickle writer writes a temporary file and then replaces the +previous file. + +What is saved +------------- + +An :class:`~sequince.checkpoint.SMCCheckpoint` contains: + +* the current :class:`~sequince.state.ParticleState`, including its particles, + log-weights, log-evidence, and metadata; +* the next iteration and observation indices; +* mutable context metadata and recorded outputs; +* the observation from the most recently completed iteration; +* optional state supplied by each configured step; +* the random-number-generator state, unless ``save_rng=False``; +* the particle-state history, unless ``include_history=False``. + +A checkpoint deliberately does not contain: + +* the configured steps, callbacks, or their callable dependencies; +* immutable ``shared_context`` values; +* the complete observation sequence. + +The sampler must therefore be reconstructed with the same ordered step +pipeline before loading a checkpoint. Shared inputs and observations must also +be supplied again when resuming. + +Resuming a run +-------------- + +Load the checkpoint, reconstruct the sampler, and use +:meth:`~sequince.smc.SequentialMonteCarlo.resume_from` to prepare a runtime. + +.. code-block:: python + + import h5py + + from sequince import SequentialMonteCarlo + from sequince.checkpoint import SMCCheckpoint + + with h5py.File("run.h5", "r") as checkpoint_file: + checkpoint = SMCCheckpoint.load_hdf5(checkpoint_file) + + smc = SequentialMonteCarlo( + steps=steps, + callbacks=callbacks, + checkpoint_writers=[writer], + ) + runtime = smc.resume_from( + checkpoint, + shared_context=shared_context, + ) + result = smc.run( + runtime=runtime, + observations=observations, + ) + +When observations were consumed before the checkpoint, ``observations`` must +be the original indexable sequence. Sequince resumes at the saved observation +index. A one-shot iterator cannot be repositioned and is therefore rejected. +Runs that do not use observations can continue with ``observations=None``. + +If the checkpoint contains RNG state, it takes precedence over the ``rng`` +argument to ``resume_from``. If ``save_rng=False`` was used, provide an RNG or +seed to ``resume_from`` instead: + +.. code-block:: python + + runtime = smc.resume_from( + checkpoint, + rng=456, + shared_context=shared_context, + ) + +Once the runtime has been prepared, RNG and context arguments belong to the +runtime. Do not pass ``rng``, ``shared_context``, ``metadata``, or ``outputs`` +again to ``run(runtime=...)``. + +Stateful steps +-------------- + +Steps are stateless by default. A step that learns or accumulates information +needed by later iterations should implement ``state_dict()`` and +``load_state_dict()``. The state dictionary should contain mutable learned +state, not constructor configuration or immutable shared inputs. + +.. code-block:: python + + from sequince.steps.base import SMCStep + + class AdaptiveStep(SMCStep): + def __init__(self): + self.scale = 1.0 + + def state_dict(self): + return {"scale": self.scale} + + def load_state_dict(self, state): + self.scale = float(state["scale"]) + + def apply(self, particle_state, context, rng): + # Update and use self.scale here. + return particle_state + +Sequince records the index, module, and qualified class name for every step. On +resume, it validates the complete ordered pipeline before loading any step +state. This prevents state from silently being restored into a different step +configuration. Normalizing flows or adaptive proposals owned by a step should +be included in that step's state dictionary when their fitted parameters are +required by subsequent iterations. + +Result history +-------------- + +With ``include_history=True``, resuming restores all recorded particle states. +This allows history-dependent diagnostics and plots to continue across the +checkpoint boundary. :meth:`~sequince.checkpoint.SMCCheckpoint.result` can also +reconstruct an :class:`~sequince.result.SMCResult` directly from a checkpoint. + +With ``include_history=False``, only the current particle state is available. +The reconstructed result therefore starts with that state. This reduces file +size but cannot recreate plots or diagnostics that require earlier states. + +HDF5 layout +----------- + +The default HDF5 path is ``/checkpoint``. The latest runtime and optional +history are stored separately: + +.. code-block:: text + + /checkpoint + ├── runtime + │ ├── version + │ ├── next_iteration + │ ├── next_observation_index + │ ├── state + │ │ ├── particles + │ │ ├── log_weights + │ │ ├── log_evidence + │ │ └── metadata + │ ├── metadata + │ ├── outputs + │ ├── rng_state + │ └── step_states + │ ├── 00000000 + │ └── 00000001 + └── history + ├── 00000000 + ├── 00000001 + └── 00000002 + +String-keyed mappings become groups, lists and tuples become numbered groups, +and NumPy values and ordinary scalars become native datasets. Strings and bytes +also use datasets, while ``None`` is represented by an empty tagged group. +Unsupported values are stored as tagged pickle datasets. Mapping keys that are +not valid HDF5 path components are escaped and restored when reading. A mapping +with non-string keys uses the pickle fallback as a single value. + +The ``history`` group is absent when history is disabled. A temporary runtime +group may be present if a process is interrupted during a write; a successful +write moves it into place as ``runtime``. + +Loading trusted files +--------------------- + +Pickle can execute arbitrary code while loading. This applies both to pickle +checkpoint files and to unsupported values stored through the HDF5 pickle +fallback. Only load checkpoint files from trusted sources. diff --git a/docs/index.rst b/docs/index.rst index 9dc70561..6e719811 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,6 +11,7 @@ sequential methods. getting_started user_guide + checkpointing examples api/index diff --git a/pyproject.toml b/pyproject.toml index 91fc2ba0..7f1b224d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "minipcn[array-api]>=0.2.0a4", "array-api-compat>=1.5", "array-api-extra", - "orng>=0.1.0a3", + "orng>=0.2.0a1", "frozendict", ] dynamic = [ @@ -43,6 +43,9 @@ emcee = [ corner = [ "corner", ] +hdf5 = [ + "h5py>=3", +] [dependency-groups] test = [ @@ -50,6 +53,7 @@ test = [ "pytest-cov>=4.0", "corner", "emcee>=3.0", + "h5py>=3", ] dev = [ "pre-commit", diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py new file mode 100644 index 00000000..0b87cebd --- /dev/null +++ b/src/sequince/checkpoint.py @@ -0,0 +1,617 @@ +from __future__ import annotations + +import copy +import pickle +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any +from urllib.parse import quote, unquote + +import numpy as np +from orng import RandomGenerator + +from .result import SMCResult +from .state import ParticleState +from .steps.context import StepContext +from .utils import resolve_rng + +CHECKPOINT_VERSION = 1 +_HDF_KIND = "_sequince_kind" +_HDF_PYTHON_TYPE = "_sequince_python_type" + + +def _attr_text(value: Any) -> str: + """Normalize an HDF5 text attribute to ``str``.""" + return value.decode("utf-8") if isinstance(value, bytes) else str(value) + + +def _pickle_bytes(value: Any, protocol: int) -> np.ndarray: + """Return a pickled value as an HDF5-compatible byte array.""" + return np.frombuffer(pickle.dumps(value, protocol=protocol), dtype="u1") + + +def _write_pickled_hdf_value( + parent: Any, name: str, value: Any, protocol: int +) -> Any: + """Write one unsupported leaf as a tagged pickle dataset.""" + if name in parent: + del parent[name] + node = parent.create_dataset(name, data=_pickle_bytes(value, protocol)) + node.attrs[_HDF_KIND] = "pickle" + return node + + +def _write_hdf_value(parent: Any, name: str, value: Any, protocol: int) -> Any: + """Write mappings, sequences, and native leaves to HDF5 recursively.""" + if name in parent: + del parent[name] + + if isinstance(value, Mapping): + if not all(isinstance(key, str) for key in value): + return _write_pickled_hdf_value(parent, name, value, protocol) + node = parent.create_group(name) + node.attrs[_HDF_KIND] = "mapping" + for key, item in value.items(): + child_name = quote(key, safe="") or "%" + _write_hdf_value(node, child_name, item, protocol) + return node + + if isinstance(value, (list, tuple)): + node = parent.create_group(name) + node.attrs[_HDF_KIND] = "list" if isinstance(value, list) else "tuple" + for index, item in enumerate(value): + _write_hdf_value(node, f"{index:08d}", item, protocol) + return node + + if value is None: + node = parent.create_group(name) + node.attrs[_HDF_KIND] = "none" + return node + + if isinstance(value, str): + import h5py + + node = parent.create_dataset( + name, data=value, dtype=h5py.string_dtype(encoding="utf-8") + ) + node.attrs[_HDF_KIND] = "string" + return node + + if isinstance(value, bytes): + node = parent.create_dataset( + name, data=np.frombuffer(value, dtype="u1") + ) + node.attrs[_HDF_KIND] = "bytes" + return node + + if isinstance(value, (np.ndarray, np.generic)): + if value.dtype.hasobject: + return _write_pickled_hdf_value(parent, name, value, protocol) + try: + node = parent.create_dataset(name, data=value) + except (TypeError, ValueError): + return _write_pickled_hdf_value(parent, name, value, protocol) + node.attrs[_HDF_KIND] = "numpy" + return node + + if isinstance(value, (bool, int, float, complex)): + try: + node = parent.create_dataset(name, data=value) + except (OverflowError, TypeError, ValueError): + return _write_pickled_hdf_value(parent, name, value, protocol) + node.attrs[_HDF_KIND] = "scalar" + node.attrs[_HDF_PYTHON_TYPE] = type(value).__name__ + return node + + return _write_pickled_hdf_value(parent, name, value, protocol) + + +def _read_hdf_value(node: Any) -> Any: + """Reconstruct a value written by :func:`_write_hdf_value`.""" + if _HDF_KIND not in node.attrs: + raise ValueError(f"HDF5 node {node.name!r} has no Sequince type tag.") + kind = _attr_text(node.attrs[_HDF_KIND]) + + if kind == "mapping": + return { + "" if name == "%" else unquote(name): _read_hdf_value(node[name]) + for name in node + } + + if kind in ("list", "tuple"): + items = [_read_hdf_value(node[name]) for name in sorted(node.keys())] + return items if kind == "list" else tuple(items) + + if kind == "none": + return None + if kind == "string": + string_value = node[()] + return ( + string_value.decode("utf-8") + if isinstance(string_value, bytes) + else str(string_value) + ) + if kind == "bytes": + return node[...].tobytes() + if kind == "numpy": + return node[()] + if kind == "scalar": + scalar_value = node[()] + if isinstance(scalar_value, np.generic): + scalar_value = scalar_value.item() + scalar_type = _attr_text(node.attrs[_HDF_PYTHON_TYPE]) + scalar_constructors = { + "bool": bool, + "int": int, + "float": float, + "complex": complex, + } + return scalar_constructors[scalar_type](scalar_value) + if kind == "pickle": + return pickle.loads(node[...].tobytes()) + raise ValueError(f"Unsupported Sequince HDF5 node kind {kind!r}.") + + +class CheckpointWriter(ABC): + """Base class for persisting committed SMC checkpoints.""" + + def __init__( + self, + *, + every: int = 1, + save_rng: bool = True, + include_history: bool = True, + ) -> None: + if every <= 0: + raise ValueError("every must be positive.") + self.every = int(every) + self.save_rng = save_rng + self.include_history = include_history + + def is_due(self, next_iteration: int) -> bool: + """Return whether a checkpoint is due after this iteration.""" + return next_iteration % self.every == 0 + + def write(self, checkpoint: "SMCCheckpoint") -> None: + """Apply this writer's options and persist the checkpoint.""" + prepared = replace( + checkpoint, + rng_state=checkpoint.rng_state if self.save_rng else None, + history=checkpoint.history if self.include_history else None, + ) + self._write(prepared) + + @abstractmethod + def _write(self, checkpoint: "SMCCheckpoint") -> None: + """Persist a checkpoint prepared for this writer.""" + + +@dataclass(frozen=True) +class StepState: + """Checkpoint state associated with one configured SMC step.""" + + index: int + module: str + qualname: str + state: dict[str, Any] + + @classmethod + def from_step(cls, index: int, step: Any) -> "StepState": + state_fn = getattr(step, "state_dict", None) + state = {} if state_fn is None else state_fn() + if not isinstance(state, Mapping): + raise TypeError( + f"Step {index} ({type(step).__name__}) state_dict() must " + "return a mapping." + ) + step_type = type(step) + return cls( + index=index, + module=step_type.__module__, + qualname=step_type.__qualname__, + state=copy.deepcopy(dict(state)), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "index": self.index, + "module": self.module, + "qualname": self.qualname, + "state": copy.deepcopy(self.state), + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "StepState": + missing = { + key + for key in ("index", "module", "qualname", "state") + if key not in payload + } + if missing: + raise ValueError( + "Step state is missing required keys: " + + ", ".join(sorted(missing)) + ) + state = payload["state"] + if not isinstance(state, Mapping): + raise TypeError("Step state 'state' must be a mapping.") + module = payload["module"] + qualname = payload["qualname"] + if not isinstance(module, str) or not isinstance(qualname, str): + raise TypeError("Step state module and qualname must be strings.") + return cls( + index=int(payload["index"]), + module=module, + qualname=qualname, + state=copy.deepcopy(dict(state)), + ) + + +def capture_step_states(steps: Sequence[Any]) -> tuple[StepState, ...]: + """Capture mutable state from a configured step pipeline.""" + return tuple( + StepState.from_step(index, step) for index, step in enumerate(steps) + ) + + +def restore_steps(steps: Sequence[Any], states: Sequence[StepState]) -> None: + """Restore step state after validating pipeline identity and order.""" + if len(steps) != len(states): + raise ValueError( + "Checkpoint step pipeline length does not match the configured " + f"pipeline ({len(states)} != {len(steps)})." + ) + for index, (step, saved) in enumerate(zip(steps, states, strict=True)): + step_type = type(step) + identity = (step_type.__module__, step_type.__qualname__) + saved_identity = (saved.module, saved.qualname) + if saved.index != index or saved_identity != identity: + raise ValueError( + f"Checkpoint step {index} is " + f"{saved.module}.{saved.qualname}, but the configured step is " + f"{step_type.__module__}.{step_type.__qualname__}." + ) + load_fn = getattr(step, "load_state_dict", None) + if load_fn is None and saved.state: + raise ValueError( + f"Step {index} ({step_type.__name__}) has checkpoint " + "state but does not implement load_state_dict()." + ) + for step, saved in zip(steps, states, strict=True): + load_fn = getattr(step, "load_state_dict", None) + if load_fn is not None: + load_fn(copy.deepcopy(saved.state)) + + +@dataclass(frozen=True) +class SMCCheckpoint: + """Detached end-of-iteration state for resuming an SMC run.""" + + state: ParticleState + next_iteration: int + next_observation_index: int + metadata: dict[str, Any] + outputs: dict[str, Any] + step_states: tuple[StepState, ...] + rng_state: dict[str, Any] | None = None + history: tuple[ParticleState, ...] | None = None + observation: Any = None + version: int = CHECKPOINT_VERSION + + @classmethod + def from_runtime( + cls, + state: ParticleState, + context: StepContext, + result: SMCResult, + steps: Sequence[Any], + rng: Any | None = None, + *, + include_history: bool = True, + ) -> "SMCCheckpoint": + """Capture a checkpoint after a completed SMC iteration.""" + state_copy = state.copy() + history = None + if include_history: + history = tuple(item.copy() for item in result.states) + return cls( + state=state_copy, + next_iteration=context.t + 1, + next_observation_index=context.next_observation_index, + metadata=copy.deepcopy( + { + key: value + for key, value in context.metadata.items() + if key != "particle_transform" + } + ), + outputs=copy.deepcopy(dict(context.outputs)), + step_states=capture_step_states(steps), + rng_state=None if rng is None else resolve_rng(rng).state_dict(), + history=history, + observation=copy.deepcopy(context.observation), + ) + + def to_dict(self) -> dict[str, Any]: + """Return a detached plain-Python checkpoint representation.""" + return { + "version": self.version, + "state": self.state.to_dict(), + "next_iteration": self.next_iteration, + "next_observation_index": self.next_observation_index, + "metadata": copy.deepcopy(self.metadata), + "outputs": copy.deepcopy(self.outputs), + "step_states": [item.to_dict() for item in self.step_states], + "rng_state": copy.deepcopy(self.rng_state), + "history": ( + None + if self.history is None + else [item.to_dict() for item in self.history] + ), + "observation": copy.deepcopy(self.observation), + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "SMCCheckpoint": + """Reconstruct a checkpoint from :meth:`to_dict` output.""" + if not isinstance(payload, Mapping): + raise TypeError("SMCCheckpoint payload must be a mapping.") + required = { + "version", + "state", + "next_iteration", + "next_observation_index", + "metadata", + "outputs", + "step_states", + } + missing = required.difference(payload) + if missing: + raise ValueError( + "SMCCheckpoint payload is missing required keys: " + + ", ".join(sorted(missing)) + ) + version = payload["version"] + if version != CHECKPOINT_VERSION: + raise ValueError( + f"Unsupported SMC checkpoint version {version!r}; expected " + f"{CHECKPOINT_VERSION}." + ) + rng_state = payload.get("rng_state") + if rng_state is not None and not isinstance(rng_state, Mapping): + raise TypeError( + "SMCCheckpoint 'rng_state' must be a mapping or None." + ) + for key in ("metadata", "outputs"): + if not isinstance(payload[key], Mapping): + raise TypeError(f"SMCCheckpoint '{key}' must be a mapping.") + step_payload = payload["step_states"] + if not isinstance(step_payload, Sequence) or isinstance( + step_payload, (str, bytes) + ): + raise TypeError("SMCCheckpoint 'step_states' must be a sequence.") + history_payload = payload.get("history") + if history_payload is not None and ( + not isinstance(history_payload, Sequence) + or isinstance(history_payload, (str, bytes)) + ): + raise TypeError( + "SMCCheckpoint 'history' must be a sequence or None." + ) + next_iteration = int(payload["next_iteration"]) + next_observation_index = int(payload["next_observation_index"]) + if next_iteration < 0 or next_observation_index < 0: + raise ValueError( + "Checkpoint iteration and observation indices must be " + "non-negative." + ) + return cls( + state=ParticleState.from_dict(copy.deepcopy(payload["state"])), + next_iteration=next_iteration, + next_observation_index=next_observation_index, + metadata=copy.deepcopy(dict(payload["metadata"])), + outputs=copy.deepcopy(dict(payload["outputs"])), + step_states=tuple( + StepState.from_dict(item) for item in step_payload + ), + rng_state=( + None if rng_state is None else copy.deepcopy(dict(rng_state)) + ), + history=( + None + if history_payload is None + else tuple( + ParticleState.from_dict(copy.deepcopy(item)) + for item in history_payload + ) + ), + observation=copy.deepcopy(payload.get("observation")), + version=version, + ) + + def restore_rng(self) -> Any | None: + """Return a live ORNG generator restored from this checkpoint.""" + if self.rng_state is None: + return None + return RandomGenerator.from_state_dict(self.rng_state) + + def result( + self, shared_context: Mapping[str, Any] | None = None + ) -> SMCResult: + """Return the result history represented by this checkpoint.""" + states = list(self.history or (self.state,)) + context = StepContext( + t=self.next_iteration - 1, + next_observation_index=self.next_observation_index, + observation=copy.deepcopy(self.observation), + shared=dict(shared_context or {}), + metadata=copy.deepcopy(self.metadata), + outputs=copy.deepcopy(self.outputs), + ) + return SMCResult(states=states, context=context) + + def save_pickle( + self, + file_path: str | Path, + protocol: int | None = None, + ) -> None: + """Save this checkpoint as a versioned pickle payload.""" + protocol = ( + pickle.HIGHEST_PROTOCOL if protocol is None else int(protocol) + ) + with Path(file_path).open("wb") as checkpoint_file: + pickle.dump(self.to_dict(), checkpoint_file, protocol=protocol) + + @classmethod + def load_pickle(cls, file_path: str | Path) -> "SMCCheckpoint": + """Load a trusted checkpoint from a pickle file.""" + with Path(file_path).open("rb") as checkpoint_file: + payload = pickle.load(checkpoint_file) + return cls.from_dict(payload) + + def save_hdf5( + self, + h5_file, + path: str = "checkpoint", + dsetname: str = "runtime", + protocol: int | None = None, + *, + replace_history: bool = False, + ) -> None: + """Save runtime state and append result history in HDF5.""" + protocol = ( + pickle.HIGHEST_PROTOCOL if protocol is None else int(protocol) + ) + group = h5_file.require_group(path) + runtime_payload = self.to_dict() + runtime_payload.pop("history") + + if self.history is None: + if "history" in group: + del group["history"] + else: + history_group = group.require_group("history") + if replace_history: + del group["history"] + history_group = group.require_group("history") + existing = len(history_group) + if existing > len(self.history): + del group["history"] + history_group = group.require_group("history") + existing = 0 + for index in range(existing, len(self.history)): + _write_hdf_value( + history_group, + f"{index:08d}", + self.history[index].to_dict(), + protocol, + ) + + temporary_name = f"{dsetname}.tmp" + _write_hdf_value(group, temporary_name, runtime_payload, protocol) + if dsetname in group: + del group[dsetname] + group.move(temporary_name, dsetname) + + @classmethod + def load_hdf5( + cls, + h5_file, + path: str = "checkpoint", + dsetname: str = "runtime", + ) -> "SMCCheckpoint": + """Load a checkpoint from its recursive HDF5 representation.""" + group = h5_file[path] + checkpoint_payload = _read_hdf_value(group[dsetname]) + if not isinstance(checkpoint_payload, dict): + raise TypeError("HDF5 checkpoint runtime must be a mapping.") + checkpoint_payload["history"] = None + if "history" in group: + history_group = group["history"] + checkpoint_payload["history"] = [ + _read_hdf_value(history_group[name]) + for name in sorted(history_group) + ] + return cls.from_dict(checkpoint_payload) + + +class HDF5CheckpointWriter(CheckpointWriter): + """Persist committed SMC checkpoints to HDF5 at a fixed cadence.""" + + def __init__( + self, + file_path: str | Path, + *, + every: int = 1, + path: str = "checkpoint", + dsetname: str = "runtime", + save_rng: bool = True, + include_history: bool = True, + ) -> None: + super().__init__( + every=every, + save_rng=save_rng, + include_history=include_history, + ) + file_path = Path(file_path) + if not file_path.name.lower().endswith((".h5", ".hdf5")): + raise ValueError( + "Checkpoint file must be an HDF5 file (.h5 or .hdf5)." + ) + self.file_path = file_path + self.path = path + self.dsetname = dsetname + self._last_iteration: int | None = None + self.file_path.parent.mkdir(parents=True, exist_ok=True) + + def _write(self, checkpoint: SMCCheckpoint) -> None: + """Write a committed checkpoint to the configured HDF5 location.""" + import h5py + + replace_history = ( + self._last_iteration is None + or checkpoint.next_iteration <= self._last_iteration + ) + with h5py.File(self.file_path, "a") as h5_file: + checkpoint.save_hdf5( + h5_file, + path=self.path, + dsetname=self.dsetname, + replace_history=replace_history, + ) + self._last_iteration = checkpoint.next_iteration + + +class PickleCheckpointWriter(CheckpointWriter): + """Atomically persist committed SMC checkpoints as pickle files.""" + + def __init__( + self, + file_path: str | Path, + *, + every: int = 1, + save_rng: bool = True, + include_history: bool = True, + protocol: int | None = None, + ) -> None: + super().__init__( + every=every, + save_rng=save_rng, + include_history=include_history, + ) + file_path = Path(file_path) + if file_path.suffix.lower() not in (".pkl", ".pickle"): + raise ValueError( + "Checkpoint file must be a pickle file (.pkl or .pickle)." + ) + self.file_path = file_path + self.protocol = protocol + self.file_path.parent.mkdir(parents=True, exist_ok=True) + + def _write(self, checkpoint: SMCCheckpoint) -> None: + """Write a committed checkpoint using an atomic file replacement.""" + temporary_path = self.file_path.with_name(self.file_path.name + ".tmp") + checkpoint.save_pickle(temporary_path, protocol=self.protocol) + temporary_path.replace(self.file_path) diff --git a/src/sequince/flows.py b/src/sequince/flows.py index be193fe2..46051953 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Protocol from ._typing import Array @@ -42,7 +43,9 @@ def log_prob( ... def sample( - self, num_samples: int = 1, context: StepContext | None = None + self, + num_samples: int = 1, + context: StepContext | None = None, ) -> Array: """Draw samples from the flow distribution.""" ... @@ -53,6 +56,31 @@ def flow_asarray(flow: NormalizingFlow, array: Any): return flow.xp.asarray(array) +def flow_state_dict(flow: NormalizingFlow) -> dict[str, Any]: + """Return persistent flow state when the adapter exposes it.""" + state_fn = getattr(flow, "state_dict", None) + if state_fn is None: + return {} + state = state_fn() + if not isinstance(state, Mapping): + raise TypeError("Flow state_dict() must return a mapping.") + return copy.deepcopy(dict(state)) + + +def load_flow_state_dict( + flow: NormalizingFlow, state: Mapping[str, Any] +) -> None: + """Restore persistent flow state when supported by the adapter.""" + load_fn = getattr(flow, "load_state_dict", None) + if load_fn is None: + if state: + raise ValueError( + f"{type(flow).__name__} does not support checkpoint state." + ) + return + load_fn(copy.deepcopy(dict(state))) + + def clone_flow(flow: NormalizingFlow) -> NormalizingFlow: """Return a fresh copy of ``flow``. diff --git a/src/sequince/proposals.py b/src/sequince/proposals.py index fa3dcd09..369e2c61 100644 --- a/src/sequince/proposals.py +++ b/src/sequince/proposals.py @@ -1,11 +1,18 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Protocol +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Protocol from ._array_api import asarray, namespace from ._typing import Array -from .flows import NormalizingFlow, clone_flow, flow_asarray +from .flows import ( + NormalizingFlow, + clone_flow, + flow_asarray, + flow_state_dict, + load_flow_state_dict, +) if TYPE_CHECKING: from .steps.context import StepContext @@ -61,6 +68,15 @@ def fit( self._flow.fit(backend_particles, context=context) return self + def state_dict(self) -> dict[str, Any]: + return {"flow": flow_state_dict(self._flow)} + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + flow_state = state.get("flow", {}) + if not isinstance(flow_state, Mapping): + raise TypeError("FlowProposal state 'flow' must be a mapping.") + load_flow_state_dict(self._flow, flow_state) + def log_prob( self, particles: Array, @@ -88,6 +104,9 @@ def sample_and_log_prob( self._flow, "sample_and_log_prob", None ) if callable(flow_sample_and_log_prob): - return flow_sample_and_log_prob(n_samples, context=context) + return flow_sample_and_log_prob( + n_samples, + context=context, + ) samples = self.sample(n_samples, context=context) return samples, self.log_prob(samples, context=context) diff --git a/src/sequince/result.py b/src/sequince/result.py index b6ab82b0..5836de37 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -1,9 +1,10 @@ from __future__ import annotations import math +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Sequence +from typing import TYPE_CHECKING, Any, Callable import numpy as np from numpy.random import Generator as NumPyGenerator @@ -11,14 +12,13 @@ from ._array_api import asarray, namespace, to_numpy from ._typing import Array from .state import ParticleState +from .steps.context import StepContext from .steps.nested import NESTED_SMC_OUTPUT_KEY from .utils import resolve_rng if TYPE_CHECKING: from orng import RandomGenerator - from .steps.base import StepContext - def _prepare_axis(ax: Any, **subplot_kwargs: Any) -> Any: if ax is not None: diff --git a/src/sequince/smc.py b/src/sequince/smc.py index b6e317d4..2a1c141c 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -1,17 +1,18 @@ from __future__ import annotations +import copy import logging -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Iterable, - List, - Mapping, - MutableMapping, - Sequence, -) +from collections.abc import Iterable, Mapping, MutableMapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable +from .checkpoint import ( + CheckpointWriter, + SMCCheckpoint, + StepState, + capture_step_states, + restore_steps, +) from .result import SMCResult from .state import ParticleState from .steps.base import StepContext @@ -24,6 +25,20 @@ logger = logging.getLogger(__name__) +@dataclass +class SMCRuntime: + """Prepared SMC runtime state for fresh or resumed execution.""" + + state: ParticleState + produced_states: list[ParticleState] + t: int + next_observation_index: int + shared_context: dict[str, Any] + metadata: MutableMapping[str, Any] + outputs: MutableMapping[str, Any] + rng: Any + + class SequentialMonteCarlo: """Composable Sequential Monte Carlo runner. @@ -41,6 +56,9 @@ class SequentialMonteCarlo: with the current state, context, and an ``SMCResult`` containing all states produced so far. Callbacks can be used for logging, diagnostics, or custom stopping criteria that require access to the full history of states. + checkpoint_writers: + Optional persistence hooks invoked with a committed checkpoint after all + ordinary callbacks have completed. auto_normalize: If ``True`` (default) automatically normalize the particle weights at the end of each iteration. @@ -55,25 +73,43 @@ def __init__( ] ], *, - callbacks: Sequence[ - Callable[[ParticleState, StepContext, SMCResult], None] - ] - | None = None, + callbacks: Sequence[Callable[..., None]] | None = None, + checkpoint_writers: Sequence[CheckpointWriter] | None = None, auto_normalize: bool = True, ): self._steps = list(steps) self._callbacks = list(callbacks) if callbacks is not None else [] + self._checkpoint_writers = list(checkpoint_writers or ()) self._auto_normalize = auto_normalize + def state_dict(self) -> dict[str, Any]: + """Return mutable state for the configured step pipeline.""" + return { + "steps": [ + item.to_dict() for item in capture_step_states(self._steps) + ] + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore mutable state for the configured step pipeline.""" + step_payload = state.get("steps", []) + if not isinstance(step_payload, Sequence): + raise TypeError("SMC runner state 'steps' must be a sequence.") + restore_steps( + self._steps, + [StepState.from_dict(item) for item in step_payload], + ) + def run( self, - initial_state: ParticleState, + initial_state: ParticleState | None = None, observations: Iterable[Any] | None = None, num_steps: int | None = None, rng: "RandomGenerator" | np.random.Generator | int | None = None, shared_context: Mapping[str, Any] | None = None, metadata: MutableMapping[str, Any] | None = None, outputs: MutableMapping[str, Any] | None = None, + runtime: SMCRuntime | None = None, ) -> SMCResult: """Run the SMC algorithm. @@ -101,11 +137,37 @@ def run( for communication between steps and callbacks. outputs: Optional mutable dictionary for storing outputs from steps. + runtime: + Optional prepared runtime returned by :meth:`resume_from`. """ - rng = resolve_rng(rng) - shared_context = shared_context or {} - metadata = metadata or {} - outputs = outputs or {} + if runtime is not None: + if initial_state is not None: + raise ValueError( + "run() accepts either initial_state or runtime, not both." + ) + if any( + value is not None + for value in (rng, shared_context, metadata, outputs) + ): + raise ValueError( + "run(runtime=...) does not accept rng, shared_context, " + "metadata, or outputs." + ) + else: + if initial_state is None: + raise ValueError( + "run() requires either initial_state or runtime." + ) + runtime = SMCRuntime( + state=initial_state, + produced_states=[], + t=0, + next_observation_index=0, + shared_context=dict(shared_context or {}), + metadata=metadata or {}, + outputs=outputs or {}, + rng=resolve_rng(rng), + ) logger.info( "Starting SMC run with %d steps (auto_normalize=%s)", @@ -118,22 +180,57 @@ def run( ) if observations is not None: logger.debug("Observations iterator supplied to SMC run.") + if ( + runtime.next_observation_index > 0 + and observations is not None + and not isinstance(observations, Sequence) + ): + raise ValueError( + "Resuming with observations requires an indexable sequence or None." + ) - produced_states: List[ParticleState] = [] - state = initial_state - + observation_sequence = ( + observations if isinstance(observations, Sequence) else None + ) observation_iterator = ( - iter(observations) if observations is not None else None + iter(observations) + if observations is not None and observation_sequence is None + else None ) + if ( + runtime.next_observation_index > 0 + and observation_sequence is None + and observations is not None + ): + raise ValueError( + "Resumed runtime only supports observations=None or sequence observations." + ) - t = 0 steps_remaining = num_steps if num_steps is not None else None + produced_states = list(runtime.produced_states) + state = runtime.state + t = runtime.t + next_observation_index = runtime.next_observation_index + shared_context = runtime.shared_context + metadata = runtime.metadata + outputs = runtime.outputs + rng = runtime.rng + context = None while steps_remaining is None or steps_remaining > 0: logger.debug("Beginning iteration %d.", t) - if observation_iterator is not None: + if observation_sequence is not None: + if next_observation_index < len(observation_sequence): + observation = observation_sequence[next_observation_index] + next_observation_index += 1 + elif steps_remaining is None: + break + else: + observation = None + elif observation_iterator is not None: try: observation = next(observation_iterator) + next_observation_index += 1 except StopIteration: if steps_remaining is None: break @@ -147,6 +244,7 @@ def run( context = StepContext( t=t, + next_observation_index=next_observation_index, observation=observation, shared=shared_context, metadata=metadata, @@ -173,6 +271,7 @@ def run( state.log_evidence, ) produced_states.append(state) + result = SMCResult(states=list(produced_states), context=context) for callback in self._callbacks: callback_name = getattr( @@ -181,11 +280,31 @@ def run( logger.debug( "Iteration %d: invoking callback %s.", t, callback_name ) - callback( - state, - context, - SMCResult(states=produced_states, context=context), + callback(state, context, result) + + next_iteration = context.t + 1 + checkpoint_writers = [ + writer + for writer in self._checkpoint_writers + if writer.is_due(next_iteration) + ] + if checkpoint_writers: + include_history = any( + writer.include_history for writer in checkpoint_writers + ) + save_rng = any( + writer.save_rng for writer in checkpoint_writers + ) + checkpoint = SMCCheckpoint.from_runtime( + state=state, + context=context, + result=result, + steps=self._steps, + rng=rng if save_rng else None, + include_history=include_history, ) + for writer in checkpoint_writers: + writer.write(checkpoint) if context.metadata["stop_requested"]: logger.info( @@ -201,4 +320,37 @@ def run( len(produced_states), len(produced_states), ) - return SMCResult(states=produced_states, context=context) + if context is None or not produced_states: + raise ValueError("SMC run produced no states.") + return SMCResult(states=list(produced_states), context=context) + + def resume_from( + self, + checkpoint: SMCCheckpoint | Mapping[str, Any], + *, + rng: "RandomGenerator" | np.random.Generator | int | None = None, + shared_context: Mapping[str, Any] | None = None, + ) -> SMCRuntime: + """Reconstruct runtime state from a checkpoint for use with :meth:`run`.""" + checkpoint_obj = ( + checkpoint + if isinstance(checkpoint, SMCCheckpoint) + else SMCCheckpoint.from_dict(checkpoint) + ) + restore_steps(self._steps, checkpoint_obj.step_states) + restored_rng = checkpoint_obj.restore_rng() + state = checkpoint_obj.state.copy() + history = checkpoint_obj.history or (checkpoint_obj.state,) + restored_metadata = copy.deepcopy(checkpoint_obj.metadata) + restored_metadata.pop("stop_requested", None) + restored_metadata.pop("stop_immediately", None) + return SMCRuntime( + state=state, + produced_states=[item.copy() for item in history], + t=checkpoint_obj.next_iteration, + next_observation_index=checkpoint_obj.next_observation_index, + shared_context=dict(shared_context or {}), + metadata=restored_metadata, + outputs=copy.deepcopy(checkpoint_obj.outputs), + rng=restored_rng if restored_rng is not None else resolve_rng(rng), + ) diff --git a/src/sequince/state.py b/src/sequince/state.py index 5a8ed77c..7419baca 100644 --- a/src/sequince/state.py +++ b/src/sequince/state.py @@ -1,6 +1,8 @@ from __future__ import annotations +import copy import math +from collections.abc import Mapping from dataclasses import dataclass, field from numbers import Real from typing import Any, MutableMapping @@ -60,7 +62,7 @@ def copy(self) -> "ParticleState": particles=array_copy(self.particles, xp=xp), log_weights=array_copy(self.log_weights, xp=xp), log_evidence=float(self.log_evidence), - metadata=dict(self.metadata), + metadata=copy.deepcopy(self.metadata), ) def normalize(self) -> tuple["ParticleState", float]: @@ -131,3 +133,35 @@ def resample(self, indices: Array) -> "ParticleState": log_evidence=self.log_evidence, metadata=dict(self.metadata, ancestors=indices), ) + + def to_dict(self) -> dict[str, Any]: + """Return a detached dictionary representation of this state.""" + copied = self.copy() + return { + "particles": copied.particles, + "log_weights": copied.log_weights, + "log_evidence": copied.log_evidence, + "metadata": copied.metadata, + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "ParticleState": + """Reconstruct a state from :meth:`to_dict` output.""" + if not isinstance(payload, Mapping): + raise TypeError("ParticleState payload must be a mapping.") + missing = { + key + for key in ("particles", "log_weights", "log_evidence", "metadata") + if key not in payload + } + if missing: + raise ValueError( + "ParticleState payload is missing required keys: " + + ", ".join(sorted(missing)) + ) + return cls( + particles=payload["particles"], + log_weights=payload["log_weights"], + log_evidence=float(payload["log_evidence"]), + metadata=copy.deepcopy(payload["metadata"]), + ) diff --git a/src/sequince/steps/base.py b/src/sequince/steps/base.py index 259fb62b..cd4dc859 100644 --- a/src/sequince/steps/base.py +++ b/src/sequince/steps/base.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any from ..utils import resolve_rng from .context import StepContext @@ -29,3 +30,14 @@ def apply( rng: "RandomGenerator", ) -> "ParticleState": raise NotImplementedError("SMCStep subclasses must implement apply().") + + def state_dict(self) -> dict[str, Any]: + """Return mutable state required by subsequent SMC iterations.""" + return {} + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore state produced by :meth:`state_dict`.""" + if state: + raise ValueError( + f"{type(self).__name__} does not support checkpoint state." + ) diff --git a/src/sequince/steps/context.py b/src/sequince/steps/context.py index 3543e9a4..b72d8ef0 100644 --- a/src/sequince/steps/context.py +++ b/src/sequince/steps/context.py @@ -20,6 +20,7 @@ class StepContext: """Runtime context passed to each SMC step.""" t: int + next_observation_index: int = 0 observation: Any = None shared: SharedContext = field(default_factory=dict) """Dictionary for shared user-defined context across steps. diff --git a/src/sequince/steps/control.py b/src/sequince/steps/control.py index e06fb0ec..f2a4c891 100644 --- a/src/sequince/steps/control.py +++ b/src/sequince/steps/control.py @@ -1,7 +1,8 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Callable, Sequence +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Callable, Sequence from .base import SMCStep, StepContext @@ -150,6 +151,15 @@ def apply( ) return state + def state_dict(self) -> dict[str, Any]: + return {"step": self.step.state_dict()} + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + nested_state = state.get("step", {}) + if not isinstance(nested_state, Mapping): + raise TypeError("PeriodicStep state 'step' must be a mapping.") + self.step.load_state_dict(nested_state) + class RepeatStep(SMCStep): """Repeat a step or sequence of steps a fixed number of times.""" diff --git a/src/sequince/steps/nested.py b/src/sequince/steps/nested.py index f3e430c1..4b396b59 100644 --- a/src/sequince/steps/nested.py +++ b/src/sequince/steps/nested.py @@ -4,6 +4,8 @@ import copy import logging +from collections.abc import Mapping +from typing import Any from orng import RandomGenerator @@ -72,3 +74,12 @@ def apply( if self.record_outputs: context.record_output(self.output_key, result) return result.final_state + + def state_dict(self) -> dict[str, Any]: + return {"runner": self.runner.state_dict()} + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + runner_state = state.get("runner", {}) + if not isinstance(runner_state, Mapping): + raise TypeError("NestedSMCStep state 'runner' must be a mapping.") + self.runner.load_state_dict(runner_state) diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index f489dcb2..7e506fcd 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -9,7 +9,13 @@ import array_api_extra as xpx from ._array_api import asarray, namespace, to_numpy -from .flows import NormalizingFlow, clone_flow, flow_asarray +from .flows import ( + NormalizingFlow, + clone_flow, + flow_asarray, + flow_state_dict, + load_flow_state_dict, +) logger = logging.getLogger(__name__) @@ -151,6 +157,15 @@ def clone(self) -> "FlowTransform": fit_kwargs=self._fit_kwargs, ) + def state_dict(self): + return {"flow": flow_state_dict(self._flow)} + + def load_state_dict(self, state): + flow_state = state.get("flow", {}) + if not isinstance(flow_state, Mapping): + raise TypeError("FlowTransform state 'flow' must be a mapping.") + load_flow_state_dict(self._flow, flow_state) + def fit(self, particles, context=None): backend_particles = flow_asarray(self._flow, particles) logger.debug( diff --git a/src/sequince/utils.py b/src/sequince/utils.py index 9e3e49b3..694c755c 100644 --- a/src/sequince/utils.py +++ b/src/sequince/utils.py @@ -20,18 +20,19 @@ def resolve_rng( rng: RNGLike = None, *, backend: str | None = None, -) -> RandomGenerator | np.random.Generator: - """Return the provided RNG or create an ``RandomGenerator`` for the requested backend.""" +) -> RandomGenerator: + """Return an ORNG generator for the supplied RNG-like value.""" if isinstance(rng, RandomGenerator): return rng if isinstance(rng, np.random.Generator): - return rng + return RandomGenerator(backend="numpy", generator=rng) if isinstance(rng, int): return RandomGenerator(backend=backend or "numpy", seed=rng) if rng is None: return RandomGenerator(backend=backend or "numpy") raise TypeError( - "rng must be an RandomGenerator, numpy.random.Generator, an int seed, or None." + "rng must be an orng.RandomGenerator, numpy.random.Generator, " + "an int seed, or None." ) @@ -43,34 +44,31 @@ def call_with_optional_context( """Call fn with base_args and append context only if required.""" signature = inspect.signature(fn) params = signature.parameters.values() - positional = [ - p - for p in params - if p.kind + parameter + for parameter in params + if parameter.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ) ] - required = sum(1 for p in positional if p.default is inspect._empty) - if required > len(base_args): - return fn(*base_args, context) - + required = sum( + parameter.default is inspect._empty for parameter in positional + ) has_varargs = any( - p.kind == inspect.Parameter.VAR_POSITIONAL for p in params + parameter.kind == inspect.Parameter.VAR_POSITIONAL + for parameter in params ) - if not has_varargs and len(base_args) > len(positional): - raise TypeError( - "Function cannot accept the provided number of positional arguments." - ) - - try: + max_positional = float("inf") if has_varargs else len(positional) + base_count = len(base_args) + if required <= base_count + 1 <= max_positional: + return fn(*base_args, context) + if required <= base_count <= max_positional: return fn(*base_args) - except TypeError: - if has_varargs or len(positional) > len(base_args): - return fn(*base_args, context) - raise + raise TypeError( + "Function cannot accept the provided positional arguments." + ) def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py new file mode 100644 index 00000000..5d672c17 --- /dev/null +++ b/tests/test_checkpointing.py @@ -0,0 +1,681 @@ +from __future__ import annotations + +import copy +from collections.abc import Sequence + +import h5py +import numpy as np +import pytest + +from sequince.checkpoint import ( + CheckpointWriter, + HDF5CheckpointWriter, + PickleCheckpointWriter, + SMCCheckpoint, +) +from sequince.result import SMCResult +from sequince.smc import SequentialMonteCarlo +from sequince.state import ParticleState +from sequince.steps.base import SMCStep, StepContext +from sequince.utils import resolve_rng + + +class StochasticStep(SMCStep): + def apply( + self, + state: ParticleState, + context: StepContext, + rng, + ) -> ParticleState: + observation_value = ( + 0.0 if context.observation is None else float(context.observation) + ) + particles = np.asarray(state.particles) + log_weights = np.asarray(state.log_weights) + noise = np.asarray(rng.normal(size=particles.shape)) + weight_noise = np.asarray(rng.normal(size=log_weights.shape)) + + metadata = copy.deepcopy(state.metadata) + metadata.setdefault("noise_sums", []).append(float(noise.sum())) + metadata["iteration"] = context.t + + context.metadata.setdefault("seen_observations", []).append( + context.observation + ) + context.metadata["last_iteration"] = context.t + context.record_output("noise_sums", float(noise.sum())) + context.record_output("observations", context.observation) + + return ParticleState( + particles=particles + noise + observation_value, + log_weights=log_weights + weight_noise + observation_value, + log_evidence=state.log_evidence, + metadata=metadata, + ) + + +class DeterministicStep(SMCStep): + def apply( + self, + state: ParticleState, + context: StepContext, + rng, + ) -> ParticleState: + metadata = copy.deepcopy(state.metadata) + metadata["iteration"] = context.t + context.record_output("iterations", context.t) + return ParticleState( + particles=np.asarray(state.particles) + 1.0, + log_weights=np.asarray(state.log_weights), + log_evidence=state.log_evidence, + metadata=metadata, + ) + + +class StatefulStep(SMCStep): + def __init__(self) -> None: + self.counter = 0 + + def apply(self, state, context, rng): + particles = np.asarray(state.particles) + self.counter + self.counter += 1 + return state.with_updates(particles=particles) + + def state_dict(self): + return {"counter": self.counter} + + def load_state_dict(self, state): + self.counter = int(state["counter"]) + + +class OutputCallback: + def __call__( + self, + state: ParticleState, + context: StepContext, + result: SMCResult, + ) -> None: + context.metadata.setdefault("callback_iterations", []).append( + context.t + ) + context.record_output("callback_iterations", context.t) + context.metadata["callback_result_length"] = len(result) + + +class CheckpointRecorder(CheckpointWriter): + def __init__(self, *, save_rng: bool = True) -> None: + super().__init__(save_rng=save_rng) + self.checkpoints: list[SMCCheckpoint] = [] + + def _write(self, checkpoint: SMCCheckpoint) -> None: + self.checkpoints.append(checkpoint) + + +def _assert_states_equal(left: ParticleState, right: ParticleState) -> None: + np.testing.assert_allclose(left.particles, right.particles) + np.testing.assert_allclose(left.log_weights, right.log_weights) + assert left.log_evidence == pytest.approx(right.log_evidence) + assert left.metadata == right.metadata + + +def _assert_results_equal( + left: SMCResult, + right: SMCResult, +) -> None: + assert len(left.states) == len(right.states) + for left_state, right_state in zip(left.states, right.states): + _assert_states_equal(left_state, right_state) + assert left.context.t == right.context.t + assert ( + left.context.next_observation_index + == right.context.next_observation_index + ) + assert left.context.observation == right.context.observation + assert dict(left.context.shared) == dict(right.context.shared) + assert left.context.metadata == right.context.metadata + assert left.context.outputs == right.context.outputs + + +def _make_smc( + *checkpoint_writers, step: SMCStep | None = None +) -> SequentialMonteCarlo: + return SequentialMonteCarlo( + steps=[step or StochasticStep()], + callbacks=[OutputCallback()], + checkpoint_writers=checkpoint_writers, + ) + + +def _run_with_checkpoint_split( + particle_state: ParticleState, + observations: Sequence[float] | None, + *, + total_steps: int, + split_after: int, +) -> tuple[SMCResult, SMCResult, SMCCheckpoint]: + full_result = _make_smc().run( + particle_state, + observations=observations, + num_steps=total_steps, + rng=resolve_rng(2024), + shared_context={"alpha": 1.5}, + metadata={"tag": "run"}, + outputs={"seed_output": ["start"]}, + ) + + recorder = CheckpointRecorder() + _make_smc(recorder).run( + particle_state, + observations=observations, + num_steps=split_after, + rng=resolve_rng(2024), + shared_context={"alpha": 1.5}, + metadata={"tag": "run"}, + outputs={"seed_output": ["start"]}, + ) + + runtime = _make_smc().resume_from( + recorder.checkpoints[-1].to_dict(), + rng=resolve_rng(999), + shared_context={"alpha": 1.5}, + ) + resumed_result = _make_smc().run( + runtime=runtime, + observations=observations, + num_steps=total_steps - split_after, + ) + + return full_result, resumed_result, recorder.checkpoints[-1] + + +def test_checkpoint_round_trip_preserves_runtime_objects( + particle_state: ParticleState, +) -> None: + context = StepContext( + t=2, + next_observation_index=3, + observation={"batch": 2}, + shared={"alpha": 1.5}, + metadata={"notes": ["x"]}, + outputs={"metric": [1.0, 2.0]}, + ) + state = particle_state.with_updates( + metadata={"temperature": 0.5, "nested": {"value": [1, 2, 3]}} + ) + result = SMCResult(states=[particle_state, state], context=context) + rng = np.random.default_rng(123) + rng.normal(size=3) + + checkpoint = SMCCheckpoint.from_runtime( + state=state, + context=context, + result=result, + steps=[], + rng=rng, + ) + restored = SMCCheckpoint.from_dict(checkpoint.to_dict()) + + _assert_states_equal(checkpoint.state, restored.state) + _assert_results_equal(checkpoint.result(), restored.result()) + assert checkpoint.metadata == restored.metadata + assert checkpoint.outputs == restored.outputs + assert restored.observation == {"batch": 2} + restored_rng = restored.restore_rng() + assert restored_rng is not None + np.testing.assert_allclose( + rng.normal(size=4), + restored_rng.normal(size=4), + ) + + +def test_checkpoint_version_one_payload_without_observation_still_loads( + particle_state: ParticleState, +) -> None: + context = StepContext(t=0, observation={"batch": 1}) + checkpoint = SMCCheckpoint.from_runtime( + state=particle_state, + context=context, + result=SMCResult(states=[particle_state], context=context), + steps=[], + ) + payload = checkpoint.to_dict() + payload.pop("observation") + + restored = SMCCheckpoint.from_dict(payload) + + assert restored.observation is None + assert restored.result().context.observation is None + + +def test_checkpoint_round_trip_allows_missing_rng_state( + particle_state: ParticleState, +) -> None: + context = StepContext(t=1, next_observation_index=2) + result = SMCResult(states=[particle_state], context=context) + + checkpoint = SMCCheckpoint.from_runtime( + state=particle_state, + context=context, + result=result, + steps=[], + rng=None, + ) + restored = SMCCheckpoint.from_dict(checkpoint.to_dict()) + + assert restored.rng_state is None + assert restored.restore_rng() is None + + +def test_checkpoint_writers_run_after_callbacks( + particle_state: ParticleState, +) -> None: + recorder = CheckpointRecorder() + result = _make_smc(recorder).run( + particle_state, + observations=[0.0, 1.0, 2.0], + num_steps=3, + rng=resolve_rng(123), + ) + + assert len(recorder.checkpoints) == 3 + assert [ + checkpoint.next_iteration for checkpoint in recorder.checkpoints + ] == [ + 1, + 2, + 3, + ] + assert recorder.checkpoints[-1].next_observation_index == 3 + assert recorder.checkpoints[-1].outputs["callback_iterations"] == [ + 0, + 1, + 2, + ] + assert recorder.checkpoints[-1].metadata["callback_iterations"] == [ + 0, + 1, + 2, + ] + assert ( + recorder.checkpoints[-1].result().context.outputs + == result.context.outputs + ) + + +def test_hdf5_checkpoint_writer_writes_latest_checkpoint( + particle_state: ParticleState, + tmp_path, +) -> None: + file_path = tmp_path / "checkpoint.h5" + writer = HDF5CheckpointWriter(file_path, path="checkpoint", every=1) + result = _make_smc(writer).run( + particle_state, + observations=[0.0, 1.0, 2.0], + num_steps=3, + rng=resolve_rng(123), + metadata={ + "nested": { + "values": [1, 2, 3], + "unsupported_leaf": range(2, 7, 2), + } + }, + ) + + with h5py.File(file_path, "r") as h5_file: + assert len(h5_file["checkpoint/history"]) == 3 + runtime = h5_file["checkpoint/runtime"] + assert isinstance(runtime, h5py.Group) + assert isinstance(runtime["state"], h5py.Group) + assert isinstance(runtime["state/particles"], h5py.Dataset) + assert isinstance(runtime["metadata/nested"], h5py.Group) + assert isinstance(runtime["metadata/nested/values"], h5py.Group) + assert ( + runtime["metadata/nested/unsupported_leaf"].attrs["_sequince_kind"] + == "pickle" + ) + assert isinstance( + h5_file["checkpoint/history/00000000/particles"], h5py.Dataset + ) + checkpoint = SMCCheckpoint.load_hdf5(h5_file, path="checkpoint") + + _assert_states_equal(checkpoint.state, result.final_state) + _assert_results_equal(checkpoint.result(), result) + assert checkpoint.next_iteration == 3 + assert checkpoint.next_observation_index == 3 + assert checkpoint.rng_state is not None + + +def test_hdf5_checkpoint_writer_respects_every_and_optional_rng( + particle_state: ParticleState, + tmp_path, +) -> None: + file_path = tmp_path / "checkpoint_every.hdf5" + writer = HDF5CheckpointWriter( + file_path, + path="checkpoint", + every=2, + save_rng=False, + ) + result = _make_smc(writer).run( + particle_state, + observations=[0.0, 1.0, 2.0], + num_steps=3, + rng=resolve_rng(123), + ) + + with h5py.File(file_path, "r") as h5_file: + checkpoint = SMCCheckpoint.load_hdf5(h5_file, path="checkpoint") + + assert checkpoint.next_iteration == 2 + assert checkpoint.next_observation_index == 2 + assert checkpoint.rng_state is None + _assert_states_equal(checkpoint.state, result.states[1]) + + +def test_hdf5_checkpoint_can_omit_result_history( + particle_state: ParticleState, + tmp_path, +) -> None: + file_path = tmp_path / "checkpoint_without_history.h5" + writer = HDF5CheckpointWriter(file_path, include_history=False) + _make_smc(writer).run( + particle_state, + num_steps=2, + rng=123, + ) + + with h5py.File(file_path, "r") as h5_file: + assert "history" not in h5_file["checkpoint"] + checkpoint = SMCCheckpoint.load_hdf5(h5_file) + + assert checkpoint.history is None + assert len(checkpoint.result()) == 1 + + +def test_hdf5_writer_replaces_history_when_reused_for_a_fresh_run( + particle_state: ParticleState, + tmp_path, +) -> None: + file_path = tmp_path / "reused_writer.h5" + writer = HDF5CheckpointWriter(file_path, every=2) + smc = SequentialMonteCarlo(steps=[], checkpoint_writers=[writer]) + + smc.run(particle_state, num_steps=2, rng=1) + replacement_state = particle_state.with_updates( + particles=np.full_like(particle_state.particles, 9.0) + ) + replacement_result = smc.run(replacement_state, num_steps=2, rng=2) + + with h5py.File(file_path, "r") as h5_file: + checkpoint = SMCCheckpoint.load_hdf5(h5_file) + + _assert_results_equal(checkpoint.result(), replacement_result) + + +def test_pickle_checkpoint_writer_writes_latest_checkpoint( + particle_state: ParticleState, + tmp_path, +) -> None: + file_path = tmp_path / "checkpoint.pkl" + writer = PickleCheckpointWriter(file_path, every=1) + result = _make_smc(writer).run( + particle_state, + observations=[0.0, 1.0, 2.0], + num_steps=3, + rng=resolve_rng(123), + ) + + checkpoint = SMCCheckpoint.load_pickle(file_path) + + _assert_states_equal(checkpoint.state, result.final_state) + _assert_results_equal(checkpoint.result(), result) + assert checkpoint.next_iteration == 3 + assert checkpoint.next_observation_index == 3 + assert checkpoint.rng_state is not None + assert not file_path.with_name(file_path.name + ".tmp").exists() + + +def test_resume_matches_uninterrupted_run_with_sequence_observations( + particle_state: ParticleState, +) -> None: + observations = [0.5, -1.0, 2.0, 0.25] + full_result, resumed_result, checkpoint = _run_with_checkpoint_split( + particle_state, + observations, + total_steps=4, + split_after=2, + ) + + _assert_results_equal(full_result, resumed_result) + assert checkpoint.next_observation_index == 2 + assert len(resumed_result.states) == len(full_result.states) + _assert_states_equal(resumed_result.states[1], full_result.states[1]) + _assert_states_equal(resumed_result.states[2], full_result.states[2]) + + +def test_resume_matches_uninterrupted_run_with_none_observations( + particle_state: ParticleState, +) -> None: + full_result, resumed_result, checkpoint = _run_with_checkpoint_split( + particle_state, + None, + total_steps=4, + split_after=2, + ) + + _assert_results_equal(full_result, resumed_result) + assert checkpoint.next_observation_index == 0 + + +def test_resume_without_saved_rng_works_for_deterministic_steps( + particle_state: ParticleState, +) -> None: + full_result = _make_smc(step=DeterministicStep()).run( + particle_state, + observations=[0.0, 1.0, 2.0, 3.0], + num_steps=4, + rng=resolve_rng(2024), + ) + + recorder = CheckpointRecorder(save_rng=False) + _make_smc(recorder, step=DeterministicStep()).run( + particle_state, + observations=[0.0, 1.0, 2.0, 3.0], + num_steps=2, + rng=resolve_rng(2024), + ) + + resumed_result = _make_smc(step=DeterministicStep()).run( + runtime=_make_smc(step=DeterministicStep()).resume_from( + recorder.checkpoints[-1], + rng=resolve_rng(999), + ), + observations=[0.0, 1.0, 2.0, 3.0], + num_steps=2, + ) + + _assert_results_equal(full_result, resumed_result) + assert recorder.checkpoints[-1].rng_state is None + + +@pytest.mark.parametrize("immediate", [False, True]) +def test_resume_clears_transient_stop_controls( + particle_state: ParticleState, + immediate: bool, +) -> None: + step = DeterministicStep() + context = StepContext(t=0) + if immediate: + context.request_stop_immediately() + else: + context.request_stop() + checkpoint = SMCCheckpoint.from_runtime( + state=particle_state, + context=context, + result=SMCResult(states=[particle_state], context=context), + steps=[step], + rng=123, + ) + resumed_smc = _make_smc(step=DeterministicStep()) + + runtime = resumed_smc.resume_from(checkpoint) + + assert not runtime.metadata.get("stop_requested", False) + assert not runtime.metadata.get("stop_immediately", False) + result = resumed_smc.run(runtime=runtime, num_steps=2) + assert len(result.states) == 3 + + +def test_resume_restores_stateful_step(particle_state: ParticleState) -> None: + full_step = StatefulStep() + full_result = _make_smc(step=full_step).run( + particle_state, + num_steps=4, + rng=123, + ) + + recorder = CheckpointRecorder() + first_step = StatefulStep() + _make_smc(recorder, step=first_step).run( + particle_state, + num_steps=2, + rng=123, + ) + + resumed_step = StatefulStep() + resumed_smc = _make_smc(step=resumed_step) + runtime = resumed_smc.resume_from(recorder.checkpoints[-1]) + resumed_result = resumed_smc.run(runtime=runtime, num_steps=2) + + _assert_results_equal(full_result, resumed_result) + assert resumed_step.counter == 4 + + +def test_emcee_resume_uses_checkpointed_master_rng( + particle_state: ParticleState, +) -> None: + pytest.importorskip("emcee") + from sequince.steps.mutation import EmceeMutation + + def log_prob(particles): + values = np.asarray(particles) + return -0.5 * np.sum(values**2, axis=-1) + + full_result = _make_smc( + step=EmceeMutation(log_prob, n_steps=2, vectorize=True) + ).run(particle_state, num_steps=3, rng=123) + + recorder = CheckpointRecorder() + _make_smc( + recorder, + step=EmceeMutation(log_prob, n_steps=2, vectorize=True), + ).run(particle_state, num_steps=1, rng=123) + + resumed_smc = _make_smc( + step=EmceeMutation(log_prob, n_steps=2, vectorize=True) + ) + runtime = resumed_smc.resume_from(recorder.checkpoints[-1]) + resumed_result = resumed_smc.run(runtime=runtime, num_steps=2) + + _assert_results_equal(full_result, resumed_result) + + +def test_resume_rejects_different_step_pipeline( + particle_state: ParticleState, +) -> None: + recorder = CheckpointRecorder() + _make_smc(recorder, step=StatefulStep()).run( + particle_state, + num_steps=1, + rng=123, + ) + + with pytest.raises(ValueError, match="configured step"): + _make_smc(step=DeterministicStep()).resume_from( + recorder.checkpoints[-1] + ) + + +def test_resume_with_non_sequence_observations_raises( + particle_state: ParticleState, +) -> None: + recorder = CheckpointRecorder() + _make_smc(recorder).run( + particle_state, + observations=[0.0, 1.0], + num_steps=1, + rng=resolve_rng(123), + ) + + with pytest.raises(ValueError, match="indexable sequence"): + runtime = _make_smc().resume_from( + recorder.checkpoints[-1], + rng=resolve_rng(123), + ) + _make_smc().run( + runtime=runtime, + observations=iter([1.0, 2.0]), + num_steps=1, + ) + + +def test_run_rejects_initial_state_and_runtime_together( + particle_state: ParticleState, +) -> None: + step = StochasticStep() + smc = _make_smc(step=step) + runtime = smc.resume_from( + SMCCheckpoint.from_runtime( + state=particle_state, + context=StepContext(t=0, next_observation_index=0), + result=SMCResult( + states=[particle_state], context=StepContext(t=0) + ), + steps=[step], + ) + ) + + with pytest.raises(ValueError, match="either initial_state or runtime"): + _make_smc().run( + particle_state, + runtime=runtime, + num_steps=1, + ) + + +def test_malformed_checkpoint_payload_raises() -> None: + with pytest.raises(ValueError, match="missing required keys"): + SMCCheckpoint.from_dict({"state": {}}) + + with pytest.raises(ValueError, match="Unsupported SMC checkpoint version"): + SMCCheckpoint.from_dict( + { + "version": 99, + "state": {}, + "next_iteration": 1, + "next_observation_index": 0, + "metadata": {}, + "outputs": {}, + "step_states": [], + } + ) + + with pytest.raises(TypeError, match="rng_state"): + SMCCheckpoint.from_dict( + { + "version": 1, + "state": { + "particles": np.zeros((2, 1)), + "log_weights": np.zeros(2), + "log_evidence": 0.0, + "metadata": {}, + }, + "next_iteration": 1, + "next_observation_index": 0, + "metadata": {}, + "outputs": {}, + "step_states": [], + "rng_state": "bad", + "history": None, + } + ) diff --git a/tests/test_posterior.py b/tests/test_posterior.py index ba01e451..69dfb84e 100644 --- a/tests/test_posterior.py +++ b/tests/test_posterior.py @@ -1,5 +1,6 @@ from __future__ import annotations +from typing import Any from unittest.mock import MagicMock import numpy as np @@ -45,9 +46,8 @@ def sample_and_log_prob( def test_flow_proposal_delegates_sampling_and_density(context_factory) -> None: class DummyFlow: - xp = np - - def __init__(self, shift=0.75): + def __init__(self, shift: float = 0.75) -> None: + self.xp: Any = np self.shift = shift self.fitted_particles = None @@ -72,6 +72,12 @@ def log_prob(self, particles, context=None): def sample(self, num_samples: int = 1, context=None): return np.full((num_samples, 2), self.shift, dtype=float) + def state_dict(self): + return {"fitted_particles": self.fitted_particles} + + def load_state_dict(self, state): + self.fitted_particles = np.asarray(state["fitted_particles"]) + training_particles = np.array([[0.0, 1.0], [2.0, 3.0]], dtype=float) proposal = FlowProposal(DummyFlow()).fit( training_particles, @@ -89,6 +95,11 @@ def sample(self, num_samples: int = 1, context=None): sampled_log_prob, proposal.log_prob(sampled_particles, context=context_factory()), ) + restored_proposal = FlowProposal(DummyFlow()) + restored_proposal.load_state_dict(proposal.state_dict()) + np.testing.assert_allclose( + getattr(restored_proposal.flow, "fitted_particles"), training_particles + ) def test_draw_additional_samples_respects_terminal_weights( diff --git a/tests/test_state.py b/tests/test_state.py index 3da8ef9a..35da6061 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -7,14 +7,17 @@ def test_copy_creates_independent_state(particle_state: ParticleState) -> None: + particle_state.metadata["nested"] = {"values": [1.0]} copied = particle_state.copy() copied.particles[0, 0] = 99.0 copied.log_weights[0] = -5.0 copied.metadata["temperature"] = 10.0 + copied.metadata["nested"]["values"].append(2.0) assert not np.array_equal(copied.particles, particle_state.particles) assert not np.array_equal(copied.log_weights, particle_state.log_weights) assert particle_state.metadata["temperature"] == 0.0 + assert particle_state.metadata["nested"] == {"values": [1.0]} def test_normalize_updates_log_weights_and_evidence( diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 181db16f..92bb1c88 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -129,6 +129,12 @@ def log_prob(self, values, context=None): def sample(self, num_samples: int = 1, context=None): return np.full((num_samples, 3), self.shift, dtype=float) + def state_dict(self): + return {"fitted_particles": self.fitted_particles} + + def load_state_dict(self, state): + self.fitted_particles = np.asarray(state["fitted_particles"]) + transform = FlowTransform(ShiftFlow()) transform.fit(particles) latent, log_det = transform.forward(particles) @@ -145,6 +151,12 @@ def sample(self, num_samples: int = 1, context=None): assert cloned is not transform assert cloned.flow is not transform.flow + restored_transform = FlowTransform(ShiftFlow()) + restored_transform.load_state_dict(transform.state_dict()) + np.testing.assert_allclose( + restored_transform.flow.fitted_particles, particles + ) + def test_flow_transform_passes_fit_kwargs(particles): class RecordingFlow: