aspire-inference 0.1.0a10__py3-none-any.whl → 0.1.0a11__py3-none-any.whl

aspire/aspire.py

@@ -1,5 +1,8 @@
+import copy
 import logging
 import multiprocessing as mp
+import pickle
+from contextlib import contextmanager
 from inspect import signature
 from typing import Any, Callable
 
@@ -8,13 +11,19 @@ import h5py
 from .flows import get_flow_wrapper
 from .flows.base import Flow
 from .history import History
+from .samplers.base import Sampler
 from .samples import Samples
 from .transforms import (
     CompositeTransform,
     FlowPreconditioningTransform,
     FlowTransform,
 )
-from .utils import recursively_save_to_h5_file
+from .utils import (
+    AspireFile,
+    load_from_h5_file,
+    recursively_save_to_h5_file,
+    resolve_xp,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -102,6 +111,7 @@ class Aspire:
         self.dtype = dtype
 
         self._flow = flow
+        self._sampler = None
 
     @property
     def flow(self):
@@ -114,7 +124,7 @@ class Aspire:
         self._flow = flow
 
     @property
-    def sampler(self):
+    def sampler(self) -> Sampler | None:
         """The sampler object."""
         return self._sampler
 
@@ -192,7 +202,29 @@ class Aspire:
             **self.flow_kwargs,
         )
 
-    def fit(self, samples: Samples, **kwargs) -> History:
+    def fit(
+        self,
+        samples: Samples,
+        checkpoint_path: str | None = None,
+        checkpoint_save_config: bool = True,
+        overwrite: bool = False,
+        **kwargs,
+    ) -> History:
+        """Fit the normalizing flow to the provided samples.
+
+        Parameters
+        ----------
+        samples : Samples
+            The samples to fit the flow to.
+        checkpoint_path : str | None
+            Path to save the checkpoint. If None, no checkpoint is saved.
+        checkpoint_save_config : bool
+            Whether to save the Aspire configuration to the checkpoint.
+        overwrite : bool
+            Whether to overwrite an existing flow in the checkpoint file.
+        kwargs : dict
+            Keyword arguments to pass to the flow's fit method.
+        """
         if self.xp is None:
             self.xp = samples.xp
 
@@ -202,6 +234,28 @@ class Aspire:
         self.training_samples = samples
         logger.info(f"Training with {len(samples.x)} samples")
         history = self.flow.fit(samples.x, **kwargs)
+        defaults = getattr(self, "_checkpoint_defaults", None)
+        if checkpoint_path is None and defaults:
+            checkpoint_path = defaults["path"]
+            checkpoint_save_config = defaults["save_config"]
+        saved_config = (
+            defaults.get("saved_config", False) if defaults else False
+        )
+        if checkpoint_path is not None:
+            with AspireFile(checkpoint_path, "a") as h5_file:
+                if checkpoint_save_config and not saved_config:
+                    if "aspire_config" in h5_file:
+                        del h5_file["aspire_config"]
+                    self.save_config(h5_file, include_sampler_config=False)
+                    if defaults is not None:
+                        defaults["saved_config"] = True
+                # Save flow only if missing or overwrite=True
+                if "flow" in h5_file:
+                    if overwrite:
+                        del h5_file["flow"]
+                        self.save_flow(h5_file)
+                else:
+                    self.save_flow(h5_file)
         return history
 
     def get_sampler_class(self, sampler_type: str) -> Callable:
@@ -241,6 +295,13 @@ class Aspire:
         ----------
         sampler_type : str
             The type of sampler to use. Options are 'importance', 'emcee', or 'smc'.
+        preconditioning: str
+            Type of preconditioning to apply in the sampler. Options are
+            'default', 'flow', or 'none'.
+        preconditioning_kwargs: dict
+            Keyword arguments to pass to the preconditioning transform.
+        kwargs : dict
+            Keyword arguments to pass to the sampler.
         """
         SamplerClass = self.get_sampler_class(sampler_type)
 
@@ -304,6 +365,9 @@ class Aspire:
         return_history: bool = False,
         preconditioning: str | None = None,
         preconditioning_kwargs: dict | None = None,
+        checkpoint_path: str | None = None,
+        checkpoint_every: int = 1,
+        checkpoint_save_config: bool = True,
         **kwargs,
     ) -> Samples:
         """Draw samples from the posterior distribution.
@@ -342,6 +406,14 @@ class Aspire:
             will default to 'none' and the other samplers to 'default'
         preconditioning_kwargs: dict
             Keyword arguments to pass to the preconditioning transform.
+        checkpoint_path : str | None
+            Path to save the checkpoint. If None, no checkpoint is saved unless
+            within an :py:meth:`auto_checkpoint` context or a custom callback
+            is provided.
+        checkpoint_every : int
+            Frequency (in number of sampler iterations) to save the checkpoint.
+        checkpoint_save_config : bool
+            Whether to save the Aspire configuration to the checkpoint.
         kwargs : dict
             Keyword arguments to pass to the sampler. These are passed
             automatically to the init method of the sampler or to the sample
