From af699ed7dc5fac1b7526cd21a85689f349cb4f32 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 7 Jul 2026 15:39:10 +0100 Subject: [PATCH 01/19] ENH: add support for checkpointing --- src/sequince/__init__.py | 5 +- src/sequince/callbacks.py | 64 +++++- src/sequince/checkpoint.py | 121 +++++++++++ src/sequince/flows.py | 29 ++- src/sequince/result.py | 38 +++- src/sequince/smc.py | 160 +++++++++++--- src/sequince/state.py | 34 +++ src/sequince/steps/context.py | 45 ++++ src/sequince/transforms.py | 204 ++++++++++++++++-- src/sequince/utils.py | 394 ++++++++++++++++++++++++++++++++-- 10 files changed, 1027 insertions(+), 67 deletions(-) create mode 100644 src/sequince/checkpoint.py diff --git a/src/sequince/__init__.py b/src/sequince/__init__.py index b3562465..4f85ba62 100644 --- a/src/sequince/__init__.py +++ b/src/sequince/__init__.py @@ -3,10 +3,11 @@ import importlib.metadata import logging +from .checkpoint import SMCCheckpoint from .posterior import TemperedPosterior from .prebuilt_samplers import make_standard_bayesian_smc from .result import SMCResult -from .smc import SequentialMonteCarlo +from .smc import SequentialMonteCarlo, SMCRuntime from .state import ParticleState try: @@ -18,6 +19,8 @@ __all__ = [ "ParticleState", + "SMCCheckpoint", + "SMCRuntime", "SequentialMonteCarlo", "MetaSequentialMonteCarlo", "SMCResult", diff --git a/src/sequince/callbacks.py b/src/sequince/callbacks.py index b8f765e1..d17d671f 100644 --- a/src/sequince/callbacks.py +++ b/src/sequince/callbacks.py @@ -4,7 +4,9 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from .checkpoint import SMCCheckpoint from .steps.mutation import PARTICLE_TRANSFORM_KEY +from .utils import call_with_optional_args if TYPE_CHECKING: from .result import SMCResult @@ -23,14 +25,20 @@ def __call__( state: "ParticleState", context: "StepContext", results: "SMCResult", + rng: Any | None = None, ) -> None: - self.apply(state, context, results) + call_with_optional_args( + self.apply, + (state, context, results), + (rng,), + ) def apply( self, state: "ParticleState", context: "StepContext", results: "SMCResult", + rng: Any | None = None, ) -> None: raise NotImplementedError( "SMCCallback subclasses must implement apply()." @@ -66,6 +74,60 @@ def __init__( Path(self.outdir).mkdir(parents=True, exist_ok=True) +class HDF5Checkpoint(SMCCallback): + """Write SMC checkpoints to an HDF5 file at a fixed iteration cadence.""" + + def __init__( + self, + file_path: str | Path, + *, + every: int = 1, + path: str = "checkpoint", + dsetname: str = "state", + save_rng: bool = True, + ) -> None: + if every <= 0: + raise ValueError("every must be positive.") + file_path = Path(file_path) + lower_path = file_path.name.lower() + if not lower_path.endswith((".h5", ".hdf5")): + raise ValueError( + "Checkpoint file must be an HDF5 file (.h5 or .hdf5)." + ) + self.file_path = file_path + self.every = int(every) + self.path = path + self.dsetname = dsetname + self.save_rng = save_rng + self.file_path.parent.mkdir(parents=True, exist_ok=True) + + def apply( + self, + state: "ParticleState", + context: "StepContext", + results: "SMCResult", + rng: Any | None = None, + ) -> None: + completed_iterations = context.t + 1 + if completed_iterations % self.every != 0: + return + checkpoint_rng = rng if self.save_rng else None + checkpoint = SMCCheckpoint.from_runtime( + state=state, + context=context, + result=results, + rng=checkpoint_rng, + ) + import h5py + + with h5py.File(self.file_path, "a") as h5_file: + checkpoint.save_hdf5( + h5_file, + path=self.path, + dsetname=self.dsetname, + ) + + class TracePlot(PlotCallback): """ Callback to plot 1-d particle locations every SMC iteration. diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py new file mode 100644 index 00000000..bb59dd76 --- /dev/null +++ b/src/sequince/checkpoint.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import copy +import pickle +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .result import SMCResult +from .state import ParticleState +from .steps.context import StepContext +from .utils import restore_rng, snapshot_rng + + +@dataclass(frozen=True) +class SMCCheckpoint: + """Detached runtime snapshot for callback-driven checkpointing.""" + + state: ParticleState + context: StepContext + result: SMCResult + rng_state: dict[str, Any] | None = None + + @classmethod + def from_runtime( + cls, + state: ParticleState, + context: StepContext, + result: SMCResult, + rng: Any | None = None, + ) -> "SMCCheckpoint": + """Capture the current SMC runtime state.""" + context_copy = StepContext.from_snapshot(context.snapshot()) + state_copy = ParticleState.from_snapshot(state.snapshot()) + result_copy = SMCResult( + states=[ + ParticleState.from_snapshot(result_state.snapshot()) + for result_state in result.states + ], + context=context_copy, + ) + return cls( + state=state_copy, + context=context_copy, + result=result_copy, + rng_state=None if rng is None else snapshot_rng(rng), + ) + + def to_dict(self) -> dict[str, Any]: + """Return a plain Python mapping representation of the checkpoint.""" + return { + "state": self.state.snapshot(), + "context": self.context.snapshot(), + "result": self.result.snapshot(), + "rng_state": copy.deepcopy(self.rng_state), + } + + @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.") + missing = { + key for key in ("state", "context", "result") if key not in payload + } + if missing: + raise ValueError( + "SMCCheckpoint payload is missing required keys: " + + ", ".join(sorted(missing)) + ) + 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." + ) + return cls( + state=ParticleState.from_snapshot(copy.deepcopy(payload["state"])), + context=StepContext.from_snapshot( + copy.deepcopy(payload["context"]) + ), + result=SMCResult.from_snapshot(copy.deepcopy(payload["result"])), + rng_state=( + None if rng_state is None else copy.deepcopy(dict(rng_state)) + ), + ) + + def restore_rng(self) -> Any | None: + """Return a live RNG restored from the stored RNG snapshot.""" + if self.rng_state is None: + return None + return restore_rng(self.rng_state) + + def save_hdf5( + self, + h5_file, + path: str = "checkpoint", + dsetname: str = "state", + protocol: int | None = None, + ) -> None: + """Save the checkpoint to an HDF5 dataset as a pickled blob.""" + protocol = ( + pickle.HIGHEST_PROTOCOL if protocol is None else int(protocol) + ) + group = h5_file.require_group(path) + if dsetname in group: + del group[dsetname] + payload = pickle.dumps(self.to_dict(), protocol=protocol) + group.create_dataset(dsetname, data=np.frombuffer(payload, dtype="u1")) + + @classmethod + def load_hdf5( + cls, + h5_file, + path: str = "checkpoint", + dsetname: str = "state", + ) -> "SMCCheckpoint": + """Load a checkpoint from a pickled HDF5 dataset.""" + payload = h5_file[path][dsetname][...].tobytes() + return cls.from_dict(pickle.loads(payload)) diff --git a/src/sequince/flows.py b/src/sequince/flows.py index be193fe2..c6ebceff 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -10,6 +10,8 @@ if TYPE_CHECKING: from .steps.context import StepContext +FLOW_FIT_OUTPUT_KEY = "flow_fit" + class NormalizingFlow(Protocol): """Minimal interface for normalizing flows used by Sequince.""" @@ -21,7 +23,13 @@ def fit( particles: Array, context: StepContext | None = None, **kwargs: Any, - ) -> "NormalizingFlow": ... + ) -> "NormalizingFlow": + """Adapt the flow in place. + + Implementations may record fit diagnostics in ``context.outputs`` under + :data:`FLOW_FIT_OUTPUT_KEY`. + """ + ... def forward( self, particles: Array, context: StepContext | None = None @@ -53,6 +61,25 @@ def flow_asarray(flow: NormalizingFlow, array: Any): return flow.xp.asarray(array) +def record_flow_fit_output( + context: "StepContext | None", + diagnostics: dict[str, Any], + *, + key: str = FLOW_FIT_OUTPUT_KEY, +) -> dict[str, Any]: + """Record flow-fit diagnostics in ``context.outputs`` when available.""" + if context is None: + return diagnostics + record_output = getattr(context, "record_output", None) + if callable(record_output): + record_output(key, diagnostics) + return diagnostics + outputs = getattr(context, "outputs", None) + if outputs is not None: + outputs.setdefault(key, []).append(diagnostics) + return diagnostics + + def clone_flow(flow: NormalizingFlow) -> NormalizingFlow: """Return a fresh copy of ``flow``. diff --git a/src/sequince/result.py b/src/sequince/result.py index b6ab82b0..5f787ccb 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -1,9 +1,11 @@ from __future__ import annotations +import copy import math +from collections.abc import Mapping, 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 +13,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: @@ -53,6 +54,37 @@ def append_state(self, state: ParticleState) -> SMCResult: context=self.context, ) + def snapshot(self) -> dict[str, Any]: + """Return a detached snapshot suitable for checkpoint payloads.""" + return { + "states": [state.snapshot() for state in self.states], + "context": self.context.snapshot(), + } + + @classmethod + def from_snapshot(cls, payload: Mapping[str, Any]) -> "SMCResult": + """Reconstruct a result from :meth:`snapshot` output.""" + if not isinstance(payload, Mapping): + raise TypeError("SMCResult snapshot must be a mapping.") + missing = {key for key in ("states", "context") if key not in payload} + if missing: + raise ValueError( + "SMCResult snapshot is missing required keys: " + + ", ".join(sorted(missing)) + ) + states_payload = payload["states"] + if not isinstance(states_payload, Sequence): + raise TypeError("SMCResult snapshot 'states' must be a sequence.") + return cls( + states=[ + ParticleState.from_snapshot(state_payload) + for state_payload in states_payload + ], + context=StepContext.from_snapshot( + copy.deepcopy(payload["context"]) + ), + ) + def get_nested_results( self, key: str = NESTED_SMC_OUTPUT_KEY, diff --git a/src/sequince/smc.py b/src/sequince/smc.py index b6e317d4..cef6409c 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -1,21 +1,15 @@ from __future__ import annotations 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 SMCCheckpoint from .result import SMCResult from .state import ParticleState from .steps.base import StepContext -from .utils import resolve_rng +from .utils import call_with_optional_args, resolve_rng if TYPE_CHECKING: import numpy as np @@ -24,6 +18,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. @@ -55,10 +63,7 @@ def __init__( ] ], *, - callbacks: Sequence[ - Callable[[ParticleState, StepContext, SMCResult], None] - ] - | None = None, + callbacks: Sequence[Callable[..., None]] | None = None, auto_normalize: bool = True, ): self._steps = list(steps) @@ -67,13 +72,14 @@ def __init__( 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 +107,36 @@ 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." + ) + elif initial_state is None: + raise ValueError("run() requires either initial_state or runtime.") + + if runtime is None: + 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 +149,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 +213,7 @@ def run( context = StepContext( t=t, + next_observation_index=next_observation_index, observation=observation, shared=shared_context, metadata=metadata, @@ -173,6 +240,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,10 +249,10 @@ def run( logger.debug( "Iteration %d: invoking callback %s.", t, callback_name ) - callback( - state, - context, - SMCResult(states=produced_states, context=context), + call_with_optional_args( + callback, + (state, context, result), + (rng,), ) if context.metadata["stop_requested"]: @@ -201,4 +269,34 @@ 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, + ) -> SMCRuntime: + """Reconstruct runtime state from a checkpoint for use with :meth:`run`.""" + checkpoint_obj = ( + checkpoint + if isinstance(checkpoint, SMCCheckpoint) + else SMCCheckpoint.from_dict(checkpoint) + ) + restored_rng = checkpoint_obj.restore_rng() + context_snapshot = checkpoint_obj.context.snapshot() + return SMCRuntime( + state=ParticleState.from_snapshot(checkpoint_obj.state.snapshot()), + produced_states=[ + ParticleState.from_snapshot(state.snapshot()) + for state in checkpoint_obj.result.states + ], + t=checkpoint_obj.context.t + 1, + next_observation_index=checkpoint_obj.context.next_observation_index, + shared_context=context_snapshot["shared"], + metadata=context_snapshot["metadata"], + outputs=context_snapshot["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..ecac4d5d 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 @@ -131,3 +133,35 @@ def resample(self, indices: Array) -> "ParticleState": log_evidence=self.log_evidence, metadata=dict(self.metadata, ancestors=indices), ) + + def snapshot(self) -> dict[str, Any]: + """Return a detached snapshot suitable for checkpoint payloads.""" + xp = namespace(self.particles, self.log_weights) + return { + "particles": array_copy(self.particles, xp=xp), + "log_weights": array_copy(self.log_weights, xp=xp), + "log_evidence": float(self.log_evidence), + "metadata": copy.deepcopy(self.metadata), + } + + @classmethod + def from_snapshot(cls, payload: Mapping[str, Any]) -> "ParticleState": + """Reconstruct a state from :meth:`snapshot` output.""" + if not isinstance(payload, Mapping): + raise TypeError("ParticleState snapshot 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 snapshot 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/context.py b/src/sequince/steps/context.py index 3543e9a4..f7bb9dc7 100644 --- a/src/sequince/steps/context.py +++ b/src/sequince/steps/context.py @@ -1,5 +1,7 @@ from __future__ import annotations +import copy +from collections.abc import Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -20,6 +22,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. @@ -61,3 +64,45 @@ def clear_stop_request(self) -> None: def record_output(self, key: str, value: Any) -> None: """Store an output value that can be accessed by callbacks.""" self.outputs.setdefault(key, []).append(value) + + def snapshot(self) -> dict[str, Any]: + """Return a detached snapshot suitable for checkpoint payloads.""" + return { + "t": int(self.t), + "next_observation_index": int(self.next_observation_index), + "observation": copy.deepcopy(self.observation), + "shared": copy.deepcopy(dict(self.shared)), + "metadata": copy.deepcopy(self.metadata), + "outputs": copy.deepcopy(self.outputs), + } + + @classmethod + def from_snapshot(cls, payload: Mapping[str, Any]) -> "StepContext": + """Reconstruct a context from :meth:`snapshot` output.""" + if not isinstance(payload, Mapping): + raise TypeError("StepContext snapshot must be a mapping.") + missing = { + key + for key in ( + "t", + "next_observation_index", + "observation", + "shared", + "metadata", + "outputs", + ) + if key not in payload + } + if missing: + raise ValueError( + "StepContext snapshot is missing required keys: " + + ", ".join(sorted(missing)) + ) + return cls( + t=int(payload["t"]), + next_observation_index=int(payload["next_observation_index"]), + observation=copy.deepcopy(payload["observation"]), + shared=copy.deepcopy(payload["shared"]), + metadata=copy.deepcopy(payload["metadata"]), + outputs=copy.deepcopy(payload["outputs"]), + ) diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index f489dcb2..e88a4f23 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -2,14 +2,17 @@ from __future__ import annotations +import importlib import logging from dataclasses import dataclass, field from typing import Any, Mapping, Sequence import array_api_extra as xpx +import numpy as np from ._array_api import asarray, namespace, to_numpy from .flows import NormalizingFlow, clone_flow, flow_asarray +from .utils import load_from_h5_file, recursively_save_to_h5_file, to_numpy logger = logging.getLogger(__name__) @@ -30,6 +33,11 @@ def _to_index(columns): """ if columns is None: return None + if hasattr(columns, "ndim"): + arr = np.asarray(columns) + if arr.ndim == 0: + return (int(arr.item()),) + return tuple(int(i) for i in arr.tolist()) if isinstance(columns, Sequence) and not isinstance(columns, (str, bytes)): return tuple(int(i) for i in columns) return (int(columns),) @@ -81,6 +89,20 @@ def _log_det_entries(value, arr, xp): return xp.full((arr.shape[0],), value, dtype=arr.dtype) +def _encode_bounds( + bounds: Mapping[int, tuple[float, float]], +) -> dict[str, list[float]]: + return { + str(key): [float(lo), float(hi)] for key, (lo, hi) in bounds.items() + } + + +def _decode_bounds( + bounds: Mapping[str, Any], +) -> dict[int, tuple[float, float]]: + return {int(key): tuple(value) for key, value in bounds.items()} + + class ParticleTransform: """Minimal interface for coordinate transforms applied to particle states. @@ -120,6 +142,57 @@ def inverse(self, values, context=None): zeros = xp.zeros(arr.shape[0], dtype=arr.dtype) return arr, zeros + def config_dict(self) -> dict[str, Any]: + """Return constructor arguments required to reconstruct the transform.""" + return {} + + def save(self, h5_file, path: str = "particle_transform") -> None: + """Serialize the transform into an HDF5 group.""" + group = h5_file.require_group(path) + group.attrs["class"] = self.__class__.__name__ + group.attrs["module"] = self.__class__.__module__ + recursively_save_to_h5_file(group, "config", self.config_dict()) + self._save_state(group) + + @classmethod + def load( + cls, + h5_file, + path: str = "particle_transform", + strict: bool = False, + ) -> "ParticleTransform": + """Load a transform serialized with :meth:`save`.""" + group = h5_file[path] + class_name = group.attrs["class"] + module_name = group.attrs.get("module", __name__) + + transform_cls = cls + if class_name != cls.__name__ or module_name != cls.__module__: + if strict: + raise ValueError( + f"Expected class {cls.__module__}.{cls.__name__}, got " + f"{module_name}.{class_name}." + ) + module = importlib.import_module(module_name) + transform_cls = getattr(module, class_name) + + return transform_cls._load_from_group(group) + + @classmethod + def _load_from_group(cls, group) -> "ParticleTransform": + config = ( + load_from_h5_file(group, "config") if "config" in group else {} + ) + obj = cls(**config) + obj._load_state(group) + return obj + + def _save_state(self, group) -> None: + """Hook for persisting fitted state.""" + + def _load_state(self, group) -> None: + """Hook for restoring fitted state.""" + class IdentityTransform(ParticleTransform): """No-op transform that simply returns the input particles.""" @@ -151,7 +224,16 @@ def clone(self) -> "FlowTransform": fit_kwargs=self._fit_kwargs, ) - def fit(self, particles, context=None): + def reset(self) -> "FlowTransform": + reset = getattr(self._flow, "reset", None) + if not callable(reset): + raise AttributeError("Underlying flow does not implement reset().") + reset() + return self + + def fit(self, particles, context=None, reset: bool = False): + if reset: + self.reset() backend_particles = flow_asarray(self._flow, particles) logger.debug( "Fitting flow transform to %d particles.", @@ -194,6 +276,17 @@ def inverse(self, values, context=None): self._restore_namespace(log_det, values), ) + def _save_state(self, group) -> None: + self._flow.save(group, "flow") + + @classmethod + def _load_from_group(cls, group) -> "FlowTransform": + flow_group = group["flow"] + module = importlib.import_module(flow_group.attrs["module"]) + flow_cls = getattr(module, flow_group.attrs["class"]) + flow = flow_cls.load(group, "flow") + return cls(flow) + @dataclass class AffineTransform(ParticleTransform): @@ -230,37 +323,68 @@ def fit(self, particles, context=None): return self def forward(self, particles, context=None): - xp = namespace(particles, self._mean, self._scale) + xp = namespace(particles) arr = asarray(particles, xp=xp) + mean = asarray(self._mean, xp=xp, like=arr) + scale = asarray(self._scale, xp=xp, like=arr) if self._column_index is None: - centered = arr - self._mean - scaled = centered / self._scale - log_det = -xp.sum(xp.log(self._scale)) + centered = arr - mean + scaled = centered / scale + log_det = -xp.sum(xp.log(scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return scaled, log_det_vec subset = arr[..., self._column_index] - centered = subset - self._mean - scaled_subset = centered / self._scale + centered = subset - mean + scaled_subset = centered / scale scaled = _replace_columns(arr, self._column_index, scaled_subset) - log_det = -xp.sum(xp.log(self._scale)) + log_det = -xp.sum(xp.log(scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return scaled, log_det_vec def inverse(self, values, context=None): - xp = namespace(values, self._mean, self._scale) + xp = namespace(values) arr = asarray(values, xp=xp) + mean = asarray(self._mean, xp=xp, like=arr) + scale = asarray(self._scale, xp=xp, like=arr) if self._column_index is None: - restored = arr * self._scale + self._mean - log_det = xp.sum(xp.log(self._scale)) + restored = arr * scale + mean + log_det = xp.sum(xp.log(scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return restored, log_det_vec subset = arr[..., self._column_index] - restored_subset = subset * self._scale + self._mean + restored_subset = subset * scale + mean restored = _replace_columns(arr, self._column_index, restored_subset) - log_det = xp.sum(xp.log(self._scale)) + log_det = xp.sum(xp.log(scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return restored, log_det_vec + def config_dict(self) -> dict[str, Any]: + return { + "eps": float(self.eps), + "columns": list(self.columns) + if self.columns is not None + else None, + } + + def _save_state(self, group) -> None: + recursively_save_to_h5_file( + group, + "state", + { + "mean": None if self._mean is None else to_numpy(self._mean), + "scale": None + if self._scale is None + else to_numpy(self._scale), + }, + ) + + def _load_state(self, group) -> None: + state = load_from_h5_file(group, "state") if "state" in group else {} + mean = state.get("mean") + scale = state.get("scale") + self._mean = None if mean is None else np.asarray(mean) + self._scale = None if scale is None else np.asarray(scale) + class BoundedToUnitHypercubeTransform(ParticleTransform): """Linearly rescale specified coordinates into the unit hypercube. @@ -330,6 +454,17 @@ def inverse(self, values, context=None): log_det_vec = _log_det_entries(log_det, arr, xp) return updated, log_det_vec + def config_dict(self) -> dict[str, Any]: + return {"bounds": _encode_bounds(self._bounds)} + + @classmethod + def _load_from_group(cls, group) -> "BoundedToUnitHypercubeTransform": + config = ( + load_from_h5_file(group, "config") if "config" in group else {} + ) + config["bounds"] = _decode_bounds(config["bounds"]) + return cls(**config) + class BoundedToUnboundedTransform(BoundedToUnitHypercubeTransform): """Send bounded coordinates to the real line via an independent logit. @@ -382,6 +517,20 @@ def inverse(self, values, context=None): log_abs_det = base_log_det + jac return restored, log_abs_det + def config_dict(self) -> dict[str, Any]: + return { + "bounds": _encode_bounds(self._bounds), + "eps": float(self._eps), + } + + @classmethod + def _load_from_group(cls, group) -> "BoundedToUnboundedTransform": + config = ( + load_from_h5_file(group, "config") if "config" in group else {} + ) + config["bounds"] = _decode_bounds(config["bounds"]) + return cls(**config) + class PeriodicBoundaryTransform(ParticleTransform): """Wrap selected coordinates so they stay within ``[lower, upper)``. @@ -428,6 +577,17 @@ def forward(self, particles, context=None): def inverse(self, values, context=None): return self.forward(values, context=context) + def config_dict(self) -> dict[str, Any]: + return {"bounds": _encode_bounds(self._bounds)} + + @classmethod + def _load_from_group(cls, group) -> "PeriodicBoundaryTransform": + config = ( + load_from_h5_file(group, "config") if "config" in group else {} + ) + config["bounds"] = _decode_bounds(config["bounds"]) + return cls(**config) + class CompositeTransform(ParticleTransform): """Compose multiple transforms while tracking cumulative log-Jacobians. @@ -493,6 +653,24 @@ def inverse(self, values, context=None): ) return data, total_log_det + def _save_state(self, group) -> None: + transforms_group = group.require_group("transforms") + for name in list(transforms_group.keys()): + del transforms_group[name] + for index, transform in enumerate(self._transforms): + transform.save(transforms_group, str(index)) + + @classmethod + def _load_from_group(cls, group) -> "CompositeTransform": + transforms_group = group["transforms"] + loaded = [ + ParticleTransform.load(transforms_group, str(index)) + for index in sorted(map(int, transforms_group.keys())) + ] + obj = cls.__new__(cls) + obj._transforms = loaded + return obj + _transform_registry = { "identity": IdentityTransform, diff --git a/src/sequince/utils.py b/src/sequince/utils.py index 9e3e49b3..da06047a 100644 --- a/src/sequince/utils.py +++ b/src/sequince/utils.py @@ -1,10 +1,16 @@ from __future__ import annotations +import copy import inspect +import logging +from collections.abc import Mapping +from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, TypeAlias +import h5py import numpy as np -from orng import RandomGenerator +from array_api_compat import is_torch_namespace +from orng import ArrayRNG, RandomGenerator from ._array_api import namespace from ._typing import Array @@ -12,6 +18,7 @@ if TYPE_CHECKING: from .steps.base import StepContext +logger = logging.getLogger(__name__) RNGLike: TypeAlias = RandomGenerator | np.random.Generator | int | None @@ -21,7 +28,7 @@ def resolve_rng( *, backend: str | None = None, ) -> RandomGenerator | np.random.Generator: - """Return the provided RNG or create an ``RandomGenerator`` for the requested backend.""" + """Return the provided RNG or create a generator for the requested backend.""" if isinstance(rng, RandomGenerator): return rng if isinstance(rng, np.random.Generator): @@ -35,12 +42,12 @@ def resolve_rng( ) -def call_with_optional_context( +def call_with_optional_args( fn: Callable[..., Any], base_args: tuple[Any, ...], - context: StepContext, + optional_args: tuple[Any, ...] = (), ) -> Any: - """Call fn with base_args and append context only if required.""" + """Call fn with optional positional arguments when its signature allows it.""" signature = inspect.signature(fn) params = signature.parameters.values() @@ -54,23 +61,201 @@ def call_with_optional_context( ) ] required = sum(1 for p in positional if p.default is inspect._empty) - if required > len(base_args): - return fn(*base_args, context) - has_varargs = any( p.kind == inspect.Parameter.VAR_POSITIONAL for p in params ) - if not has_varargs and len(base_args) > len(positional): - raise TypeError( - "Function cannot accept the provided number of positional arguments." - ) - try: + base_count = len(base_args) + full_count = base_count + len(optional_args) + max_positional = float("inf") if has_varargs else len(positional) + + if required <= full_count <= max_positional: + return fn(*base_args, *optional_args) + 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 _snapshot_numpy_generator_state( + rng: np.random.Generator, +) -> dict[str, Any]: + return { + "kind": "numpy.random.Generator", + "bit_generator_state": copy.deepcopy(rng.bit_generator.state), + } + + +def _restore_numpy_generator(state: Mapping[str, Any]) -> np.random.Generator: + try: + bit_generator_state = copy.deepcopy(state["bit_generator_state"]) + except KeyError as exc: + raise ValueError( + "NumPy RNG snapshot is missing 'bit_generator_state'." + ) from exc + rng = np.random.default_rng() + rng.bit_generator.state = bit_generator_state + return rng + + +def _snapshot_orng_state(backend: str, state: Any) -> dict[str, Any]: + normalized_backend = backend.lower() + if normalized_backend == "numpy": + generator = state + if isinstance(generator, np.random.Generator): + return { + "kind": "numpy", + "bit_generator_state": copy.deepcopy( + generator.bit_generator.state + ), + } + if normalized_backend in {"torch", "pytorch"}: + generator = state + if hasattr(generator, "get_state"): + generator_device = getattr(generator, "device", None) + return { + "kind": "torch", + "generator_state": generator.get_state().clone(), + "device": None + if generator_device is None + else str(generator_device), + } + if normalized_backend == "cupy": + generator = state + if hasattr(generator, "bit_generator"): + return { + "kind": "cupy", + "bit_generator_state": copy.deepcopy( + generator.bit_generator.state + ), + } + if normalized_backend == "jax": + xp = namespace(state) + return { + "kind": "jax", + "key": xp.asarray(state), + } + return { + "kind": normalized_backend, + "state": copy.deepcopy(state), + } + + +def _restore_orng_backend_state( + backend: str, + state: Mapping[str, Any], +) -> Any: + normalized_backend = backend.lower() + state_kind = state.get("kind") + if normalized_backend == "numpy": + if state_kind != "numpy": + raise ValueError(f"Expected NumPy ORNG state, got '{state_kind}'.") + rng = np.random.default_rng() + rng.bit_generator.state = copy.deepcopy(state["bit_generator_state"]) + return rng + if normalized_backend in {"torch", "pytorch"}: + if state_kind != "torch": + raise ValueError(f"Expected Torch ORNG state, got '{state_kind}'.") + try: + import torch + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "Restoring a torch-backed ORNG RNG requires 'torch'." + ) from exc + device = state.get("device") or "cpu" + generator = torch.Generator(device=device) + generator_state = state["generator_state"] + if hasattr(generator_state, "clone"): + generator_state = generator_state.clone() + generator.set_state(generator_state) + return generator + if normalized_backend == "cupy": + if state_kind != "cupy": + raise ValueError(f"Expected CuPy ORNG state, got '{state_kind}'.") + try: + import cupy as cp + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "Restoring a cupy-backed ORNG RNG requires 'cupy'." + ) from exc + generator = cp.random.default_rng() + generator.bit_generator.state = copy.deepcopy( + state["bit_generator_state"] + ) + return generator + if normalized_backend == "jax": + if state_kind != "jax": + raise ValueError(f"Expected JAX ORNG state, got '{state_kind}'.") + return state["key"] + if "state" not in state: + raise ValueError( + f"Unsupported ORNG RNG snapshot payload for backend '{backend}'." + ) + return copy.deepcopy(state["state"]) + + +def snapshot_rng(rng: RandomGenerator | np.random.Generator) -> dict[str, Any]: + """Capture the current RNG state in a plain Python mapping.""" + if isinstance(rng, np.random.Generator): + return _snapshot_numpy_generator_state(rng) + if isinstance(rng, RandomGenerator): + _, backend_state = rng.to_functional(pure=False) + return { + "kind": "orng.RandomGenerator", + "class_name": type(rng).__name__, + "backend": rng.backend, + "device": None if rng.device is None else str(rng.device), + "state": _snapshot_orng_state(rng.backend, backend_state), + } + raise TypeError( + "rng must be a numpy.random.Generator or orng.RandomGenerator." + ) + + +def restore_rng( + snapshot: Mapping[str, Any], +) -> RandomGenerator | np.random.Generator: + """Restore an RNG from a snapshot created by :func:`snapshot_rng`.""" + if not isinstance(snapshot, Mapping): + raise TypeError("RNG snapshot must be a mapping.") + kind = snapshot.get("kind") + if kind == "numpy.random.Generator": + return _restore_numpy_generator(snapshot) + if kind == "orng.RandomGenerator": + try: + backend = snapshot["backend"] + state = snapshot["state"] + except KeyError as exc: + raise ValueError( + "ORNG RNG snapshot must include 'backend' and 'state'." + ) from exc + if not isinstance(state, Mapping): + raise TypeError("ORNG RNG snapshot 'state' must be a mapping.") + rng_type = ( + ArrayRNG + if snapshot.get("class_name") == "ArrayRNG" + else RandomGenerator + ) + generator = _restore_orng_backend_state(backend, state) + device = snapshot.get("device") + kwargs: dict[str, Any] = { + "backend": backend, + "generator": generator, + } + if device is not None: + kwargs["device"] = device + return rng_type(**kwargs) + raise ValueError(f"Unsupported RNG snapshot kind '{kind}'.") + + +def call_with_optional_context( + fn: Callable[..., Any], + base_args: tuple[Any, ...], + context: StepContext, +) -> Any: + """Call fn with base_args and append context only if required.""" + return call_with_optional_args(fn, base_args, (context,)) def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: @@ -102,3 +287,178 @@ def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: "Bounds array must have the same number of dimensions as particles." ) return xp.all((particles >= bounds[0]) & (particles <= bounds[1]), axis=-1) + + +def _dtype_to_name(dtype: Any | str | None) -> str | None: + if dtype is None: + return None + if isinstance(dtype, str): + name = dtype + elif hasattr(dtype, "name") and getattr(dtype, "name"): + name = dtype.name + elif hasattr(dtype, "__name__"): + name = dtype.__name__ + else: + text = str(dtype) + if text.startswith(""): + text = text.split("'")[1] + if text.startswith("dtype(") and text.endswith(")"): + text = text[6:-1].strip("'\" ") + name = text + name = name.split(".")[-1] + return name.strip(" '\"<>").lower() + + +def resolve_dtype(dtype: Any | str | None, xp: Any) -> Any | None: + if dtype is None or xp is None: + return dtype + if isinstance(dtype, str): + dtype_name = _dtype_to_name(dtype) + if is_torch_namespace(xp): + resolved = getattr(xp, dtype_name, None) + if resolved is None: + raise ValueError( + f"Unknown dtype '{dtype}' for namespace {xp.__name__}" + ) + return resolved + try: + return xp.dtype(dtype_name) + except (AttributeError, TypeError, ValueError): + resolved = getattr(xp, dtype_name, None) + if resolved is not None: + return resolved + raise ValueError( + f"Unknown dtype '{dtype}' for namespace " + f"{getattr(xp, '__name__', xp)}" + ) + if is_torch_namespace(xp): + return dtype + try: + return xp.dtype(dtype) + except (AttributeError, TypeError, ValueError): + return dtype + + +def encode_dtype(xp: Any, dtype: Any) -> dict[str, Any] | None: + if dtype is None: + return None + return { + "__dtype__": True, + "xp": xp.__name__, + "dtype": _dtype_to_name(dtype), + } + + +def decode_dtype(xp: Any, encoded_dtype: Any) -> Any: + if isinstance(encoded_dtype, dict) and encoded_dtype.get("__dtype__"): + if encoded_dtype["xp"] != xp.__name__: + raise ValueError( + f"Encoded dtype xp {encoded_dtype['xp']} does not match " + f"current xp {xp.__name__}" + ) + if is_torch_namespace(xp): + return getattr(xp, encoded_dtype["dtype"].split(".")[-1]) + return xp.dtype(encoded_dtype["dtype"].split(".")[-1]) + return encoded_dtype + + +def to_numpy(value: Any) -> Any: + if hasattr(value, "detach") and hasattr(value, "cpu"): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def encode_for_hdf5(value: Any) -> Any: + if hasattr(value, "detach") and hasattr(value, "cpu"): + return value.detach().cpu().numpy() + if isinstance(value, np.ndarray): + return value + if isinstance(value, Path): + return str(value) + if isinstance(value, (int, float, str, bool)): + return value + if isinstance(value, (list, tuple)): + if all(isinstance(v, str) for v in value): + dt = h5py.string_dtype(encoding="utf-8") + return np.array(value, dtype=dt) + return [encode_for_hdf5(v) for v in value] + if isinstance(value, dict): + if not value: + return "__empty_dict__" + return {k: encode_for_hdf5(v) for k, v in value.items()} + if value is None: + return "__none__" + return value + + +def decode_from_hdf5(value: Any) -> Any: + if isinstance(value, bytes): + value = value.decode("utf-8") + if isinstance(value, str): + if value == "__none__": + return None + if value == "__empty_dict__": + return {} + return value + if isinstance(value, np.ndarray): + if value.shape == (): + return value.item() + if value.dtype.kind in {"S", "O", "U"}: + try: + return value.astype(str).tolist() + except Exception: + return value + return value + if isinstance(value, list): + return [decode_from_hdf5(v) for v in value] + if isinstance(value, tuple): + return tuple(decode_from_hdf5(v) for v in value) + if isinstance(value, dict): + return { + ( + k.decode("utf-8") if isinstance(k, bytes) else k + ): decode_from_hdf5(v) + for k, v in value.items() + } + return value + + +def recursively_save_to_h5_file( + h5_file, path: str, dictionary: dict[str, Any] +) -> None: + """Save a nested dictionary to an HDF5 group using flattened keys.""" + group = h5_file.require_group(path) + + def _save_flattened(g, prefix: str, mapping: dict[str, Any]) -> None: + for key, value in mapping.items(): + full_key = f"{prefix}.{key}" if prefix else key + if isinstance(value, dict): + _save_flattened(g, full_key, value) + continue + if full_key in g: + del g[full_key] + try: + g.create_dataset(full_key, data=encode_for_hdf5(value)) + except (TypeError, ValueError) as error: + logger.warning( + "Falling back to string encoding for key %s: %s", + full_key, + error, + ) + dt = h5py.string_dtype(encoding="utf-8") + g.create_dataset(full_key, data=np.array(str(value), dtype=dt)) + + _save_flattened(group, "", dictionary) + + +def load_from_h5_file(h5_file, path: str) -> dict[str, Any]: + """Load a nested dictionary saved by ``recursively_save_to_h5_file``.""" + group = h5_file[path] + result: dict[str, Any] = {} + for key, dataset in group.items(): + parts = key.split(".") + cursor = result + for part in parts[:-1]: + cursor = cursor.setdefault(part, {}) + cursor[parts[-1]] = decode_from_hdf5(dataset[()]) + return result From 2df4f7dc34173ff5af2208eaf5492a8bd3bd4aa5 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 7 Jul 2026 15:39:42 +0100 Subject: [PATCH 02/19] BLD: add h5py dependency --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 91fc2ba0..33323710 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "scipy>=1.8", "matplotlib>=3.5", "numpy>=1.22", + "h5py>=3", "minipcn[array-api]>=0.2.0a4", "array-api-compat>=1.5", "array-api-extra", From c719874e9251ec880f0e1d73dac53745ac27fb9f Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 15:38:51 +0100 Subject: [PATCH 03/19] MAINT: remove unrelated changes --- src/sequince/callbacks.py | 1 - src/sequince/checkpoint.py | 2 +- src/sequince/flows.py | 29 +---- src/sequince/smc.py | 9 +- src/sequince/transforms.py | 206 +++------------------------------- src/sequince/utils.py | 222 +++---------------------------------- 6 files changed, 37 insertions(+), 432 deletions(-) diff --git a/src/sequince/callbacks.py b/src/sequince/callbacks.py index d17d671f..9768b435 100644 --- a/src/sequince/callbacks.py +++ b/src/sequince/callbacks.py @@ -38,7 +38,6 @@ def apply( state: "ParticleState", context: "StepContext", results: "SMCResult", - rng: Any | None = None, ) -> None: raise NotImplementedError( "SMCCallback subclasses must implement apply()." diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py index bb59dd76..175fb3c8 100644 --- a/src/sequince/checkpoint.py +++ b/src/sequince/checkpoint.py @@ -16,7 +16,7 @@ @dataclass(frozen=True) class SMCCheckpoint: - """Detached runtime snapshot for callback-driven checkpointing.""" + """Snapshot of the SMC runtime for checkpointing and resuming.""" state: ParticleState context: StepContext diff --git a/src/sequince/flows.py b/src/sequince/flows.py index c6ebceff..be193fe2 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -10,8 +10,6 @@ if TYPE_CHECKING: from .steps.context import StepContext -FLOW_FIT_OUTPUT_KEY = "flow_fit" - class NormalizingFlow(Protocol): """Minimal interface for normalizing flows used by Sequince.""" @@ -23,13 +21,7 @@ def fit( particles: Array, context: StepContext | None = None, **kwargs: Any, - ) -> "NormalizingFlow": - """Adapt the flow in place. - - Implementations may record fit diagnostics in ``context.outputs`` under - :data:`FLOW_FIT_OUTPUT_KEY`. - """ - ... + ) -> "NormalizingFlow": ... def forward( self, particles: Array, context: StepContext | None = None @@ -61,25 +53,6 @@ def flow_asarray(flow: NormalizingFlow, array: Any): return flow.xp.asarray(array) -def record_flow_fit_output( - context: "StepContext | None", - diagnostics: dict[str, Any], - *, - key: str = FLOW_FIT_OUTPUT_KEY, -) -> dict[str, Any]: - """Record flow-fit diagnostics in ``context.outputs`` when available.""" - if context is None: - return diagnostics - record_output = getattr(context, "record_output", None) - if callable(record_output): - record_output(key, diagnostics) - return diagnostics - outputs = getattr(context, "outputs", None) - if outputs is not None: - outputs.setdefault(key, []).append(diagnostics) - return diagnostics - - def clone_flow(flow: NormalizingFlow) -> NormalizingFlow: """Return a fresh copy of ``flow``. diff --git a/src/sequince/smc.py b/src/sequince/smc.py index cef6409c..93a2fe8e 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -123,10 +123,11 @@ def run( "run(runtime=...) does not accept rng, shared_context, " "metadata, or outputs." ) - elif initial_state is None: - raise ValueError("run() requires either initial_state or runtime.") - - if runtime is None: + else: + if initial_state is None: + raise ValueError( + "run() requires either initial_state or runtime." + ) runtime = SMCRuntime( state=initial_state, produced_states=[], diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index e88a4f23..677807ec 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -2,17 +2,14 @@ from __future__ import annotations -import importlib import logging from dataclasses import dataclass, field -from typing import Any, Mapping, Sequence +from typing import Mapping, Sequence import array_api_extra as xpx -import numpy as np from ._array_api import asarray, namespace, to_numpy from .flows import NormalizingFlow, clone_flow, flow_asarray -from .utils import load_from_h5_file, recursively_save_to_h5_file, to_numpy logger = logging.getLogger(__name__) @@ -33,11 +30,6 @@ def _to_index(columns): """ if columns is None: return None - if hasattr(columns, "ndim"): - arr = np.asarray(columns) - if arr.ndim == 0: - return (int(arr.item()),) - return tuple(int(i) for i in arr.tolist()) if isinstance(columns, Sequence) and not isinstance(columns, (str, bytes)): return tuple(int(i) for i in columns) return (int(columns),) @@ -89,20 +81,6 @@ def _log_det_entries(value, arr, xp): return xp.full((arr.shape[0],), value, dtype=arr.dtype) -def _encode_bounds( - bounds: Mapping[int, tuple[float, float]], -) -> dict[str, list[float]]: - return { - str(key): [float(lo), float(hi)] for key, (lo, hi) in bounds.items() - } - - -def _decode_bounds( - bounds: Mapping[str, Any], -) -> dict[int, tuple[float, float]]: - return {int(key): tuple(value) for key, value in bounds.items()} - - class ParticleTransform: """Minimal interface for coordinate transforms applied to particle states. @@ -142,57 +120,6 @@ def inverse(self, values, context=None): zeros = xp.zeros(arr.shape[0], dtype=arr.dtype) return arr, zeros - def config_dict(self) -> dict[str, Any]: - """Return constructor arguments required to reconstruct the transform.""" - return {} - - def save(self, h5_file, path: str = "particle_transform") -> None: - """Serialize the transform into an HDF5 group.""" - group = h5_file.require_group(path) - group.attrs["class"] = self.__class__.__name__ - group.attrs["module"] = self.__class__.__module__ - recursively_save_to_h5_file(group, "config", self.config_dict()) - self._save_state(group) - - @classmethod - def load( - cls, - h5_file, - path: str = "particle_transform", - strict: bool = False, - ) -> "ParticleTransform": - """Load a transform serialized with :meth:`save`.""" - group = h5_file[path] - class_name = group.attrs["class"] - module_name = group.attrs.get("module", __name__) - - transform_cls = cls - if class_name != cls.__name__ or module_name != cls.__module__: - if strict: - raise ValueError( - f"Expected class {cls.__module__}.{cls.__name__}, got " - f"{module_name}.{class_name}." - ) - module = importlib.import_module(module_name) - transform_cls = getattr(module, class_name) - - return transform_cls._load_from_group(group) - - @classmethod - def _load_from_group(cls, group) -> "ParticleTransform": - config = ( - load_from_h5_file(group, "config") if "config" in group else {} - ) - obj = cls(**config) - obj._load_state(group) - return obj - - def _save_state(self, group) -> None: - """Hook for persisting fitted state.""" - - def _load_state(self, group) -> None: - """Hook for restoring fitted state.""" - class IdentityTransform(ParticleTransform): """No-op transform that simply returns the input particles.""" @@ -224,16 +151,7 @@ def clone(self) -> "FlowTransform": fit_kwargs=self._fit_kwargs, ) - def reset(self) -> "FlowTransform": - reset = getattr(self._flow, "reset", None) - if not callable(reset): - raise AttributeError("Underlying flow does not implement reset().") - reset() - return self - - def fit(self, particles, context=None, reset: bool = False): - if reset: - self.reset() + def fit(self, particles, context=None): backend_particles = flow_asarray(self._flow, particles) logger.debug( "Fitting flow transform to %d particles.", @@ -276,17 +194,6 @@ def inverse(self, values, context=None): self._restore_namespace(log_det, values), ) - def _save_state(self, group) -> None: - self._flow.save(group, "flow") - - @classmethod - def _load_from_group(cls, group) -> "FlowTransform": - flow_group = group["flow"] - module = importlib.import_module(flow_group.attrs["module"]) - flow_cls = getattr(module, flow_group.attrs["class"]) - flow = flow_cls.load(group, "flow") - return cls(flow) - @dataclass class AffineTransform(ParticleTransform): @@ -323,68 +230,37 @@ def fit(self, particles, context=None): return self def forward(self, particles, context=None): - xp = namespace(particles) + xp = namespace(particles, self._mean, self._scale) arr = asarray(particles, xp=xp) - mean = asarray(self._mean, xp=xp, like=arr) - scale = asarray(self._scale, xp=xp, like=arr) if self._column_index is None: - centered = arr - mean - scaled = centered / scale - log_det = -xp.sum(xp.log(scale)) + centered = arr - self._mean + scaled = centered / self._scale + log_det = -xp.sum(xp.log(self._scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return scaled, log_det_vec subset = arr[..., self._column_index] - centered = subset - mean - scaled_subset = centered / scale + centered = subset - self._mean + scaled_subset = centered / self._scale scaled = _replace_columns(arr, self._column_index, scaled_subset) - log_det = -xp.sum(xp.log(scale)) + log_det = -xp.sum(xp.log(self._scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return scaled, log_det_vec def inverse(self, values, context=None): - xp = namespace(values) + xp = namespace(values, self._mean, self._scale) arr = asarray(values, xp=xp) - mean = asarray(self._mean, xp=xp, like=arr) - scale = asarray(self._scale, xp=xp, like=arr) if self._column_index is None: - restored = arr * scale + mean - log_det = xp.sum(xp.log(scale)) + restored = arr * self._scale + self._mean + log_det = xp.sum(xp.log(self._scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return restored, log_det_vec subset = arr[..., self._column_index] - restored_subset = subset * scale + mean + restored_subset = subset * self._scale + self._mean restored = _replace_columns(arr, self._column_index, restored_subset) - log_det = xp.sum(xp.log(scale)) + log_det = xp.sum(xp.log(self._scale)) log_det_vec = _log_det_entries(log_det, arr, xp) return restored, log_det_vec - def config_dict(self) -> dict[str, Any]: - return { - "eps": float(self.eps), - "columns": list(self.columns) - if self.columns is not None - else None, - } - - def _save_state(self, group) -> None: - recursively_save_to_h5_file( - group, - "state", - { - "mean": None if self._mean is None else to_numpy(self._mean), - "scale": None - if self._scale is None - else to_numpy(self._scale), - }, - ) - - def _load_state(self, group) -> None: - state = load_from_h5_file(group, "state") if "state" in group else {} - mean = state.get("mean") - scale = state.get("scale") - self._mean = None if mean is None else np.asarray(mean) - self._scale = None if scale is None else np.asarray(scale) - class BoundedToUnitHypercubeTransform(ParticleTransform): """Linearly rescale specified coordinates into the unit hypercube. @@ -454,17 +330,6 @@ def inverse(self, values, context=None): log_det_vec = _log_det_entries(log_det, arr, xp) return updated, log_det_vec - def config_dict(self) -> dict[str, Any]: - return {"bounds": _encode_bounds(self._bounds)} - - @classmethod - def _load_from_group(cls, group) -> "BoundedToUnitHypercubeTransform": - config = ( - load_from_h5_file(group, "config") if "config" in group else {} - ) - config["bounds"] = _decode_bounds(config["bounds"]) - return cls(**config) - class BoundedToUnboundedTransform(BoundedToUnitHypercubeTransform): """Send bounded coordinates to the real line via an independent logit. @@ -517,20 +382,6 @@ def inverse(self, values, context=None): log_abs_det = base_log_det + jac return restored, log_abs_det - def config_dict(self) -> dict[str, Any]: - return { - "bounds": _encode_bounds(self._bounds), - "eps": float(self._eps), - } - - @classmethod - def _load_from_group(cls, group) -> "BoundedToUnboundedTransform": - config = ( - load_from_h5_file(group, "config") if "config" in group else {} - ) - config["bounds"] = _decode_bounds(config["bounds"]) - return cls(**config) - class PeriodicBoundaryTransform(ParticleTransform): """Wrap selected coordinates so they stay within ``[lower, upper)``. @@ -577,17 +428,6 @@ def forward(self, particles, context=None): def inverse(self, values, context=None): return self.forward(values, context=context) - def config_dict(self) -> dict[str, Any]: - return {"bounds": _encode_bounds(self._bounds)} - - @classmethod - def _load_from_group(cls, group) -> "PeriodicBoundaryTransform": - config = ( - load_from_h5_file(group, "config") if "config" in group else {} - ) - config["bounds"] = _decode_bounds(config["bounds"]) - return cls(**config) - class CompositeTransform(ParticleTransform): """Compose multiple transforms while tracking cumulative log-Jacobians. @@ -653,24 +493,6 @@ def inverse(self, values, context=None): ) return data, total_log_det - def _save_state(self, group) -> None: - transforms_group = group.require_group("transforms") - for name in list(transforms_group.keys()): - del transforms_group[name] - for index, transform in enumerate(self._transforms): - transform.save(transforms_group, str(index)) - - @classmethod - def _load_from_group(cls, group) -> "CompositeTransform": - transforms_group = group["transforms"] - loaded = [ - ParticleTransform.load(transforms_group, str(index)) - for index in sorted(map(int, transforms_group.keys())) - ] - obj = cls.__new__(cls) - obj._transforms = loaded - return obj - _transform_registry = { "identity": IdentityTransform, diff --git a/src/sequince/utils.py b/src/sequince/utils.py index da06047a..971947ec 100644 --- a/src/sequince/utils.py +++ b/src/sequince/utils.py @@ -2,14 +2,10 @@ import copy import inspect -import logging from collections.abc import Mapping -from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, TypeAlias -import h5py import numpy as np -from array_api_compat import is_torch_namespace from orng import ArrayRNG, RandomGenerator from ._array_api import namespace @@ -18,7 +14,6 @@ if TYPE_CHECKING: from .steps.base import StepContext -logger = logging.getLogger(__name__) RNGLike: TypeAlias = RandomGenerator | np.random.Generator | int | None @@ -47,7 +42,7 @@ def call_with_optional_args( base_args: tuple[Any, ...], optional_args: tuple[Any, ...] = (), ) -> Any: - """Call fn with optional positional arguments when its signature allows it.""" + """Call ``fn`` with optional positional arguments when accepted.""" signature = inspect.signature(fn) params = signature.parameters.values() @@ -64,7 +59,6 @@ def call_with_optional_args( has_varargs = any( p.kind == inspect.Parameter.VAR_POSITIONAL for p in params ) - base_count = len(base_args) full_count = base_count + len(optional_args) max_positional = float("inf") if has_varargs else len(positional) @@ -78,6 +72,15 @@ def call_with_optional_args( ) +def call_with_optional_context( + fn: Callable[..., Any], + base_args: tuple[Any, ...], + context: StepContext, +) -> Any: + """Call fn with base_args and append context only if required.""" + return call_with_optional_args(fn, base_args, (context,)) + + def _snapshot_numpy_generator_state( rng: np.random.Generator, ) -> dict[str, Any]: @@ -117,9 +120,9 @@ def _snapshot_orng_state(backend: str, state: Any) -> dict[str, Any]: return { "kind": "torch", "generator_state": generator.get_state().clone(), - "device": None - if generator_device is None - else str(generator_device), + "device": ( + None if generator_device is None else str(generator_device) + ), } if normalized_backend == "cupy": generator = state @@ -132,14 +135,8 @@ def _snapshot_orng_state(backend: str, state: Any) -> dict[str, Any]: } if normalized_backend == "jax": xp = namespace(state) - return { - "kind": "jax", - "key": xp.asarray(state), - } - return { - "kind": normalized_backend, - "state": copy.deepcopy(state), - } + return {"kind": "jax", "key": xp.asarray(state)} + return {"kind": normalized_backend, "state": copy.deepcopy(state)} def _restore_orng_backend_state( @@ -239,25 +236,13 @@ def restore_rng( ) generator = _restore_orng_backend_state(backend, state) device = snapshot.get("device") - kwargs: dict[str, Any] = { - "backend": backend, - "generator": generator, - } + kwargs: dict[str, Any] = {"backend": backend, "generator": generator} if device is not None: kwargs["device"] = device return rng_type(**kwargs) raise ValueError(f"Unsupported RNG snapshot kind '{kind}'.") -def call_with_optional_context( - fn: Callable[..., Any], - base_args: tuple[Any, ...], - context: StepContext, -) -> Any: - """Call fn with base_args and append context only if required.""" - return call_with_optional_args(fn, base_args, (context,)) - - def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: """Check if particles are within the specified bounds. @@ -287,178 +272,3 @@ def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: "Bounds array must have the same number of dimensions as particles." ) return xp.all((particles >= bounds[0]) & (particles <= bounds[1]), axis=-1) - - -def _dtype_to_name(dtype: Any | str | None) -> str | None: - if dtype is None: - return None - if isinstance(dtype, str): - name = dtype - elif hasattr(dtype, "name") and getattr(dtype, "name"): - name = dtype.name - elif hasattr(dtype, "__name__"): - name = dtype.__name__ - else: - text = str(dtype) - if text.startswith(""): - text = text.split("'")[1] - if text.startswith("dtype(") and text.endswith(")"): - text = text[6:-1].strip("'\" ") - name = text - name = name.split(".")[-1] - return name.strip(" '\"<>").lower() - - -def resolve_dtype(dtype: Any | str | None, xp: Any) -> Any | None: - if dtype is None or xp is None: - return dtype - if isinstance(dtype, str): - dtype_name = _dtype_to_name(dtype) - if is_torch_namespace(xp): - resolved = getattr(xp, dtype_name, None) - if resolved is None: - raise ValueError( - f"Unknown dtype '{dtype}' for namespace {xp.__name__}" - ) - return resolved - try: - return xp.dtype(dtype_name) - except (AttributeError, TypeError, ValueError): - resolved = getattr(xp, dtype_name, None) - if resolved is not None: - return resolved - raise ValueError( - f"Unknown dtype '{dtype}' for namespace " - f"{getattr(xp, '__name__', xp)}" - ) - if is_torch_namespace(xp): - return dtype - try: - return xp.dtype(dtype) - except (AttributeError, TypeError, ValueError): - return dtype - - -def encode_dtype(xp: Any, dtype: Any) -> dict[str, Any] | None: - if dtype is None: - return None - return { - "__dtype__": True, - "xp": xp.__name__, - "dtype": _dtype_to_name(dtype), - } - - -def decode_dtype(xp: Any, encoded_dtype: Any) -> Any: - if isinstance(encoded_dtype, dict) and encoded_dtype.get("__dtype__"): - if encoded_dtype["xp"] != xp.__name__: - raise ValueError( - f"Encoded dtype xp {encoded_dtype['xp']} does not match " - f"current xp {xp.__name__}" - ) - if is_torch_namespace(xp): - return getattr(xp, encoded_dtype["dtype"].split(".")[-1]) - return xp.dtype(encoded_dtype["dtype"].split(".")[-1]) - return encoded_dtype - - -def to_numpy(value: Any) -> Any: - if hasattr(value, "detach") and hasattr(value, "cpu"): - return value.detach().cpu().numpy() - return np.asarray(value) - - -def encode_for_hdf5(value: Any) -> Any: - if hasattr(value, "detach") and hasattr(value, "cpu"): - return value.detach().cpu().numpy() - if isinstance(value, np.ndarray): - return value - if isinstance(value, Path): - return str(value) - if isinstance(value, (int, float, str, bool)): - return value - if isinstance(value, (list, tuple)): - if all(isinstance(v, str) for v in value): - dt = h5py.string_dtype(encoding="utf-8") - return np.array(value, dtype=dt) - return [encode_for_hdf5(v) for v in value] - if isinstance(value, dict): - if not value: - return "__empty_dict__" - return {k: encode_for_hdf5(v) for k, v in value.items()} - if value is None: - return "__none__" - return value - - -def decode_from_hdf5(value: Any) -> Any: - if isinstance(value, bytes): - value = value.decode("utf-8") - if isinstance(value, str): - if value == "__none__": - return None - if value == "__empty_dict__": - return {} - return value - if isinstance(value, np.ndarray): - if value.shape == (): - return value.item() - if value.dtype.kind in {"S", "O", "U"}: - try: - return value.astype(str).tolist() - except Exception: - return value - return value - if isinstance(value, list): - return [decode_from_hdf5(v) for v in value] - if isinstance(value, tuple): - return tuple(decode_from_hdf5(v) for v in value) - if isinstance(value, dict): - return { - ( - k.decode("utf-8") if isinstance(k, bytes) else k - ): decode_from_hdf5(v) - for k, v in value.items() - } - return value - - -def recursively_save_to_h5_file( - h5_file, path: str, dictionary: dict[str, Any] -) -> None: - """Save a nested dictionary to an HDF5 group using flattened keys.""" - group = h5_file.require_group(path) - - def _save_flattened(g, prefix: str, mapping: dict[str, Any]) -> None: - for key, value in mapping.items(): - full_key = f"{prefix}.{key}" if prefix else key - if isinstance(value, dict): - _save_flattened(g, full_key, value) - continue - if full_key in g: - del g[full_key] - try: - g.create_dataset(full_key, data=encode_for_hdf5(value)) - except (TypeError, ValueError) as error: - logger.warning( - "Falling back to string encoding for key %s: %s", - full_key, - error, - ) - dt = h5py.string_dtype(encoding="utf-8") - g.create_dataset(full_key, data=np.array(str(value), dtype=dt)) - - _save_flattened(group, "", dictionary) - - -def load_from_h5_file(h5_file, path: str) -> dict[str, Any]: - """Load a nested dictionary saved by ``recursively_save_to_h5_file``.""" - group = h5_file[path] - result: dict[str, Any] = {} - for key, dataset in group.items(): - parts = key.split(".") - cursor = result - for part in parts[:-1]: - cursor = cursor.setdefault(part, {}) - cursor[parts[-1]] = decode_from_hdf5(dataset[()]) - return result From 357b1b3895e1e8a0b0c29b3939d43e18090428a5 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:25:20 +0100 Subject: [PATCH 04/19] ENH: add state_dict/load_state_dict to steps --- src/sequince/steps/base.py | 14 +++++++++++++- src/sequince/steps/control.py | 11 ++++++++++- src/sequince/steps/nested.py | 11 +++++++++++ 3 files changed, 34 insertions(+), 2 deletions(-) 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/control.py b/src/sequince/steps/control.py index e06fb0ec..935fd9fe 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,14 @@ 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) From bbf6a9772f846317dcefb5e317aaa5519bcd53c7 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:29:45 +0100 Subject: [PATCH 05/19] MAINT: detach checkpointing from callbacks --- src/sequince/__init__.py | 5 +- src/sequince/callbacks.py | 63 +--- src/sequince/checkpoint.py | 626 ++++++++++++++++++++++++++++++++++--- src/sequince/smc.py | 81 ++++- 4 files changed, 659 insertions(+), 116 deletions(-) diff --git a/src/sequince/__init__.py b/src/sequince/__init__.py index 4f85ba62..b3562465 100644 --- a/src/sequince/__init__.py +++ b/src/sequince/__init__.py @@ -3,11 +3,10 @@ import importlib.metadata import logging -from .checkpoint import SMCCheckpoint from .posterior import TemperedPosterior from .prebuilt_samplers import make_standard_bayesian_smc from .result import SMCResult -from .smc import SequentialMonteCarlo, SMCRuntime +from .smc import SequentialMonteCarlo from .state import ParticleState try: @@ -19,8 +18,6 @@ __all__ = [ "ParticleState", - "SMCCheckpoint", - "SMCRuntime", "SequentialMonteCarlo", "MetaSequentialMonteCarlo", "SMCResult", diff --git a/src/sequince/callbacks.py b/src/sequince/callbacks.py index 9768b435..b8f765e1 100644 --- a/src/sequince/callbacks.py +++ b/src/sequince/callbacks.py @@ -4,9 +4,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from .checkpoint import SMCCheckpoint from .steps.mutation import PARTICLE_TRANSFORM_KEY -from .utils import call_with_optional_args if TYPE_CHECKING: from .result import SMCResult @@ -25,13 +23,8 @@ def __call__( state: "ParticleState", context: "StepContext", results: "SMCResult", - rng: Any | None = None, ) -> None: - call_with_optional_args( - self.apply, - (state, context, results), - (rng,), - ) + self.apply(state, context, results) def apply( self, @@ -73,60 +66,6 @@ def __init__( Path(self.outdir).mkdir(parents=True, exist_ok=True) -class HDF5Checkpoint(SMCCallback): - """Write SMC checkpoints to an HDF5 file at a fixed iteration cadence.""" - - def __init__( - self, - file_path: str | Path, - *, - every: int = 1, - path: str = "checkpoint", - dsetname: str = "state", - save_rng: bool = True, - ) -> None: - if every <= 0: - raise ValueError("every must be positive.") - file_path = Path(file_path) - lower_path = file_path.name.lower() - if not lower_path.endswith((".h5", ".hdf5")): - raise ValueError( - "Checkpoint file must be an HDF5 file (.h5 or .hdf5)." - ) - self.file_path = file_path - self.every = int(every) - self.path = path - self.dsetname = dsetname - self.save_rng = save_rng - self.file_path.parent.mkdir(parents=True, exist_ok=True) - - def apply( - self, - state: "ParticleState", - context: "StepContext", - results: "SMCResult", - rng: Any | None = None, - ) -> None: - completed_iterations = context.t + 1 - if completed_iterations % self.every != 0: - return - checkpoint_rng = rng if self.save_rng else None - checkpoint = SMCCheckpoint.from_runtime( - state=state, - context=context, - result=results, - rng=checkpoint_rng, - ) - import h5py - - with h5py.File(self.file_path, "a") as h5_file: - checkpoint.save_hdf5( - h5_file, - path=self.path, - dsetname=self.dsetname, - ) - - class TracePlot(PlotCallback): """ Callback to plot 1-d particle locations every SMC iteration. diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py index 175fb3c8..91c737ad 100644 --- a/src/sequince/checkpoint.py +++ b/src/sequince/checkpoint.py @@ -2,26 +2,363 @@ import copy import pickle -from collections.abc import Mapping -from dataclasses import dataclass +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 import numpy as np from .result import SMCResult from .state import ParticleState from .steps.context import StepContext -from .utils import restore_rng, snapshot_rng +from .utils import resolve_rng, restore_rng, snapshot_rng + +CHECKPOINT_VERSION = 1 +_HDF_KIND = "_sequince_kind" +_HDF_KEY_KIND = "_sequince_key_kind" +_HDF_KEY = "_sequince_key" +_HDF_SCALAR_TYPE = "_sequince_scalar_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: + """Recursively write Python containers using native HDF5 structures.""" + if name in parent: + del parent[name] + + if isinstance(value, Mapping): + node = parent.create_group(name) + node.attrs[_HDF_KIND] = "mapping" + for index, (key, item) in enumerate(value.items()): + child_name = ( + quote(key, safe="") or "%" + if isinstance(key, str) + else f"@{index:08d}" + ) + child = _write_hdf_value(node, child_name, item, protocol) + if isinstance(key, str): + try: + child.attrs[_HDF_KEY] = key + child.attrs[_HDF_KEY_KIND] = "string" + continue + except (TypeError, ValueError): + if _HDF_KEY in child.attrs: + del child.attrs[_HDF_KEY] + child.attrs[_HDF_KEY] = np.void( + pickle.dumps(key, protocol=protocol) + ) + child.attrs[_HDF_KEY_KIND] = "pickle" + return node + + if isinstance(value, (list, tuple, set, frozenset)): + node = parent.create_group(name) + container_kind = { + list: "list", + tuple: "tuple", + set: "set", + frozenset: "frozenset", + } + node.attrs[_HDF_KIND] = container_kind[type(value)] + 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): + 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] = "array" + return node + + if isinstance(value, np.generic): + 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_scalar" + return node + + if isinstance(value, (bool, int, float, complex)): + try: + node = parent.create_dataset(name, data=value) + except OverflowError: + if isinstance(value, int): + node = parent.create_dataset( + name, + data=str(value), + dtype="S" + str(len(str(value))), + ) + node.attrs[_HDF_KIND] = "bigint" + return node + return _write_pickled_hdf_value(parent, name, value, protocol) + except (TypeError, ValueError): + return _write_pickled_hdf_value(parent, name, value, protocol) + node.attrs[_HDF_KIND] = "scalar" + node.attrs[_HDF_SCALAR_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": + mapping_value: dict[Any, Any] = {} + for child_name in sorted(node.keys()): + child = node[child_name] + key_kind = _attr_text(child.attrs[_HDF_KEY_KIND]) + if key_kind == "string": + key = child.attrs[_HDF_KEY] + if isinstance(key, bytes): + key = key.decode("utf-8") + elif key_kind == "pickle": + key = pickle.loads(bytes(child.attrs[_HDF_KEY])) + else: + raise ValueError( + f"Unsupported HDF5 mapping key kind {key_kind!r}." + ) + mapping_value[key] = _read_hdf_value(child) + return mapping_value + + if kind in ("list", "tuple", "set", "frozenset"): + items = [_read_hdf_value(node[name]) for name in sorted(node.keys())] + container_constructors = { + "list": list, + "tuple": tuple, + "set": set, + "frozenset": frozenset, + } + return container_constructors[kind](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 == "array": + return node[...] + if kind == "numpy_scalar": + return node[()] + if kind == "bigint": + return int(node[()].decode("ascii")) + if kind == "scalar": + scalar_value = node[()] + if isinstance(scalar_value, np.generic): + scalar_value = scalar_value.item() + scalar_type = _attr_text(node.attrs[_HDF_SCALAR_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 snapshot_steps(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: - """Snapshot of the SMC runtime for checkpointing and resuming.""" + """Detached end-of-iteration state for resuming an SMC run.""" state: ParticleState - context: StepContext - result: SMCResult + 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 + version: int = CHECKPOINT_VERSION @classmethod def from_runtime( @@ -29,32 +366,52 @@ def from_runtime( state: ParticleState, context: StepContext, result: SMCResult, + steps: Sequence[Any], rng: Any | None = None, + *, + include_history: bool = True, ) -> "SMCCheckpoint": - """Capture the current SMC runtime state.""" - context_copy = StepContext.from_snapshot(context.snapshot()) + """Capture a checkpoint after a completed SMC iteration.""" state_copy = ParticleState.from_snapshot(state.snapshot()) - result_copy = SMCResult( - states=[ - ParticleState.from_snapshot(result_state.snapshot()) - for result_state in result.states - ], - context=context_copy, - ) + history = None + if include_history: + history = tuple( + ParticleState.from_snapshot(item.snapshot()) + for item in result.states + ) return cls( state=state_copy, - context=context_copy, - result=result_copy, - rng_state=None if rng is None else snapshot_rng(rng), + 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=snapshot_steps(steps), + rng_state=None if rng is None else snapshot_rng(resolve_rng(rng)), + history=history, ) def to_dict(self) -> dict[str, Any]: - """Return a plain Python mapping representation of the checkpoint.""" + """Return a detached plain-Python checkpoint representation.""" return { + "version": self.version, "state": self.state.snapshot(), - "context": self.context.snapshot(), - "result": self.result.snapshot(), + "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.snapshot() for item in self.history] + ), } @classmethod @@ -62,60 +419,259 @@ 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.") - missing = { - key for key in ("state", "context", "result") if key not in payload + 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_snapshot(copy.deepcopy(payload["state"])), - context=StepContext.from_snapshot( - copy.deepcopy(payload["context"]) + 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 ), - result=SMCResult.from_snapshot(copy.deepcopy(payload["result"])), 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_snapshot(copy.deepcopy(item)) + for item in history_payload + ) + ), + version=version, ) def restore_rng(self) -> Any | None: - """Return a live RNG restored from the stored RNG snapshot.""" + """Return a live ORNG generator restored from this checkpoint.""" if self.rng_state is None: return None return restore_rng(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, + 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 = "state", + dsetname: str = "runtime", protocol: int | None = None, + *, + replace_history: bool = False, ) -> None: - """Save the checkpoint to an HDF5 dataset as a pickled blob.""" + """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"] + runtime_payload["history_length"] = 0 + 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].snapshot(), + protocol, + ) + history_group.attrs["length"] = len(self.history) + runtime_payload["history_length"] = len(self.history) + + temporary_name = f"{dsetname}.tmp" + _write_hdf_value(group, temporary_name, runtime_payload, protocol) if dsetname in group: del group[dsetname] - payload = pickle.dumps(self.to_dict(), protocol=protocol) - group.create_dataset(dsetname, data=np.frombuffer(payload, dtype="u1")) + group.move(temporary_name, dsetname) @classmethod def load_hdf5( cls, h5_file, path: str = "checkpoint", - dsetname: str = "state", + dsetname: str = "runtime", ) -> "SMCCheckpoint": - """Load a checkpoint from a pickled HDF5 dataset.""" - payload = h5_file[path][dsetname][...].tobytes() - return cls.from_dict(pickle.loads(payload)) + """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.") + history_length = int(checkpoint_payload.pop("history_length", 0)) + if history_length: + if "history" not in group: + raise ValueError("Checkpoint history is missing from HDF5.") + history_group = group["history"] + checkpoint_payload["history"] = [ + _read_hdf_value(history_group[f"{index:08d}"]) + for index in range(history_length) + ] + 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._has_written = False + 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 + + with h5py.File(self.file_path, "a") as h5_file: + checkpoint.save_hdf5( + h5_file, + path=self.path, + dsetname=self.dsetname, + replace_history=not self._has_written, + ) + self._has_written = True + + +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/smc.py b/src/sequince/smc.py index 93a2fe8e..74aef296 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -1,15 +1,22 @@ from __future__ import annotations +import copy import logging from collections.abc import Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable -from .checkpoint import SMCCheckpoint +from .checkpoint import ( + CheckpointWriter, + SMCCheckpoint, + StepState, + restore_steps, + snapshot_steps, +) from .result import SMCResult from .state import ParticleState from .steps.base import StepContext -from .utils import call_with_optional_args, resolve_rng +from .utils import resolve_rng if TYPE_CHECKING: import numpy as np @@ -49,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. @@ -64,12 +74,30 @@ def __init__( ], *, 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 snapshot_steps(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 | None = None, @@ -250,11 +278,31 @@ def run( logger.debug( "Iteration %d: invoking callback %s.", t, callback_name ) - call_with_optional_args( - callback, - (state, context, result), - (rng,), + 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( @@ -279,6 +327,7 @@ def resume_from( 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 = ( @@ -286,18 +335,20 @@ def resume_from( if isinstance(checkpoint, SMCCheckpoint) else SMCCheckpoint.from_dict(checkpoint) ) + restore_steps(self._steps, checkpoint_obj.step_states) restored_rng = checkpoint_obj.restore_rng() - context_snapshot = checkpoint_obj.context.snapshot() + state = ParticleState.from_snapshot(checkpoint_obj.state.snapshot()) + history = checkpoint_obj.history or (checkpoint_obj.state,) return SMCRuntime( - state=ParticleState.from_snapshot(checkpoint_obj.state.snapshot()), + state=state, produced_states=[ - ParticleState.from_snapshot(state.snapshot()) - for state in checkpoint_obj.result.states + ParticleState.from_snapshot(item.snapshot()) + for item in history ], - t=checkpoint_obj.context.t + 1, - next_observation_index=checkpoint_obj.context.next_observation_index, - shared_context=context_snapshot["shared"], - metadata=context_snapshot["metadata"], - outputs=context_snapshot["outputs"], + t=checkpoint_obj.next_iteration, + next_observation_index=checkpoint_obj.next_observation_index, + shared_context=dict(shared_context or {}), + metadata=copy.deepcopy(checkpoint_obj.metadata), + outputs=copy.deepcopy(checkpoint_obj.outputs), rng=restored_rng if restored_rng is not None else resolve_rng(rng), ) From 5391a28f45869a4cd12eee45d8a078be5c9cf60e Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:30:03 +0100 Subject: [PATCH 06/19] MAINT: use RNG encoding from orng --- src/sequince/utils.py | 183 ++++-------------------------------------- 1 file changed, 15 insertions(+), 168 deletions(-) diff --git a/src/sequince/utils.py b/src/sequince/utils.py index 971947ec..f3707687 100644 --- a/src/sequince/utils.py +++ b/src/sequince/utils.py @@ -1,12 +1,11 @@ from __future__ import annotations -import copy import inspect from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Callable, TypeAlias import numpy as np -from orng import ArrayRNG, RandomGenerator +from orng import RandomGenerator from ._array_api import namespace from ._typing import Array @@ -22,18 +21,19 @@ def resolve_rng( rng: RNGLike = None, *, backend: str | None = None, -) -> RandomGenerator | np.random.Generator: - """Return the provided RNG or create a generator 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." ) @@ -60,13 +60,12 @@ def call_with_optional_args( p.kind == inspect.Parameter.VAR_POSITIONAL for p in params ) base_count = len(base_args) - full_count = base_count + len(optional_args) max_positional = float("inf") if has_varargs else len(positional) - if required <= full_count <= max_positional: - return fn(*base_args, *optional_args) - if required <= base_count <= max_positional: - return fn(*base_args) + for optional_count in range(len(optional_args), -1, -1): + argument_count = base_count + optional_count + if required <= argument_count <= max_positional: + return fn(*base_args, *optional_args[:optional_count]) raise TypeError( "Function cannot accept the provided positional arguments." ) @@ -81,166 +80,14 @@ def call_with_optional_context( return call_with_optional_args(fn, base_args, (context,)) -def _snapshot_numpy_generator_state( - rng: np.random.Generator, -) -> dict[str, Any]: - return { - "kind": "numpy.random.Generator", - "bit_generator_state": copy.deepcopy(rng.bit_generator.state), - } - - -def _restore_numpy_generator(state: Mapping[str, Any]) -> np.random.Generator: - try: - bit_generator_state = copy.deepcopy(state["bit_generator_state"]) - except KeyError as exc: - raise ValueError( - "NumPy RNG snapshot is missing 'bit_generator_state'." - ) from exc - rng = np.random.default_rng() - rng.bit_generator.state = bit_generator_state - return rng - - -def _snapshot_orng_state(backend: str, state: Any) -> dict[str, Any]: - normalized_backend = backend.lower() - if normalized_backend == "numpy": - generator = state - if isinstance(generator, np.random.Generator): - return { - "kind": "numpy", - "bit_generator_state": copy.deepcopy( - generator.bit_generator.state - ), - } - if normalized_backend in {"torch", "pytorch"}: - generator = state - if hasattr(generator, "get_state"): - generator_device = getattr(generator, "device", None) - return { - "kind": "torch", - "generator_state": generator.get_state().clone(), - "device": ( - None if generator_device is None else str(generator_device) - ), - } - if normalized_backend == "cupy": - generator = state - if hasattr(generator, "bit_generator"): - return { - "kind": "cupy", - "bit_generator_state": copy.deepcopy( - generator.bit_generator.state - ), - } - if normalized_backend == "jax": - xp = namespace(state) - return {"kind": "jax", "key": xp.asarray(state)} - return {"kind": normalized_backend, "state": copy.deepcopy(state)} - - -def _restore_orng_backend_state( - backend: str, - state: Mapping[str, Any], -) -> Any: - normalized_backend = backend.lower() - state_kind = state.get("kind") - if normalized_backend == "numpy": - if state_kind != "numpy": - raise ValueError(f"Expected NumPy ORNG state, got '{state_kind}'.") - rng = np.random.default_rng() - rng.bit_generator.state = copy.deepcopy(state["bit_generator_state"]) - return rng - if normalized_backend in {"torch", "pytorch"}: - if state_kind != "torch": - raise ValueError(f"Expected Torch ORNG state, got '{state_kind}'.") - try: - import torch - except ImportError as exc: # pragma: no cover - optional dependency - raise ImportError( - "Restoring a torch-backed ORNG RNG requires 'torch'." - ) from exc - device = state.get("device") or "cpu" - generator = torch.Generator(device=device) - generator_state = state["generator_state"] - if hasattr(generator_state, "clone"): - generator_state = generator_state.clone() - generator.set_state(generator_state) - return generator - if normalized_backend == "cupy": - if state_kind != "cupy": - raise ValueError(f"Expected CuPy ORNG state, got '{state_kind}'.") - try: - import cupy as cp - except ImportError as exc: # pragma: no cover - optional dependency - raise ImportError( - "Restoring a cupy-backed ORNG RNG requires 'cupy'." - ) from exc - generator = cp.random.default_rng() - generator.bit_generator.state = copy.deepcopy( - state["bit_generator_state"] - ) - return generator - if normalized_backend == "jax": - if state_kind != "jax": - raise ValueError(f"Expected JAX ORNG state, got '{state_kind}'.") - return state["key"] - if "state" not in state: - raise ValueError( - f"Unsupported ORNG RNG snapshot payload for backend '{backend}'." - ) - return copy.deepcopy(state["state"]) - - -def snapshot_rng(rng: RandomGenerator | np.random.Generator) -> dict[str, Any]: - """Capture the current RNG state in a plain Python mapping.""" - if isinstance(rng, np.random.Generator): - return _snapshot_numpy_generator_state(rng) - if isinstance(rng, RandomGenerator): - _, backend_state = rng.to_functional(pure=False) - return { - "kind": "orng.RandomGenerator", - "class_name": type(rng).__name__, - "backend": rng.backend, - "device": None if rng.device is None else str(rng.device), - "state": _snapshot_orng_state(rng.backend, backend_state), - } - raise TypeError( - "rng must be a numpy.random.Generator or orng.RandomGenerator." - ) +def snapshot_rng(rng: RandomGenerator) -> dict[str, Any]: + """Capture RNG state through ORNG's public state protocol.""" + return rng.state_dict() -def restore_rng( - snapshot: Mapping[str, Any], -) -> RandomGenerator | np.random.Generator: +def restore_rng(snapshot: Mapping[str, Any]) -> RandomGenerator: """Restore an RNG from a snapshot created by :func:`snapshot_rng`.""" - if not isinstance(snapshot, Mapping): - raise TypeError("RNG snapshot must be a mapping.") - kind = snapshot.get("kind") - if kind == "numpy.random.Generator": - return _restore_numpy_generator(snapshot) - if kind == "orng.RandomGenerator": - try: - backend = snapshot["backend"] - state = snapshot["state"] - except KeyError as exc: - raise ValueError( - "ORNG RNG snapshot must include 'backend' and 'state'." - ) from exc - if not isinstance(state, Mapping): - raise TypeError("ORNG RNG snapshot 'state' must be a mapping.") - rng_type = ( - ArrayRNG - if snapshot.get("class_name") == "ArrayRNG" - else RandomGenerator - ) - generator = _restore_orng_backend_state(backend, state) - device = snapshot.get("device") - kwargs: dict[str, Any] = {"backend": backend, "generator": generator} - if device is not None: - kwargs["device"] = device - return rng_type(**kwargs) - raise ValueError(f"Unsupported RNG snapshot kind '{kind}'.") + return RandomGenerator.from_state_dict(snapshot) def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: From e0035807abaa37e0376b5a0366627d6b0f25f0a9 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:30:12 +0100 Subject: [PATCH 07/19] BLD: bump orng version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 33323710..609714dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "minipcn[array-api]>=0.2.0a4", "array-api-compat>=1.5", "array-api-extra", - "orng>=0.1.0a3", + "orng>=0.2", "frozendict", ] dynamic = [ From b2c42be2e1f179d4110393812daa3768a7e1f40b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:30:39 +0100 Subject: [PATCH 08/19] ENH: add state dict methods to flows --- src/sequince/flows.py | 34 +++++++++++++++++++++++- src/sequince/proposals.py | 53 +++++++++++++++++++++++++++++++++----- src/sequince/transforms.py | 19 ++++++++++++-- 3 files changed, 97 insertions(+), 9 deletions(-) diff --git a/src/sequince/flows.py b/src/sequince/flows.py index be193fe2..c65d3b2e 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -3,11 +3,14 @@ from __future__ import annotations import copy +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Protocol from ._typing import Array if TYPE_CHECKING: + from orng import RandomGenerator + from .steps.context import StepContext @@ -20,6 +23,7 @@ def fit( self, particles: Array, context: StepContext | None = None, + rng: RandomGenerator | None = None, **kwargs: Any, ) -> "NormalizingFlow": ... @@ -42,7 +46,10 @@ def log_prob( ... def sample( - self, num_samples: int = 1, context: StepContext | None = None + self, + num_samples: int = 1, + context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> Array: """Draw samples from the flow distribution.""" ... @@ -53,6 +60,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..e78e6e10 100644 --- a/src/sequince/proposals.py +++ b/src/sequince/proposals.py @@ -1,13 +1,22 @@ 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 orng import RandomGenerator + from .steps.context import StepContext logger = logging.getLogger(__name__) @@ -27,6 +36,7 @@ def sample( self, n_samples: int, context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> Array: """Draw samples from the proposal distribution.""" @@ -34,6 +44,7 @@ def sample_and_log_prob( self, n_samples: int, context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> tuple[Array, Array]: """Draw samples and return their log proposal densities.""" @@ -55,12 +66,28 @@ def fit( self, particles: Array, context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> "FlowProposal": logger.info("Fitting flow proposal to %d particles", len(particles)) backend_particles = flow_asarray(self._flow, particles) - self._flow.fit(backend_particles, context=context) + from .utils import call_with_optional_args + + call_with_optional_args( + self._flow.fit, + (backend_particles,), + (context, rng), + ) 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, @@ -76,18 +103,32 @@ def sample( self, n_samples: int, context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> Array: - return self._flow.sample(n_samples, context=context) + from .utils import call_with_optional_args + + return call_with_optional_args( + self._flow.sample, + (n_samples,), + (context, rng), + ) def sample_and_log_prob( self, n_samples: int, context: StepContext | None = None, + rng: RandomGenerator | None = None, ) -> tuple[Array, Array]: flow_sample_and_log_prob = getattr( self._flow, "sample_and_log_prob", None ) if callable(flow_sample_and_log_prob): - return flow_sample_and_log_prob(n_samples, context=context) - samples = self.sample(n_samples, context=context) + from .utils import call_with_optional_args + + return call_with_optional_args( + flow_sample_and_log_prob, + (n_samples,), + (context, rng), + ) + samples = self.sample(n_samples, context=context, rng=rng) return samples, self.log_prob(samples, context=context) diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index 677807ec..7e506fcd 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -4,12 +4,18 @@ import logging from dataclasses import dataclass, field -from typing import Mapping, Sequence +from typing import Any, Mapping, Sequence 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( From ba709d25587815d563b69a1fc490ba17b797917c Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:34:16 +0100 Subject: [PATCH 09/19] MAINT: remove unintentional API changes --- src/sequince/flows.py | 4 ---- src/sequince/proposals.py | 34 ++++++---------------------------- 2 files changed, 6 insertions(+), 32 deletions(-) diff --git a/src/sequince/flows.py b/src/sequince/flows.py index c65d3b2e..46051953 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -9,8 +9,6 @@ from ._typing import Array if TYPE_CHECKING: - from orng import RandomGenerator - from .steps.context import StepContext @@ -23,7 +21,6 @@ def fit( self, particles: Array, context: StepContext | None = None, - rng: RandomGenerator | None = None, **kwargs: Any, ) -> "NormalizingFlow": ... @@ -49,7 +46,6 @@ def sample( self, num_samples: int = 1, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> Array: """Draw samples from the flow distribution.""" ... diff --git a/src/sequince/proposals.py b/src/sequince/proposals.py index e78e6e10..369e2c61 100644 --- a/src/sequince/proposals.py +++ b/src/sequince/proposals.py @@ -15,8 +15,6 @@ ) if TYPE_CHECKING: - from orng import RandomGenerator - from .steps.context import StepContext logger = logging.getLogger(__name__) @@ -36,7 +34,6 @@ def sample( self, n_samples: int, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> Array: """Draw samples from the proposal distribution.""" @@ -44,7 +41,6 @@ def sample_and_log_prob( self, n_samples: int, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> tuple[Array, Array]: """Draw samples and return their log proposal densities.""" @@ -66,17 +62,10 @@ def fit( self, particles: Array, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> "FlowProposal": logger.info("Fitting flow proposal to %d particles", len(particles)) backend_particles = flow_asarray(self._flow, particles) - from .utils import call_with_optional_args - - call_with_optional_args( - self._flow.fit, - (backend_particles,), - (context, rng), - ) + self._flow.fit(backend_particles, context=context) return self def state_dict(self) -> dict[str, Any]: @@ -103,32 +92,21 @@ def sample( self, n_samples: int, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> Array: - from .utils import call_with_optional_args - - return call_with_optional_args( - self._flow.sample, - (n_samples,), - (context, rng), - ) + return self._flow.sample(n_samples, context=context) def sample_and_log_prob( self, n_samples: int, context: StepContext | None = None, - rng: RandomGenerator | None = None, ) -> tuple[Array, Array]: flow_sample_and_log_prob = getattr( self._flow, "sample_and_log_prob", None ) if callable(flow_sample_and_log_prob): - from .utils import call_with_optional_args - - return call_with_optional_args( - flow_sample_and_log_prob, - (n_samples,), - (context, rng), + return flow_sample_and_log_prob( + n_samples, + context=context, ) - samples = self.sample(n_samples, context=context, rng=rng) + samples = self.sample(n_samples, context=context) return samples, self.log_prob(samples, context=context) From 379a960bfbce82eb53e3972f2d2fcb5c6a8d91c4 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 09:50:28 +0100 Subject: [PATCH 10/19] MAINT: remove unused function --- src/sequince/utils.py | 51 ++++++++++++++----------------------------- 1 file changed, 16 insertions(+), 35 deletions(-) diff --git a/src/sequince/utils.py b/src/sequince/utils.py index f3707687..694c755c 100644 --- a/src/sequince/utils.py +++ b/src/sequince/utils.py @@ -1,7 +1,6 @@ from __future__ import annotations import inspect -from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Callable, TypeAlias import numpy as np @@ -37,59 +36,41 @@ def resolve_rng( ) -def call_with_optional_args( +def call_with_optional_context( fn: Callable[..., Any], base_args: tuple[Any, ...], - optional_args: tuple[Any, ...] = (), + context: StepContext, ) -> Any: - """Call ``fn`` with optional positional arguments when accepted.""" + """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) + 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 ) - base_count = len(base_args) max_positional = float("inf") if has_varargs else len(positional) - - for optional_count in range(len(optional_args), -1, -1): - argument_count = base_count + optional_count - if required <= argument_count <= max_positional: - return fn(*base_args, *optional_args[:optional_count]) + 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) raise TypeError( "Function cannot accept the provided positional arguments." ) -def call_with_optional_context( - fn: Callable[..., Any], - base_args: tuple[Any, ...], - context: StepContext, -) -> Any: - """Call fn with base_args and append context only if required.""" - return call_with_optional_args(fn, base_args, (context,)) - - -def snapshot_rng(rng: RandomGenerator) -> dict[str, Any]: - """Capture RNG state through ORNG's public state protocol.""" - return rng.state_dict() - - -def restore_rng(snapshot: Mapping[str, Any]) -> RandomGenerator: - """Restore an RNG from a snapshot created by :func:`snapshot_rng`.""" - return RandomGenerator.from_state_dict(snapshot) - - def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: """Check if particles are within the specified bounds. From 096bff07009ac97ab6182fafd2c1fe22f8f8d73b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 09:50:47 +0100 Subject: [PATCH 11/19] MAINT: simplify checkpointing --- src/sequince/checkpoint.py | 28 +++++++++++----------- src/sequince/result.py | 34 +-------------------------- src/sequince/smc.py | 13 +++++------ src/sequince/state.py | 24 +++++++++---------- src/sequince/steps/context.py | 44 ----------------------------------- 5 files changed, 32 insertions(+), 111 deletions(-) diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py index 91c737ad..042802cb 100644 --- a/src/sequince/checkpoint.py +++ b/src/sequince/checkpoint.py @@ -10,11 +10,12 @@ from urllib.parse import quote 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, restore_rng, snapshot_rng +from .utils import resolve_rng CHECKPOINT_VERSION = 1 _HDF_KIND = "_sequince_kind" @@ -310,7 +311,7 @@ def from_dict(cls, payload: Mapping[str, Any]) -> "StepState": ) -def snapshot_steps(steps: Sequence[Any]) -> tuple[StepState, ...]: +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) @@ -372,13 +373,10 @@ def from_runtime( include_history: bool = True, ) -> "SMCCheckpoint": """Capture a checkpoint after a completed SMC iteration.""" - state_copy = ParticleState.from_snapshot(state.snapshot()) + state_copy = state.copy() history = None if include_history: - history = tuple( - ParticleState.from_snapshot(item.snapshot()) - for item in result.states - ) + history = tuple(item.copy() for item in result.states) return cls( state=state_copy, next_iteration=context.t + 1, @@ -391,8 +389,8 @@ def from_runtime( } ), outputs=copy.deepcopy(dict(context.outputs)), - step_states=snapshot_steps(steps), - rng_state=None if rng is None else snapshot_rng(resolve_rng(rng)), + step_states=capture_step_states(steps), + rng_state=None if rng is None else resolve_rng(rng).state_dict(), history=history, ) @@ -400,7 +398,7 @@ def to_dict(self) -> dict[str, Any]: """Return a detached plain-Python checkpoint representation.""" return { "version": self.version, - "state": self.state.snapshot(), + "state": self.state.to_dict(), "next_iteration": self.next_iteration, "next_observation_index": self.next_observation_index, "metadata": copy.deepcopy(self.metadata), @@ -410,7 +408,7 @@ def to_dict(self) -> dict[str, Any]: "history": ( None if self.history is None - else [item.snapshot() for item in self.history] + else [item.to_dict() for item in self.history] ), } @@ -469,7 +467,7 @@ def from_dict(cls, payload: Mapping[str, Any]) -> "SMCCheckpoint": "non-negative." ) return cls( - state=ParticleState.from_snapshot(copy.deepcopy(payload["state"])), + state=ParticleState.from_dict(copy.deepcopy(payload["state"])), next_iteration=next_iteration, next_observation_index=next_observation_index, metadata=copy.deepcopy(dict(payload["metadata"])), @@ -484,7 +482,7 @@ def from_dict(cls, payload: Mapping[str, Any]) -> "SMCCheckpoint": None if history_payload is None else tuple( - ParticleState.from_snapshot(copy.deepcopy(item)) + ParticleState.from_dict(copy.deepcopy(item)) for item in history_payload ) ), @@ -495,7 +493,7 @@ def restore_rng(self) -> Any | None: """Return a live ORNG generator restored from this checkpoint.""" if self.rng_state is None: return None - return restore_rng(self.rng_state) + return RandomGenerator.from_state_dict(self.rng_state) def result( self, shared_context: Mapping[str, Any] | None = None @@ -565,7 +563,7 @@ def save_hdf5( _write_hdf_value( history_group, f"{index:08d}", - self.history[index].snapshot(), + self.history[index].to_dict(), protocol, ) history_group.attrs["length"] = len(self.history) diff --git a/src/sequince/result.py b/src/sequince/result.py index 5f787ccb..5836de37 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -1,8 +1,7 @@ from __future__ import annotations -import copy import math -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Callable @@ -54,37 +53,6 @@ def append_state(self, state: ParticleState) -> SMCResult: context=self.context, ) - def snapshot(self) -> dict[str, Any]: - """Return a detached snapshot suitable for checkpoint payloads.""" - return { - "states": [state.snapshot() for state in self.states], - "context": self.context.snapshot(), - } - - @classmethod - def from_snapshot(cls, payload: Mapping[str, Any]) -> "SMCResult": - """Reconstruct a result from :meth:`snapshot` output.""" - if not isinstance(payload, Mapping): - raise TypeError("SMCResult snapshot must be a mapping.") - missing = {key for key in ("states", "context") if key not in payload} - if missing: - raise ValueError( - "SMCResult snapshot is missing required keys: " - + ", ".join(sorted(missing)) - ) - states_payload = payload["states"] - if not isinstance(states_payload, Sequence): - raise TypeError("SMCResult snapshot 'states' must be a sequence.") - return cls( - states=[ - ParticleState.from_snapshot(state_payload) - for state_payload in states_payload - ], - context=StepContext.from_snapshot( - copy.deepcopy(payload["context"]) - ), - ) - def get_nested_results( self, key: str = NESTED_SMC_OUTPUT_KEY, diff --git a/src/sequince/smc.py b/src/sequince/smc.py index 74aef296..b876e839 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -10,8 +10,8 @@ CheckpointWriter, SMCCheckpoint, StepState, + capture_step_states, restore_steps, - snapshot_steps, ) from .result import SMCResult from .state import ParticleState @@ -85,7 +85,9 @@ def __init__( def state_dict(self) -> dict[str, Any]: """Return mutable state for the configured step pipeline.""" return { - "steps": [item.to_dict() for item in snapshot_steps(self._steps)] + "steps": [ + item.to_dict() for item in capture_step_states(self._steps) + ] } def load_state_dict(self, state: Mapping[str, Any]) -> None: @@ -337,14 +339,11 @@ def resume_from( ) restore_steps(self._steps, checkpoint_obj.step_states) restored_rng = checkpoint_obj.restore_rng() - state = ParticleState.from_snapshot(checkpoint_obj.state.snapshot()) + state = checkpoint_obj.state.copy() history = checkpoint_obj.history or (checkpoint_obj.state,) return SMCRuntime( state=state, - produced_states=[ - ParticleState.from_snapshot(item.snapshot()) - for item in history - ], + 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 {}), diff --git a/src/sequince/state.py b/src/sequince/state.py index ecac4d5d..7419baca 100644 --- a/src/sequince/state.py +++ b/src/sequince/state.py @@ -62,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]: @@ -134,21 +134,21 @@ def resample(self, indices: Array) -> "ParticleState": metadata=dict(self.metadata, ancestors=indices), ) - def snapshot(self) -> dict[str, Any]: - """Return a detached snapshot suitable for checkpoint payloads.""" - xp = namespace(self.particles, self.log_weights) + def to_dict(self) -> dict[str, Any]: + """Return a detached dictionary representation of this state.""" + copied = self.copy() return { - "particles": array_copy(self.particles, xp=xp), - "log_weights": array_copy(self.log_weights, xp=xp), - "log_evidence": float(self.log_evidence), - "metadata": copy.deepcopy(self.metadata), + "particles": copied.particles, + "log_weights": copied.log_weights, + "log_evidence": copied.log_evidence, + "metadata": copied.metadata, } @classmethod - def from_snapshot(cls, payload: Mapping[str, Any]) -> "ParticleState": - """Reconstruct a state from :meth:`snapshot` output.""" + 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 snapshot must be a mapping.") + raise TypeError("ParticleState payload must be a mapping.") missing = { key for key in ("particles", "log_weights", "log_evidence", "metadata") @@ -156,7 +156,7 @@ def from_snapshot(cls, payload: Mapping[str, Any]) -> "ParticleState": } if missing: raise ValueError( - "ParticleState snapshot is missing required keys: " + "ParticleState payload is missing required keys: " + ", ".join(sorted(missing)) ) return cls( diff --git a/src/sequince/steps/context.py b/src/sequince/steps/context.py index f7bb9dc7..b72d8ef0 100644 --- a/src/sequince/steps/context.py +++ b/src/sequince/steps/context.py @@ -1,7 +1,5 @@ from __future__ import annotations -import copy -from collections.abc import Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -64,45 +62,3 @@ def clear_stop_request(self) -> None: def record_output(self, key: str, value: Any) -> None: """Store an output value that can be accessed by callbacks.""" self.outputs.setdefault(key, []).append(value) - - def snapshot(self) -> dict[str, Any]: - """Return a detached snapshot suitable for checkpoint payloads.""" - return { - "t": int(self.t), - "next_observation_index": int(self.next_observation_index), - "observation": copy.deepcopy(self.observation), - "shared": copy.deepcopy(dict(self.shared)), - "metadata": copy.deepcopy(self.metadata), - "outputs": copy.deepcopy(self.outputs), - } - - @classmethod - def from_snapshot(cls, payload: Mapping[str, Any]) -> "StepContext": - """Reconstruct a context from :meth:`snapshot` output.""" - if not isinstance(payload, Mapping): - raise TypeError("StepContext snapshot must be a mapping.") - missing = { - key - for key in ( - "t", - "next_observation_index", - "observation", - "shared", - "metadata", - "outputs", - ) - if key not in payload - } - if missing: - raise ValueError( - "StepContext snapshot is missing required keys: " - + ", ".join(sorted(missing)) - ) - return cls( - t=int(payload["t"]), - next_observation_index=int(payload["next_observation_index"]), - observation=copy.deepcopy(payload["observation"]), - shared=copy.deepcopy(payload["shared"]), - metadata=copy.deepcopy(payload["metadata"]), - outputs=copy.deepcopy(payload["outputs"]), - ) From 1adccf4e58415b734767f79ced304db2252d327e Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 10:11:46 +0100 Subject: [PATCH 12/19] TST: add tests for checkpointing --- tests/test_checkpointing.py | 620 ++++++++++++++++++++++++++++++++++++ 1 file changed, 620 insertions(+) create mode 100644 tests/test_checkpointing.py diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py new file mode 100644 index 00000000..00d960b3 --- /dev/null +++ b/tests/test_checkpointing.py @@ -0,0 +1,620 @@ +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, + *, + compare_observation: bool = True, +) -> 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 + ) + if compare_observation: + 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 + 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_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, compare_observation=False + ) + 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_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, compare_observation=False + ) + 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 + + +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, + } + ) From 6d8636c81a05c6de189736efd9659a3482f4c0ff Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 10:28:46 +0100 Subject: [PATCH 13/19] TST: improve test coverage --- tests/test_posterior.py | 17 ++++++++++++++--- tests/test_state.py | 3 +++ tests/test_transforms.py | 12 ++++++++++++ 3 files changed, 29 insertions(+), 3 deletions(-) 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: From 79efab5083d99f911edf3e06fd972869d82c3155 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 10:32:40 +0100 Subject: [PATCH 14/19] BLD: specify orng alpha release --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 609714dc..725592fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "minipcn[array-api]>=0.2.0a4", "array-api-compat>=1.5", "array-api-extra", - "orng>=0.2", + "orng>=0.2.0a1", "frozendict", ] dynamic = [ From fa994c1a10e93cd00bd86fe53f881d24d1b868ca Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 11:17:25 +0100 Subject: [PATCH 15/19] ENH: simplify HDF5 saving --- src/sequince/checkpoint.py | 119 ++++++++----------------------------- 1 file changed, 26 insertions(+), 93 deletions(-) diff --git a/src/sequince/checkpoint.py b/src/sequince/checkpoint.py index 042802cb..14e25873 100644 --- a/src/sequince/checkpoint.py +++ b/src/sequince/checkpoint.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, replace from pathlib import Path from typing import Any -from urllib.parse import quote +from urllib.parse import quote, unquote import numpy as np from orng import RandomGenerator @@ -19,9 +19,7 @@ CHECKPOINT_VERSION = 1 _HDF_KIND = "_sequince_kind" -_HDF_KEY_KIND = "_sequince_key_kind" -_HDF_KEY = "_sequince_key" -_HDF_SCALAR_TYPE = "_sequince_scalar_type" +_HDF_PYTHON_TYPE = "_sequince_python_type" def _attr_text(value: Any) -> str: @@ -46,43 +44,23 @@ def _write_pickled_hdf_value( def _write_hdf_value(parent: Any, name: str, value: Any, protocol: int) -> Any: - """Recursively write Python containers using native HDF5 structures.""" + """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 index, (key, item) in enumerate(value.items()): - child_name = ( - quote(key, safe="") or "%" - if isinstance(key, str) - else f"@{index:08d}" - ) - child = _write_hdf_value(node, child_name, item, protocol) - if isinstance(key, str): - try: - child.attrs[_HDF_KEY] = key - child.attrs[_HDF_KEY_KIND] = "string" - continue - except (TypeError, ValueError): - if _HDF_KEY in child.attrs: - del child.attrs[_HDF_KEY] - child.attrs[_HDF_KEY] = np.void( - pickle.dumps(key, protocol=protocol) - ) - child.attrs[_HDF_KEY_KIND] = "pickle" + 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, set, frozenset)): + if isinstance(value, (list, tuple)): node = parent.create_group(name) - container_kind = { - list: "list", - tuple: "tuple", - set: "set", - frozenset: "frozenset", - } - node.attrs[_HDF_KIND] = container_kind[type(value)] + 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 @@ -108,41 +86,23 @@ def _write_hdf_value(parent: Any, name: str, value: Any, protocol: int) -> Any: node.attrs[_HDF_KIND] = "bytes" return node - if isinstance(value, np.ndarray): + 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] = "array" - return node - - if isinstance(value, np.generic): - 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_scalar" + node.attrs[_HDF_KIND] = "numpy" return node if isinstance(value, (bool, int, float, complex)): try: node = parent.create_dataset(name, data=value) - except OverflowError: - if isinstance(value, int): - node = parent.create_dataset( - name, - data=str(value), - dtype="S" + str(len(str(value))), - ) - node.attrs[_HDF_KIND] = "bigint" - return node - return _write_pickled_hdf_value(parent, name, value, protocol) - except (TypeError, ValueError): + except (OverflowError, TypeError, ValueError): return _write_pickled_hdf_value(parent, name, value, protocol) node.attrs[_HDF_KIND] = "scalar" - node.attrs[_HDF_SCALAR_TYPE] = type(value).__name__ + node.attrs[_HDF_PYTHON_TYPE] = type(value).__name__ return node return _write_pickled_hdf_value(parent, name, value, protocol) @@ -155,32 +115,14 @@ def _read_hdf_value(node: Any) -> Any: kind = _attr_text(node.attrs[_HDF_KIND]) if kind == "mapping": - mapping_value: dict[Any, Any] = {} - for child_name in sorted(node.keys()): - child = node[child_name] - key_kind = _attr_text(child.attrs[_HDF_KEY_KIND]) - if key_kind == "string": - key = child.attrs[_HDF_KEY] - if isinstance(key, bytes): - key = key.decode("utf-8") - elif key_kind == "pickle": - key = pickle.loads(bytes(child.attrs[_HDF_KEY])) - else: - raise ValueError( - f"Unsupported HDF5 mapping key kind {key_kind!r}." - ) - mapping_value[key] = _read_hdf_value(child) - return mapping_value + return { + "" if name == "%" else unquote(name): _read_hdf_value(node[name]) + for name in node + } - if kind in ("list", "tuple", "set", "frozenset"): + if kind in ("list", "tuple"): items = [_read_hdf_value(node[name]) for name in sorted(node.keys())] - container_constructors = { - "list": list, - "tuple": tuple, - "set": set, - "frozenset": frozenset, - } - return container_constructors[kind](items) + return items if kind == "list" else tuple(items) if kind == "none": return None @@ -193,17 +135,13 @@ def _read_hdf_value(node: Any) -> Any: ) if kind == "bytes": return node[...].tobytes() - if kind == "array": - return node[...] - if kind == "numpy_scalar": + if kind == "numpy": return node[()] - if kind == "bigint": - return int(node[()].decode("ascii")) if kind == "scalar": scalar_value = node[()] if isinstance(scalar_value, np.generic): scalar_value = scalar_value.item() - scalar_type = _attr_text(node.attrs[_HDF_SCALAR_TYPE]) + scalar_type = _attr_text(node.attrs[_HDF_PYTHON_TYPE]) scalar_constructors = { "bool": bool, "int": int, @@ -548,7 +486,6 @@ def save_hdf5( if self.history is None: if "history" in group: del group["history"] - runtime_payload["history_length"] = 0 else: history_group = group.require_group("history") if replace_history: @@ -566,8 +503,6 @@ def save_hdf5( self.history[index].to_dict(), protocol, ) - history_group.attrs["length"] = len(self.history) - runtime_payload["history_length"] = len(self.history) temporary_name = f"{dsetname}.tmp" _write_hdf_value(group, temporary_name, runtime_payload, protocol) @@ -587,14 +522,12 @@ def load_hdf5( checkpoint_payload = _read_hdf_value(group[dsetname]) if not isinstance(checkpoint_payload, dict): raise TypeError("HDF5 checkpoint runtime must be a mapping.") - history_length = int(checkpoint_payload.pop("history_length", 0)) - if history_length: - if "history" not in group: - raise ValueError("Checkpoint history is missing from HDF5.") + checkpoint_payload["history"] = None + if "history" in group: history_group = group["history"] checkpoint_payload["history"] = [ - _read_hdf_value(history_group[f"{index:08d}"]) - for index in range(history_length) + _read_hdf_value(history_group[name]) + for name in sorted(history_group) ] return cls.from_dict(checkpoint_payload) From 41d064b4de66f511f6119bff7ea2e6504f9e5a49 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 14 Jul 2026 11:17:50 +0100 Subject: [PATCH 16/19] DOC: add checkpointing docs --- docs/checkpointing.rst | 221 +++++++++++++++++++++++++++++++++++++++++ docs/index.rst | 1 + 2 files changed, 222 insertions(+) create mode 100644 docs/checkpointing.rst diff --git a/docs/checkpointing.rst b/docs/checkpointing.rst new file mode 100644 index 00000000..d3499bf9 --- /dev/null +++ b/docs/checkpointing.rst @@ -0,0 +1,221 @@ +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. + +``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; +* 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 observation sequence itself. + +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 From 79d0f0b0a2e314ab795f348451ea37edb8653414 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Fri, 17 Jul 2026 10:37:05 +0100 Subject: [PATCH 17/19] DOC: update api --- docs/api/core.rst | 8 ++++++++ 1 file changed, 8 insertions(+) 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 ------------------- From 43b8f9973adc05476576952a57daee66be1ecbba Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Wed, 29 Jul 2026 17:55:44 +0100 Subject: [PATCH 18/19] FIX: make checkpoint resume state robust --- docs/checkpointing.rst | 6 ++- pyproject.toml | 5 ++- src/sequince/checkpoint.py | 15 +++++-- src/sequince/smc.py | 5 ++- tests/test_checkpointing.py | 81 ++++++++++++++++++++++++++++++++----- 5 files changed, 95 insertions(+), 17 deletions(-) diff --git a/docs/checkpointing.rst b/docs/checkpointing.rst index d3499bf9..8ead034a 100644 --- a/docs/checkpointing.rst +++ b/docs/checkpointing.rst @@ -43,7 +43,8 @@ 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. + 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 @@ -63,6 +64,7 @@ An :class:`~sequince.checkpoint.SMCCheckpoint` contains: 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``. @@ -71,7 +73,7 @@ A checkpoint deliberately does not contain: * the configured steps, callbacks, or their callable dependencies; * immutable ``shared_context`` values; -* the observation sequence itself. +* 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 diff --git a/pyproject.toml b/pyproject.toml index 725592fd..7f1b224d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,6 @@ dependencies = [ "scipy>=1.8", "matplotlib>=3.5", "numpy>=1.22", - "h5py>=3", "minipcn[array-api]>=0.2.0a4", "array-api-compat>=1.5", "array-api-extra", @@ -44,6 +43,9 @@ emcee = [ corner = [ "corner", ] +hdf5 = [ + "h5py>=3", +] [dependency-groups] test = [ @@ -51,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 index 14e25873..0b87cebd 100644 --- a/src/sequince/checkpoint.py +++ b/src/sequince/checkpoint.py @@ -297,6 +297,7 @@ class SMCCheckpoint: step_states: tuple[StepState, ...] rng_state: dict[str, Any] | None = None history: tuple[ParticleState, ...] | None = None + observation: Any = None version: int = CHECKPOINT_VERSION @classmethod @@ -330,6 +331,7 @@ def from_runtime( 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]: @@ -348,6 +350,7 @@ def to_dict(self) -> dict[str, Any]: if self.history is None else [item.to_dict() for item in self.history] ), + "observation": copy.deepcopy(self.observation), } @classmethod @@ -424,6 +427,7 @@ def from_dict(cls, payload: Mapping[str, Any]) -> "SMCCheckpoint": for item in history_payload ) ), + observation=copy.deepcopy(payload.get("observation")), version=version, ) @@ -441,6 +445,7 @@ def result( 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), @@ -558,21 +563,25 @@ def __init__( self.file_path = file_path self.path = path self.dsetname = dsetname - self._has_written = False + 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=not self._has_written, + replace_history=replace_history, ) - self._has_written = True + self._last_iteration = checkpoint.next_iteration class PickleCheckpointWriter(CheckpointWriter): diff --git a/src/sequince/smc.py b/src/sequince/smc.py index b876e839..2a1c141c 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -341,13 +341,16 @@ def resume_from( 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=copy.deepcopy(checkpoint_obj.metadata), + 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/tests/test_checkpointing.py b/tests/test_checkpointing.py index 00d960b3..5d672c17 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -121,8 +121,6 @@ def _assert_states_equal(left: ParticleState, right: ParticleState) -> None: def _assert_results_equal( left: SMCResult, right: SMCResult, - *, - compare_observation: bool = True, ) -> None: assert len(left.states) == len(right.states) for left_state, right_state in zip(left.states, right.states): @@ -132,8 +130,7 @@ def _assert_results_equal( left.context.next_observation_index == right.context.next_observation_index ) - if compare_observation: - assert left.context.observation == right.context.observation + 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 @@ -222,6 +219,7 @@ def test_checkpoint_round_trip_preserves_runtime_objects( _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( @@ -230,6 +228,25 @@ def test_checkpoint_round_trip_preserves_runtime_objects( ) +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: @@ -322,9 +339,7 @@ def test_hdf5_checkpoint_writer_writes_latest_checkpoint( checkpoint = SMCCheckpoint.load_hdf5(h5_file, path="checkpoint") _assert_states_equal(checkpoint.state, result.final_state) - _assert_results_equal( - checkpoint.result(), result, compare_observation=False - ) + _assert_results_equal(checkpoint.result(), result) assert checkpoint.next_iteration == 3 assert checkpoint.next_observation_index == 3 assert checkpoint.rng_state is not None @@ -377,6 +392,26 @@ def test_hdf5_checkpoint_can_omit_result_history( 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, @@ -393,9 +428,7 @@ def test_pickle_checkpoint_writer_writes_latest_checkpoint( checkpoint = SMCCheckpoint.load_pickle(file_path) _assert_states_equal(checkpoint.state, result.final_state) - _assert_results_equal( - checkpoint.result(), result, compare_observation=False - ) + _assert_results_equal(checkpoint.result(), result) assert checkpoint.next_iteration == 3 assert checkpoint.next_observation_index == 3 assert checkpoint.rng_state is not None @@ -465,6 +498,34 @@ def test_resume_without_saved_rng_works_for_deterministic_steps( 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( From fb600fd2b475bbc6342edc1df2998a988d80ef4f Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 3 Aug 2026 10:24:45 +0100 Subject: [PATCH 19/19] STY: run ruff --- src/sequince/steps/control.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/sequince/steps/control.py b/src/sequince/steps/control.py index 935fd9fe..f2a4c891 100644 --- a/src/sequince/steps/control.py +++ b/src/sequince/steps/control.py @@ -160,6 +160,7 @@ def load_state_dict(self, state: Mapping[str, Any]) -> None: 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."""