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

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

@@ -3,17 +3,21 @@ from functools import partial
 import numpy as np
 
 from ...samples import SMCSamples
-from ...utils import to_numpy, track_calls
-from .base import NumpySMCSampler
+from ...utils import (
+    asarray,
+    determine_backend_name,
+    track_calls,
+)
+from .base import SMCSampler
 
 
-class MiniPCNSMC(NumpySMCSampler):
+class MiniPCNSMC(SMCSampler):
     """MiniPCN SMC sampler."""
 
     rng = None
 
     def log_prob(self, x, beta=None):
-        return to_numpy(super().log_prob(x, beta))
+        return super().log_prob(x, beta)
 
     @track_calls
     def sample(
@@ -28,12 +32,19 @@ 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
+
         self.sampler_kwargs = sampler_kwargs or {}
         self.sampler_kwargs.setdefault("n_steps", 5 * self.dims)
         self.sampler_kwargs.setdefault("target_acceptance_rate", 0.234)
         self.sampler_kwargs.setdefault("step_fn", "tpcn")
-        self.rng = rng or np.random.default_rng()
+        self.backend_str = determine_backend_name(xp=self.xp)
+        self.rng = rng or ArrayRNG(backend=self.backend_str)
         return super().sample(
             n_samples,
             n_steps=n_steps,
@@ -43,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):
@@ -58,9 +73,14 @@ class MiniPCNSMC(NumpySMCSampler):
             target_acceptance_rate=self.sampler_kwargs[
                 "target_acceptance_rate"
             ],
+            xp=self.xp,
         )
         # Map to transformed dimension for sampling
-        z = to_numpy(self.fit_preconditioning_transform(particles.x))
+        z = asarray(
+            self.fit_preconditioning_transform(particles.x),
+            xp=self.xp,
+            dtype=self.dtype,
+        )
         chain, history = sampler.sample(
             z,
             n_steps=n_steps or self.sampler_kwargs["n_steps"],

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)
 
 
@@ -425,19 +427,23 @@ class Samples(BaseSamples):
 
     def to_namespace(self, xp):
         return self.__class__(
-            x=asarray(self.x, xp),
+            x=asarray(self.x, xp, dtype=self.dtype),
             parameters=self.parameters,
-            log_likelihood=asarray(self.log_likelihood, xp)
+            log_likelihood=asarray(self.log_likelihood, xp, dtype=self.dtype)
             if self.log_likelihood is not None
             else None,
-            log_prior=asarray(self.log_prior, xp)
+            log_prior=asarray(self.log_prior, xp, dtype=self.dtype)
             if self.log_prior is not None
             else None,
-            log_q=asarray(self.log_q, xp) if self.log_q is not None else None,
-            log_evidence=asarray(self.log_evidence, xp)
+            log_q=asarray(self.log_q, xp, dtype=self.dtype)
+            if self.log_q is not None
+            else None,
+            log_evidence=asarray(self.log_evidence, xp, dtype=self.dtype)
             if self.log_evidence is not None
             else None,
-            log_evidence_error=asarray(self.log_evidence_error, xp)
+            log_evidence_error=asarray(
+                self.log_evidence_error, xp, dtype=self.dtype
+            )
             if self.log_evidence_error is not None
             else None,
         )

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
@@ -12,7 +14,13 @@ import h5py
 import wrapt
 from array_api_compat import (
     array_namespace,
+    is_cupy_namespace,
+    is_dask_namespace,
     is_jax_array,
+    is_jax_namespace,
+    is_ndonnx_namespace,
+    is_numpy_namespace,
+    is_pydata_sparse_namespace,
     is_torch_array,
     is_torch_namespace,
     to_device,
@@ -28,6 +36,17 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+IS_NAMESPACE_FUNCTIONS = {
+    "numpy": is_numpy_namespace,
+    "torch": is_torch_namespace,
+    "jax": is_jax_namespace,
+    "cupy": is_cupy_namespace,
+    "dask": is_dask_namespace,
+    "pydata_sparse": is_pydata_sparse_namespace,
+    "ndonnx": is_ndonnx_namespace,
+}
+
+
 def configure_logger(
     log_level: str | int = "INFO",
     additional_loggers: list[str] = None,
@@ -234,7 +253,7 @@ def to_numpy(x: Array, **kwargs) -> np.ndarray:
         return np.asarray(x, **kwargs)
 
 
-def asarray(x, xp: Any = None, **kwargs) -> Array:
+def asarray(x, xp: Any = None, dtype: Any | None = None, **kwargs) -> Array:
     """Convert an array to the specified array API.
 
     Parameters
@@ -244,13 +263,51 @@ def asarray(x, xp: Any = None, **kwargs) -> Array:
     xp : Any
         The array API to use for the conversion. If None, the array API
         is inferred from the input array.
+    dtype : Any | str | None
+        The dtype to use for the conversion. If None, the dtype is not changed.
     kwargs : dict
         Additional keyword arguments to pass to xp.asarray.
     """
+    # Handle DLPack conversion from JAX to PyTorch to avoid shape issues when
+    # passing JAX arrays directly to torch.asarray.
     if is_jax_array(x) and is_torch_namespace(xp):
-        return xp.utils.dlpack.from_dlpack(x)
-    else:
-        return xp.asarray(x, **kwargs)
+        tensor = xp.utils.dlpack.from_dlpack(x)
+        if dtype is not None:
+            tensor = tensor.to(resolve_dtype(dtype, xp=xp))
+        return tensor
+
+    if dtype is not None:
+        kwargs["dtype"] = resolve_dtype(dtype, xp=xp)
+    return xp.asarray(x, **kwargs)
+
+
+def determine_backend_name(
+    x: Array | None = None, xp: Any | None = None
+) -> str:
+    """Determine the backend name from an array or array API module.
+
+    Parameters
+    ----------
+    x : Array or None
+        The array to infer the backend from. If None, xp must be provided.
+    xp : Any or None
+        The array API module to infer the backend from. If None, x must be provided.
+
+    Returns
+    -------
+    str
+        The name of the backend. If the backend cannot be determined, returns "unknown".
+    """
+    if x is not None:
+        xp = array_namespace(x)
+    if xp is None:
+        raise ValueError(
+            "Either x or xp must be provided to determine backend."
+        )
+    for name, is_namespace_fn in IS_NAMESPACE_FUNCTIONS.items():
+        if is_namespace_fn(xp):
+            return name
+    return "unknown"
 
 
 def resolve_dtype(dtype: Any | str | None, xp: Any) -> Any | None:
@@ -546,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
@@ -574,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)