@@ -352,6 +424,22 @@ class Aspire:
         samples : Samples
             Samples object contain samples and their corresponding weights.
         """
+        if (
+            sampler == "importance"
+            and hasattr(self, "_resume_sampler_type")
+            and self._resume_sampler_type
+        ):
+            sampler = self._resume_sampler_type
+
+        if "resume_from" not in kwargs and hasattr(
+            self, "_resume_from_default"
+        ):
+            kwargs["resume_from"] = self._resume_from_default
+            if hasattr(self, "_resume_overrides"):
+                kwargs.update(self._resume_overrides)
+            if hasattr(self, "_resume_n_samples") and n_samples == 1000:
+                n_samples = self._resume_n_samples
+
         SamplerClass = self.get_sampler_class(sampler)
         # Determine sampler initialization parameters
         # and remove them from kwargs
@@ -373,7 +461,73 @@ class Aspire:
             preconditioning_kwargs=preconditioning_kwargs,
             **sampler_kwargs,
         )
+        self._last_sampler_type = sampler
+        # Auto-checkpoint convenience: set defaults for checkpointing to a single file
+        defaults = getattr(self, "_checkpoint_defaults", None)
+        if checkpoint_path is None and defaults:
+            checkpoint_path = defaults["path"]
+            checkpoint_every = defaults["every"]
+            checkpoint_save_config = defaults["save_config"]
+        saved_flow = defaults.get("saved_flow", False) if defaults else False
+        saved_config = (
+            defaults.get("saved_config", False) if defaults else False
+        )
+        if checkpoint_path is not None:
+            kwargs.setdefault("checkpoint_file_path", checkpoint_path)
+            kwargs.setdefault("checkpoint_every", checkpoint_every)
+            with AspireFile(checkpoint_path, "a") as h5_file:
+                if checkpoint_save_config:
+                    if "aspire_config" in h5_file:
+                        del h5_file["aspire_config"]
+                    self.save_config(
+                        h5_file,
+                        include_sampler_config=True,
+                        include_sample_calls=False,
+                    )
+                    saved_config = True
+                    if defaults is not None:
+                        defaults["saved_config"] = True
+                if (
+                    self.flow is not None
+                    and not saved_flow
+                    and "flow" not in h5_file
+                ):
+                    self.save_flow(h5_file)
+                    saved_flow = True
+                    if defaults is not None:
+                        defaults["saved_flow"] = True
+
         samples = self._sampler.sample(n_samples, **kwargs)
+        self._last_sample_posterior_kwargs = {
+            "n_samples": n_samples,
+            "sampler": sampler,
+            "xp": xp,
+            "return_history": return_history,
+            "preconditioning": preconditioning,
+            "preconditioning_kwargs": preconditioning_kwargs,
+            "sampler_init_kwargs": sampler_kwargs,
+            "sample_kwargs": copy.deepcopy(kwargs),
+        }
+        if checkpoint_path is not None:
+            with AspireFile(checkpoint_path, "a") as h5_file:
+                if checkpoint_save_config and not saved_config:
+                    if "aspire_config" in h5_file:
+                        del h5_file["aspire_config"]
+                    self.save_config(
+                        h5_file,
+                        include_sampler_config=True,
+                        include_sample_calls=False,
+                    )
+                    if defaults is not None:
+                        defaults["saved_config"] = True
+                if (
+                    self.flow is not None
+                    and not saved_flow
+                    and "flow" not in h5_file
+                ):
+                    self.save_flow(h5_file)
+                    if defaults is not None:
+                        defaults["saved_flow"] = True
         if xp is not None:
             samples = samples.to_namespace(xp)
         samples.parameters = self.parameters
@@ -388,6 +542,122 @@ class Aspire:
         else:
             return samples
 
+    @classmethod
+    def resume_from_file(
+        cls,
+        file_path: str,
+        *,
+        log_likelihood: Callable,
+        log_prior: Callable,
+        sampler: str | None = None,
+        checkpoint_path: str = "checkpoint",
+        checkpoint_dset: str = "state",
+        flow_path: str = "flow",
+        config_path: str = "aspire_config",
+        resume_kwargs: dict | None = None,
+    ):
+        """
+        Recreate an Aspire object from a single file and prepare to resume sampling.
+
+        Parameters
+        ----------
+        file_path : str
+            Path to the HDF5 file containing config, flow, and checkpoint.
+        log_likelihood : Callable
+            Log-likelihood function (required, not pickled).
+        log_prior : Callable
+            Log-prior function (required, not pickled).
+        sampler : str
+            Sampler type to use (e.g., 'smc', 'minipcn_smc', 'emcee_smc'). If None,
+            will attempt to infer from saved config or checkpoint metadata.
+        checkpoint_path : str
+            HDF5 group path where the checkpoint is stored.
+        checkpoint_dset : str
+            Dataset name within the checkpoint group.
+        flow_path : str
+            HDF5 path to the saved flow.
+        config_path : str
+            HDF5 path to the saved Aspire config.
+        resume_kwargs : dict | None
+            Optional overrides to apply when resuming (e.g., checkpoint_every).
+        """
+        (
+            aspire,
+            checkpoint_bytes,
+            checkpoint_state,
+            sampler_config,
+            saved_sampler_type,
+            n_samples,
+        ) = cls._build_aspire_from_file(
+            file_path=file_path,
+            log_likelihood=log_likelihood,
+            log_prior=log_prior,
+            checkpoint_path=checkpoint_path,
+            checkpoint_dset=checkpoint_dset,
+            flow_path=flow_path,
+            config_path=config_path,
+        )
+
+        sampler_config = sampler_config or {}
+        sampler_config.pop("sampler_class", None)
+
+        if checkpoint_bytes is not None:
+            aspire._resume_from_default = checkpoint_bytes
+            aspire._resume_sampler_type = (
+                sampler
+                or saved_sampler_type
+                or (
+                    checkpoint_state.get("sampler")
+                    if checkpoint_state
+                    else None
+                )
+            )
+            aspire._resume_n_samples = n_samples
+            aspire._resume_overrides = resume_kwargs or {}
+            aspire._resume_sampler_config = sampler_config
+        aspire._checkpoint_defaults = {
+            "path": file_path,
+            "every": 1,
+            "save_config": False,
+            "save_flow": False,
+            "saved_config": False,
+            "saved_flow": False,
+        }
+        return aspire
+
+    @contextmanager
+    def auto_checkpoint(
+        self,
+        path: str,
+        every: int = 1,
+        save_config: bool = True,
+        save_flow: bool = True,
+    ):
+        """
+        Context manager to auto-save checkpoints, config, and flow to a file.
+
+        Within the context, sample_posterior will default to writing checkpoints
+        to the given path with the specified frequency, and will append config/flow
+        after sampling.
+        """
+        prev = getattr(self, "_checkpoint_defaults", None)
+        self._checkpoint_defaults = {
+            "path": path,
+            "every": every,
+            "save_config": save_config,
+            "save_flow": save_flow,
+            "saved_config": False,
+            "saved_flow": False,
+        }
+        try:
+            yield self
+        finally:
+            if prev is None:
+                if hasattr(self, "_checkpoint_defaults"):
+                    delattr(self, "_checkpoint_defaults")
+            else:
+                self._checkpoint_defaults = prev
+
     def enable_pool(self, pool: mp.Pool, **kwargs):
         """Context manager to temporarily replace the log_likelihood method
         with a version that uses a multiprocessing pool to parallelize
@@ -432,12 +702,16 @@ class Aspire:
             "flow_kwargs": self.flow_kwargs,
             "eps": self.eps,
         }
+        if hasattr(self, "_last_sampler_type"):
+            config["sampler_type"] = self._last_sampler_type
         if include_sampler_config:
+            if self.sampler is None:
+                raise ValueError("Sampler has not been initialized.")
             config["sampler_config"] = self.sampler.config_dict(**kwargs)
         return config
 
     def save_config(
-        self, h5_file: h5py.File, path="aspire_config", **kwargs
+        self, h5_file: h5py.File | AspireFile, path="aspire_config", **kwargs
     ) -> None:
         """Save the configuration to an HDF5 file.
 
@@ -484,6 +758,7 @@ class Aspire:
         FlowClass, xp = get_flow_wrapper(
             backend=self.flow_backend, flow_matching=self.flow_matching
         )
+        logger.debug(f"Loading flow of type {FlowClass} from {path}")
         self._flow = FlowClass.load(h5_file, path=path)
 
     def save_config_to_json(self, filename: str) -> None:
@@ -504,3 +779,80 @@ class Aspire:
         x, log_q = self.flow.sample_and_log_prob(n_samples)
         samples = Samples(x=x, log_q=log_q, xp=xp, parameters=self.parameters)
         return samples
+
+    # --- Resume helpers ---
+    @staticmethod
+    def _build_aspire_from_file(
+        file_path: str,
+        log_likelihood: Callable,
+        log_prior: Callable,
+        checkpoint_path: str,
+        checkpoint_dset: str,
+        flow_path: str,
+        config_path: str,
+    ):
+        """Construct an Aspire instance, load flow, and gather checkpoint metadata from file."""
+        with AspireFile(file_path, "r") as h5_file:
+            if config_path not in h5_file:
+                raise ValueError(
+                    f"Config path '{config_path}' not found in {file_path}"
+                )
+            config_dict = load_from_h5_file(h5_file, config_path)
+            try:
+                checkpoint_bytes = h5_file[checkpoint_path][checkpoint_dset][
+                    ...
+                ].tobytes()
+            except Exception:
+                logger.warning(
+                    "Checkpoint not found at %s/%s in %s; will resume without a checkpoint.",
+                    checkpoint_path,
+                    checkpoint_dset,
+                    file_path,
+                )
+                checkpoint_bytes = None
+
+        sampler_config = config_dict.pop("sampler_config", None)
+        saved_sampler_type = config_dict.pop("sampler_type", None)
+        if isinstance(config_dict.get("xp"), str):
+            config_dict["xp"] = resolve_xp(config_dict["xp"])
+        config_dict["log_likelihood"] = log_likelihood
+        config_dict["log_prior"] = log_prior
+
+        aspire = Aspire(**config_dict)
+
+        with AspireFile(file_path, "r") as h5_file:
+            if flow_path in h5_file:
+                logger.info(f"Loading flow from {flow_path} in {file_path}")
+                aspire.load_flow(h5_file, path=flow_path)
+            else:
+                raise ValueError(
+                    f"Flow path '{flow_path}' not found in {file_path}"
+                )
+
+        n_samples = None
+        checkpoint_state = None
+        if checkpoint_bytes is not None:
+            try:
+                checkpoint_state = pickle.loads(checkpoint_bytes)
+                samples_saved = (
+                    checkpoint_state.get("samples")
+                    if checkpoint_state
+                    else None
+                )
+                if samples_saved is not None:
+                    n_samples = len(samples_saved)
+                    if aspire.xp is None and hasattr(samples_saved, "xp"):
+                        aspire.xp = samples_saved.xp
+            except Exception:
+                logger.warning(
+                    "Failed to decode checkpoint; proceeding without resume state."
+                )
+
+        return (
+            aspire,
+            checkpoint_bytes,
+            checkpoint_state,
+            sampler_config,
+            saved_sampler_type,
+            n_samples,
+        )

aspire/flows/torch/flows.py

@@ -92,7 +92,7 @@ class BaseTorchFlow(Flow):
         config = load_from_h5_file(flow_grp, "config")
         config["dtype"] = decode_dtype(torch, config.get("dtype"))
         if "data_transform" in flow_grp:
-            from ..transforms import BaseTransform
+            from ...transforms import BaseTransform
 
             data_transform = BaseTransform.load(
                 flow_grp,

aspire/samplers/base.py

@@ -1,10 +1,12 @@
 import logging
+import pickle
+from pathlib import Path
 from typing import Any, Callable
 
 from ..flows.base import Flow
 from ..samples import Samples
 from ..transforms import IdentityTransform
-from ..utils import asarray, track_calls
+from ..utils import AspireFile, asarray, dump_state, track_calls
 
 logger = logging.getLogger(__name__)
 
@@ -49,6 +51,8 @@ class Sampler:
         self.parameters = parameters
         self.history = None
         self.n_likelihood_evaluations = 0
+        self._last_checkpoint_state: dict | None = None
+        self._last_checkpoint_bytes: bytes | None = None
         if preconditioning_transform is None:
             self.preconditioning_transform = IdentityTransform(xp=self.xp)
         else:
@@ -75,7 +79,7 @@ class Sampler:
         self.n_likelihood_evaluations += len(samples)
         return self._log_likelihood(samples)
 
-    def config_dict(self, include_sample_calls: bool = True) -> dict:
+    def config_dict(self, include_sample_calls: bool = False) -> dict:
         """
         Returns a dictionary with the configuration of the sampler.
 
@@ -83,9 +87,9 @@ class Sampler:
         ----------
         include_sample_calls : bool
             Whether to include the sample calls in the configuration.
-            Default is True.
+            Default is False.
         """
-        config = {}
+        config = {"sampler_class": self.__class__.__name__}
         if include_sample_calls:
             if hasattr(self, "sample") and hasattr(self.sample, "calls"):
                 config["sample_calls"] = self.sample.calls.to_dict(
@@ -96,3 +100,139 @@ class Sampler:
                     "Sampler does not have a sample method with calls attribute."
                 )
         return config
+
+    # --- Checkpointing helpers shared across samplers ---
+    def _checkpoint_extra_state(self) -> dict:
+        """Sampler-specific extras for checkpointing (override in subclasses)."""
+        return {}
+
+    def _restore_extra_state(self, state: dict) -> None:
+        """Restore sampler-specific extras (override in subclasses)."""
+        _ = state  # no-op for base
+
+    def build_checkpoint_state(
+        self,
+        samples: Samples,
+        iteration: int | None = None,
+        meta: dict | None = None,
+    ) -> dict:
+        """Prepare a serializable checkpoint payload for the sampler state."""
+        checkpoint_samples = samples
+        base_state = {
+            "sampler": self.__class__.__name__,
+            "iteration": iteration,
+            "samples": checkpoint_samples,
+            "config": self.config_dict(include_sample_calls=False),
+            "parameters": self.parameters,
+            "meta": meta or {},
+        }
+        base_state.update(self._checkpoint_extra_state())
+        return base_state
+
+    def serialize_checkpoint(
+        self, state: dict, protocol: int | None = None
+    ) -> bytes:
+        """Serialize a checkpoint state to bytes with pickle."""
+        protocol = (
+            pickle.HIGHEST_PROTOCOL if protocol is None else int(protocol)
+        )
+        return pickle.dumps(state, protocol=protocol)
+
+    def default_checkpoint_callback(self, state: dict) -> None:
+        """Store the latest checkpoint (state + pickled bytes) on the sampler."""
+        self._last_checkpoint_state = state
+        self._last_checkpoint_bytes = self.serialize_checkpoint(state)
+
+    def default_file_checkpoint_callback(
+        self, file_path: str | Path | None
+    ) -> Callable[[dict], None]:
+        """Return a simple default callback that overwrites an HDF5 file."""
+        if file_path is None:
+            return self.default_checkpoint_callback
+        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)."
+            )
+
+        def _callback(state: dict) -> None:
+            with AspireFile(file_path, "a") as h5_file:
+                self.save_checkpoint_to_hdf(
+                    state, h5_file, path="checkpoint", dsetname="state"
+                )
+            self.default_checkpoint_callback(state)
+
+        return _callback
+
+    def save_checkpoint_to_hdf(
+        self,
+        state: dict,
+        h5_file,
+        path: str = "sampler_checkpoints",
+        dsetname: str | None = None,
+        protocol: int | None = None,
+    ) -> None:
+        """Save a checkpoint state into an HDF5 file as a pickled blob."""
+        if dsetname is None:
+            iter_str = state.get("iteration", "unknown")
+            dsetname = f"iter_{iter_str}"
+        dump_state(
+            state,
+            h5_file,
+            path=path,
+            dsetname=dsetname,
+            protocol=protocol or pickle.HIGHEST_PROTOCOL,
+        )
+
+    def load_checkpoint_from_file(
+        self,
+        file_path: str | Path,
+        h5_path: str = "checkpoint",
+        dsetname: str = "state",
+    ) -> dict:
+        """Load a checkpoint dictionary from .pkl or .hdf5 file."""
+        file_path = Path(file_path)
+        lower_path = file_path.name.lower()
+        if lower_path.endswith((".h5", ".hdf5")):
+            with AspireFile(file_path, "r") as h5_file:
+                data = h5_file[h5_path][dsetname][...]
+                checkpoint_bytes = data.tobytes()
+        else:
+            with open(file_path, "rb") as f:
+                checkpoint_bytes = f.read()
+        return pickle.loads(checkpoint_bytes)
+
+    def restore_from_checkpoint(
+        self, source: str | bytes | dict
+    ) -> tuple[Samples, dict]:
+        """Restore sampler state from a checkpoint source."""
+        if isinstance(source, str):
+            state = self.load_checkpoint_from_file(source)
+        elif isinstance(source, bytes):
+            state = pickle.loads(source)
+        elif isinstance(source, dict):
+            state = source
+        else:
+            raise TypeError("Unsupported checkpoint source type.")
+
+        samples_saved = state.get("samples")
+        if samples_saved is None:
+            raise ValueError("Checkpoint missing samples.")
+
+        samples = Samples.from_samples(
+            samples_saved, xp=self.xp, dtype=self.dtype
+        )
+        # Allow subclasses to restore sampler-specific components
+        self._restore_extra_state(state)
+        return samples, state
+
+    @property
+    def last_checkpoint_state(self) -> dict | None:
+        """Return the most recent checkpoint state stored by the default callback."""
+        return self._last_checkpoint_state
+
+    @property
+    def last_checkpoint_bytes(self) -> bytes | None:
+        """Return the most recent pickled checkpoint produced by the default callback."""
+        return self._last_checkpoint_bytes

aspire/samplers/smc/base.py

@@ -1,3 +1,4 @@
+import copy
 import logging
 from typing import Any, Callable
 
@@ -158,11 +159,24 @@ class SMCSampler(MCMCSampler):
         target_efficiency: float = 0.5,
         target_efficiency_rate: float = 1.0,
         n_final_samples: int | None = None,
+        checkpoint_callback: Callable[[dict], None] | None = None,
+        checkpoint_every: int | None = None,
+        checkpoint_file_path: str | None = None,
+        resume_from: str | bytes | dict | None = None,
     ) -> SMCSamples:
-        samples = self.draw_initial_samples(n_samples)
-        samples = SMCSamples.from_samples(
-            samples, xp=self.xp, beta=0.0, dtype=self.dtype
-        )
+        resumed = resume_from is not None
+        if resumed:
+            samples, beta, iterations = self.restore_from_checkpoint(
+                resume_from
+            )
+        else:
+            samples = self.draw_initial_samples(n_samples)
+            samples = SMCSamples.from_samples(
+                samples, xp=self.xp, beta=0.0, dtype=self.dtype
+            )
+            beta = 0.0
+            iterations = 0
+            self.history = SMCHistory()
         self.fit_preconditioning_transform(samples.x)
 
         if self.xp.isnan(samples.log_q).any():
@@ -178,8 +192,6 @@ class SMCSampler(MCMCSampler):
         self.sampler_kwargs = self.sampler_kwargs or {}
         n_final_steps = self.sampler_kwargs.pop("n_final_steps", None)
 
-        self.history = SMCHistory()
-
         self.target_efficiency = target_efficiency
         self.target_efficiency_rate = target_efficiency_rate
 
@@ -190,7 +202,6 @@ class SMCSampler(MCMCSampler):
         else:
             beta_step = np.nan
         self.adaptive = adaptive
-        beta = 0.0
 
         if min_step is None:
             if max_n_steps is None:
@@ -202,55 +213,85 @@ class SMCSampler(MCMCSampler):
         else:
             self.adaptive_min_step = False
 
-        iterations = 0
-        while True:
-            iterations += 1
-
-            beta, min_step = self.determine_beta(
-                samples,
-                beta,
-                beta_step,
-                min_step,
+        iterations = iterations or 0
+        if checkpoint_callback is None and checkpoint_every is not None:
+            checkpoint_callback = self.default_file_checkpoint_callback(
+                checkpoint_file_path
             )
-            self.history.eff_target.append(
-                self.current_target_efficiency(beta)
+        if checkpoint_callback is not None and checkpoint_every is None:
+            checkpoint_every = 1
+
+        run_smc_loop = True
+        if resumed:
+            last_beta = self.history.beta[-1] if self.history.beta else beta
+            if last_beta >= 1.0:
+                run_smc_loop = False
+
+        def maybe_checkpoint(force: bool = False):
+            if checkpoint_callback is None:
+                return
+            should_checkpoint = force or (
+                checkpoint_every is not None
+                and checkpoint_every > 0
+                and iterations % checkpoint_every == 0
             )
+            if not should_checkpoint:
+                return
+            state = self.build_checkpoint_state(samples, iterations, beta)
+            checkpoint_callback(state)
+
+        if run_smc_loop:
+            while True:
+                iterations += 1
+
+                beta, min_step = self.determine_beta(
+                    samples,
+                    beta,
+                    beta_step,
+                    min_step,
+                )
+                self.history.eff_target.append(
+                    self.current_target_efficiency(beta)
+                )
 
-            logger.info(f"it {iterations} - beta: {beta}")
-            self.history.beta.append(beta)
-
-            ess = effective_sample_size(samples.log_weights(beta))
-            eff = ess / len(samples)
-            if eff < 0.1:
-                logger.warning(
-                    f"it {iterations} - Low sample efficiency: {eff:.2f}"
+                logger.info(f"it {iterations} - beta: {beta}")
+                self.history.beta.append(beta)
+
+                ess = effective_sample_size(samples.log_weights(beta))
+                eff = ess / len(samples)
+                if eff < 0.1:
+                    logger.warning(
+                        f"it {iterations} - Low sample efficiency: {eff:.2f}"
+                    )
+                self.history.ess.append(ess)
+                logger.info(
+                    f"it {iterations} - ESS: {ess:.1f} ({eff:.2f} efficiency)"
+                )
+                self.history.ess_target.append(
+                    effective_sample_size(samples.log_weights(1.0))
                 )
-            self.history.ess.append(ess)
-            logger.info(
-                f"it {iterations} - ESS: {ess:.1f} ({eff:.2f} efficiency)"
-            )
-            self.history.ess_target.append(
-                effective_sample_size(samples.log_weights(1.0))
-            )
 
-            log_evidence_ratio = samples.log_evidence_ratio(beta)
-            log_evidence_ratio_var = samples.log_evidence_ratio_variance(beta)
-            self.history.log_norm_ratio.append(log_evidence_ratio)
-            self.history.log_norm_ratio_var.append(log_evidence_ratio_var)
-            logger.info(
-                f"it {iterations} - Log evidence ratio: {log_evidence_ratio:.2f} +/- {np.sqrt(log_evidence_ratio_var):.2f}"
-            )
+                log_evidence_ratio = samples.log_evidence_ratio(beta)
+                log_evidence_ratio_var = samples.log_evidence_ratio_variance(
+                    beta
+                )
+                self.history.log_norm_ratio.append(log_evidence_ratio)
+                self.history.log_norm_ratio_var.append(log_evidence_ratio_var)
+                logger.info(
+                    f"it {iterations} - Log evidence ratio: {log_evidence_ratio:.2f} +/- {np.sqrt(log_evidence_ratio_var):.2f}"
+                )
 
-            samples = samples.resample(beta, rng=self.rng)
+                samples = samples.resample(beta, rng=self.rng)
 
-            samples = self.mutate(samples, beta)
-            if beta == 1.0 or (
-                max_n_steps is not None and iterations >= max_n_steps
-            ):
-                break
+                samples = self.mutate(samples, beta)
+                maybe_checkpoint()
+                if beta == 1.0 or (
+                    max_n_steps is not None and iterations >= max_n_steps
+                ):
+                    break
 
-        # If n_final_samples is not None, perform additional mutations steps
-        if n_final_samples is not None:
+        # If n_final_samples is specified and differs, perform additional mutation steps
+        if n_final_samples is not None and len(samples.x) != n_final_samples:
             logger.info(f"Generating {n_final_samples} final samples")
             final_samples = samples.resample(
                 1.0, n_samples=n_final_samples, rng=self.rng
@@ -263,6 +304,7 @@ class SMCSampler(MCMCSampler):
         samples.log_evidence_error = samples.xp.sqrt(
             samples.xp.sum(asarray(self.history.log_norm_ratio_var, self.xp))
         )
+        maybe_checkpoint(force=True)
 
         final_samples = samples.to_standard_samples()
         logger.info(
@@ -289,6 +331,49 @@ class SMCSampler(MCMCSampler):
         )
         return log_prob
 
+    def build_checkpoint_state(
+        self, samples: SMCSamples, iteration: int, beta: float
+    ) -> dict:
+        """Prepare a serializable checkpoint payload for the sampler state."""
+        return super().build_checkpoint_state(
+            samples,
+            iteration,
+            meta={"beta": beta},
+        )
+
+    def _checkpoint_extra_state(self) -> dict:
+        history_copy = copy.deepcopy(self.history)
+        rng_state = (
+            self.rng.bit_generator.state
+            if hasattr(self.rng, "bit_generator")
+            else None
+        )
+        return {
+            "history": history_copy,
+            "rng_state": rng_state,
+            "sampler_kwargs": getattr(self, "sampler_kwargs", None),
+        }
+
+    def restore_from_checkpoint(
+        self, source: str | bytes | dict
+    ) -> tuple[SMCSamples, float, int]:
+        samples, state = super().restore_from_checkpoint(source)
+        meta = state.get("meta", {}) if isinstance(state, dict) else {}
+        beta = None
+        if isinstance(meta, dict):
+            beta = meta.get("beta", None)
+        if beta is None:
+            beta = state.get("beta", 0.0)
+        iteration = state.get("iteration", 0)
+        self.history = state.get("history", SMCHistory())
+        rng_state = state.get("rng_state")
+        if rng_state is not None and hasattr(self.rng, "bit_generator"):
+            self.rng.bit_generator.state = rng_state
+        samples = SMCSamples.from_samples(
+            samples, xp=self.xp, beta=beta, dtype=self.dtype
+        )
+        return samples, beta, iteration
+
 
 class NumpySMCSampler(SMCSampler):
     def __init__(

aspire/samplers/smc/blackjax.py

@@ -84,6 +84,10 @@ class BlackJAXSMC(SMCSampler):
         n_final_samples: int | None = None,
         sampler_kwargs: dict | None = None,
         rng_key=None,
+        checkpoint_callback=None,
+        checkpoint_every: int | None = None,
+        checkpoint_file_path: str | None = None,
+        resume_from: str | bytes | dict | None = None,
     ):
         """Sample using BlackJAX SMC.
 
@@ -132,6 +136,10 @@ class BlackJAXSMC(SMCSampler):
             target_efficiency=target_efficiency,
             target_efficiency_rate=target_efficiency_rate,
             n_final_samples=n_final_samples,
+            checkpoint_callback=checkpoint_callback,
+            checkpoint_every=checkpoint_every,
+            checkpoint_file_path=checkpoint_file_path,
+            resume_from=resume_from,
         )
 
     def mutate(self, particles, beta, n_steps=None):

aspire/samplers/smc/emcee.py

@@ -21,6 +21,10 @@ class EmceeSMC(NumpySMCSampler):
         target_efficiency_rate: float = 1.0,
         sampler_kwargs: dict | None = None,
         n_final_samples: int | None = None,
+        checkpoint_callback=None,
+        checkpoint_every: int | None = None,
+        checkpoint_file_path: str | None = None,
+        resume_from: str | bytes | dict | None = None,
     ):
         self.sampler_kwargs = sampler_kwargs or {}
         self.sampler_kwargs.setdefault("nsteps", 5 * self.dims)
@@ -33,6 +37,10 @@ class EmceeSMC(NumpySMCSampler):
             target_efficiency=target_efficiency,
             target_efficiency_rate=target_efficiency_rate,
             n_final_samples=n_final_samples,
+            checkpoint_callback=checkpoint_callback,
+            checkpoint_every=checkpoint_every,
+            checkpoint_file_path=checkpoint_file_path,
+            resume_from=resume_from,
         )
 
     def mutate(self, particles, beta, n_steps=None):

aspire/samplers/smc/minipcn.py

@@ -8,10 +8,10 @@ from ...utils import (
     determine_backend_name,
     track_calls,
 )
-from .base import NumpySMCSampler
+from .base import SMCSampler
 
 
-class MiniPCNSMC(NumpySMCSampler):
+class MiniPCNSMC(SMCSampler):
     """MiniPCN SMC sampler."""
 
     rng = None
@@ -32,6 +32,10 @@ class MiniPCNSMC(NumpySMCSampler):
         n_final_samples: int | None = None,
         sampler_kwargs: dict | None = None,
         rng: np.random.Generator | None = None,
+        checkpoint_callback=None,
+        checkpoint_every: int | None = None,
+        checkpoint_file_path: str | None = None,
+        resume_from: str | bytes | dict | None = None,
     ):
         from orng import ArrayRNG
 
@@ -50,6 +54,10 @@ class MiniPCNSMC(NumpySMCSampler):
             n_final_samples=n_final_samples,
             min_step=min_step,
             max_n_steps=max_n_steps,
+            checkpoint_callback=checkpoint_callback,
+            checkpoint_every=checkpoint_every,
+            checkpoint_file_path=checkpoint_file_path,
+            resume_from=resume_from,
         )
 
     def mutate(self, particles, beta, n_steps=None):

aspire/samples.py

@@ -9,19 +9,18 @@ from typing import Any, Callable
 import numpy as np
 from array_api_compat import (
     array_namespace,
-    is_numpy_namespace,
-    to_device,
 )
-from array_api_compat import device as api_device
 from array_api_compat.common._typing import Array
 from matplotlib.figure import Figure
 
 from .utils import (
     asarray,
     convert_dtype,
+    infer_device,
     logsumexp,
     recursively_save_to_h5_file,
     resolve_dtype,
+    safe_to_device,
     to_numpy,
 )
 
@@ -67,8 +66,6 @@ class BaseSamples:
         if self.xp is None:
             self.xp = array_namespace(self.x)
         # Numpy arrays need to be on the CPU before being converted
-        if is_numpy_namespace(self.xp):
-            self.device = "cpu"
         if self.dtype is not None:
             self.dtype = resolve_dtype(self.dtype, self.xp)
         else:
@@ -76,7 +73,7 @@ class BaseSamples:
             self.dtype = None
         self.x = self.array_to_namespace(self.x, dtype=self.dtype)
         if self.device is None:
-            self.device = api_device(self.x)
+            self.device = infer_device(self.x, self.xp)
         if self.log_likelihood is not None:
             self.log_likelihood = self.array_to_namespace(
                 self.log_likelihood, dtype=self.dtype
@@ -140,8 +137,7 @@ class BaseSamples:
         else:
             kwargs["dtype"] = self.dtype
         x = asarray(x, self.xp, **kwargs)
-        if self.device:
-            x = to_device(x, self.device)
+        x = safe_to_device(x, self.device, self.xp)
         return x
 
     def to_dict(self, flat: bool = True):
@@ -174,7 +170,6 @@ class BaseSamples:
         ----------
         parameters : list[str] | None
             List of parameters to plot. If None, all parameters are plotted.
-        fig : matplotlib.figure.Figure | None
             Figure to plot on. If None, a new figure is created.
         **kwargs : dict
             Additional keyword arguments to pass to corner.corner(). Kwargs
@@ -300,6 +295,13 @@ class BaseSamples:
     def __setstate__(self, state):
         # Restore xp by checking the namespace of x
         state["xp"] = array_namespace(state["x"])
+        # device may be string; leave as-is or None
+        device = state.get("device")
+        if device is not None and "jax" in getattr(
+            state["xp"], "__name__", ""
+        ):
+            device = None
+        state["device"] = device
         self.__dict__.update(state)
 
 

aspire/utils.py

@@ -2,9 +2,11 @@ from __future__ import annotations
 
 import inspect
 import logging
+import pickle
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
+from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
 import array_api_compat.numpy as np
@@ -601,6 +603,7 @@ def decode_from_hdf5(value: Any) -> Any:
             return None
         if value == "__empty_dict__":
             return {}
+        return value
 
     if isinstance(value, np.ndarray):
         # Try to collapse 0-D arrays into scalars
@@ -629,6 +632,101 @@ def decode_from_hdf5(value: Any) -> Any:
     return value
 
 
+def dump_pickle_to_hdf(memfp, fp, path=None, dsetname="state"):
+    """Dump pickled data to an HDF5 file object."""
+    memfp.seek(0)
+    bdata = np.frombuffer(memfp.read(), dtype="S1")
+    target = fp.require_group(path) if path is not None else fp
+    if dsetname not in target:
+        target.create_dataset(
+            dsetname, shape=bdata.shape, maxshape=(None,), dtype=bdata.dtype
+        )
+    elif bdata.size != target[dsetname].shape[0]:
+        target[dsetname].resize((bdata.size,))
+    target[dsetname][:] = bdata
+
+
+def dump_state(
+    state,
+    fp,
+    path=None,
+    dsetname="state",
+    protocol=pickle.HIGHEST_PROTOCOL,
+):
+    """Pickle a state object and store it in an HDF5 dataset."""
+    memfp = BytesIO()
+    pickle.dump(state, memfp, protocol=protocol)
+    dump_pickle_to_hdf(memfp, fp, path=path, dsetname=dsetname)
+
+
+def resolve_xp(xp_name: str | None):
+    """
+    Resolve a backend name to the corresponding array_api_compat module.
+
+    Returns None if the name is None or cannot be resolved.
+    """
+    if xp_name is None:
+        return None
+    name = xp_name.lower()
+    if name.startswith("array_api_compat."):
+        name = name.removeprefix("array_api_compat.")
+    try:
+        if name in {"numpy", "numpy.ndarray"}:
+            import array_api_compat.numpy as np_xp
+
+            return np_xp
+        if name in {"jax", "jax.numpy"}:
+            import jax.numpy as jnp
+
+            return jnp
+        if name in {"torch"}:
+            import array_api_compat.torch as torch_xp
+
+            return torch_xp
+    except Exception:
+        logger.warning(
+            "Failed to resolve xp '%s', defaulting to None", xp_name
+        )
+    return None
+
+
+def infer_device(x, xp):
+    """
+    Best-effort device inference that avoids non-portable identifiers.
+
+    Returns None for numpy/jax backends; returns the backend device object
+    for torch/cupy if available.
+    """
+    if xp is None or is_numpy_namespace(xp) or is_jax_namespace(xp):
+        return None
+    try:
+        from array_api_compat import device
+
+        return device(x)
+    except Exception:
+        return None
+
+
+def safe_to_device(x, device, xp):
+    """
+    Move to device if specified; otherwise return input.
+
+    Skips moves for numpy/jax/None devices; logs and returns input on failure.
+    """
+    if device is None:
+        return x
+    if xp is None or is_numpy_namespace(xp) or is_jax_namespace(xp):
+        return x
+    try:
+        return to_device(x, device)
+    except Exception:
+        logger.warning(
+            "Failed to move array to device %s; leaving on current device",
+            device,
+        )
+        return x
+
+
 def recursively_save_to_h5_file(h5_file, path, dictionary):
     """Save a dictionary to an HDF5 file with flattened keys under a given group path."""
     # Ensure the group exists (or open it if already present)

{aspire_inference-0.1.0a10.dist-info → aspire_inference-0.1.0a11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a10
+Version: 0.1.0a11
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT

{aspire_inference-0.1.0a10.dist-info → aspire_inference-0.1.0a11.dist-info}/RECORD

@@ -1,28 +1,28 @@
 aspire/__init__.py,sha256=B2IETvlpB0oBh57prRYLdi8jB5yFGw8qVviGdf1NcnE,409
-aspire/aspire.py,sha256=M5o-QxLthE_5daa1trgUfWxPz-g4rmpEUKimKosw4lw,17400
+aspire/aspire.py,sha256=lr0bD5GDWdlAGfODGzj4BoELKUF5HYAMb8yYGvRR_y0,30860
 aspire/history.py,sha256=l_j-riZKbTWK7Wz9zvvD_mTk9psNCKItiveYhr_pYv8,4313
 aspire/plot.py,sha256=oXwUDOb_953_ADm2KLk41JIfpE3JeiiQiSYKvUVwLqw,1423
-aspire/samples.py,sha256=z5x5hpWuVFH1hYhltmROAe8pbWxGD2UvHi3vcc132dg,19399
+aspire/samples.py,sha256=v7y8DkirUCHOJbCE-o9y2K7xzU2HicIo_O0CdFhLgXE,19478
 aspire/transforms.py,sha256=CHrfPQHEyHQ9I0WWiAgUplWwxYypZsD7uCYIHUbSFtY,24974
-aspire/utils.py,sha256=eKGmchpuoNL15Xbu-AGoeZ00PcQEykEQiDZMnnRyV6A,24234
+aspire/utils.py,sha256=87avRkTce9QvELpIcqlKauSoUZSYi1fqe1asC97TzqA,26947
 aspire/flows/__init__.py,sha256=GUZToPVNJoTwULpbeW10UijfQukNrILoAQ_ubeq7G3w,2110
 aspire/flows/base.py,sha256=5UWKAiXDXLJ6Sg6a380ajLrGFaZSQyOnFEihQiiA4ko,2237
 aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
 aspire/flows/jax/flows.py,sha256=1HnVgQ1GUXNcvxiZqEV19H2QI9Th5bWX_QbNfGaUhuA,6625
 aspire/flows/jax/utils.py,sha256=5T6UrgpARG9VywC9qmTl45LjyZWuEdkW3XUladE6xJE,1518
 aspire/flows/torch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/flows/torch/flows.py,sha256=0_YkiMT49QolyQnEFsh28tfKLnURVF0Z6aTnaWLIUDI,11672
+aspire/flows/torch/flows.py,sha256=QcQOcFZEsLWHPwbQUFGOFdfEslyc59Vf_UEsS0xAGPo,11673
 aspire/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/base.py,sha256=VEHawyVA33jXHMo63p5hBHkp9k2qxU_bOxh5iaZSXew,3011
+aspire/samplers/base.py,sha256=ygqrvqSedWSb0cz8DQ_MHokOOxi6aBRdHxf_qoEPwUE,8243
 aspire/samplers/importance.py,sha256=opn_jY-V8snUz0JztLBtnaTT3WfrZ5OSpHBV5WAuM3M,676
 aspire/samplers/mcmc.py,sha256=ihHgzqvSoy1oxdFBjyqNUbCuRX7CqWjlshCUZcgEL5E,5151
 aspire/samplers/smc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/smc/base.py,sha256=66f_ORUvcKRqMIW35qjhUc-c0PFuY87lJa91MpSaTZI,10729
-aspire/samplers/smc/blackjax.py,sha256=IcTguAETiPmgFofmVW2GN40P5HBIxkmyd2VR8AU8f4k,12115
-aspire/samplers/smc/emcee.py,sha256=Wm0vvAlCcRhJMBt7_fU2ZnjDb8SN8jgUOTXLzNstRpA,2516
-aspire/samplers/smc/minipcn.py,sha256=iQUBBwHZ_D4CxNjARMngklRvx6yTlEDKdeidyYCgqM4,3003
-aspire_inference-0.1.0a10.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
-aspire_inference-0.1.0a10.dist-info/METADATA,sha256=SdvlXKjQn0uJQPpWzoZKAH3oMJSKpZnvlUzxPsIwNlY,3869
-aspire_inference-0.1.0a10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aspire_inference-0.1.0a10.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
-aspire_inference-0.1.0a10.dist-info/RECORD,,
+aspire/samplers/smc/base.py,sha256=40A9yVuKS1F8cPzbfVQ9rNk3y07mnkfbuyRDIh_fy5A,14122
+aspire/samplers/smc/blackjax.py,sha256=2riWDSRmpL5lGmnhNtdieiRs0oYC6XZA2X-nVlQaqpE,12490
+aspire/samplers/smc/emcee.py,sha256=4CI9GvH69FCoLiFBbKKYwYocYyiM95IijC5EvrcAmUo,2891
+aspire/samplers/smc/minipcn.py,sha256=IJ5466VvARd4qZCWXXl-l3BPaKW1AgcwmbP3ISL2bto,3368
+aspire_inference-0.1.0a11.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
+aspire_inference-0.1.0a11.dist-info/METADATA,sha256=lLd2d5HR-t942wKLyYbdJ1DL9CKl7tkpBul_vX8DU4M,3869
+aspire_inference-0.1.0a11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aspire_inference-0.1.0a11.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
+aspire_inference-0.1.0a11.dist-info/RECORD,,