aspire-inference 0.1.0a5__tar.gz → 0.1.0a6__tar.gz

{aspire_inference-0.1.0a5/aspire_inference.egg-info → aspire_inference-0.1.0a6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a5
+Version: 0.1.0a6
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -33,6 +33,7 @@ Requires-Dist: blackjax; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6/aspire_inference.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a5
+Version: 0.1.0a6
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -33,6 +33,7 @@ Requires-Dist: blackjax; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/aspire_inference.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ aspire_inference.egg-info/dependency_links.txt
 aspire_inference.egg-info/requires.txt
 aspire_inference.egg-info/top_level.txt
 examples/basic_example.py
+examples/smc_example.py
 src/aspire/__init__.py
 src/aspire/aspire.py
 src/aspire/history.py
@@ -36,6 +37,9 @@ src/aspire/samplers/smc/blackjax.py
 src/aspire/samplers/smc/emcee.py
 src/aspire/samplers/smc/minipcn.py
 tests/conftest.py
+tests/test_samples.py
+tests/test_transforms.py
+tests/test_utils.py
 tests/integration_tests/conftest.py
 tests/integration_tests/test_integration.py
 tests/test_flows/test_jax_flows/test_flowjax_flows.py

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/aspire_inference.egg-info/requires.txt

@@ -24,6 +24,7 @@ scipy
 [test]
 pytest
 pytest-requires
+pytest-cov
 
 [torch]
 torch

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/examples/basic_example.py

@@ -6,11 +6,12 @@ likelihood with a uniform prior.
 import math
 from pathlib import Path
 
+from scipy.stats import norm, uniform
+
 from aspire import Aspire
 from aspire.plot import plot_comparison
 from aspire.samples import Samples
 from aspire.utils import AspireFile, configure_logger
-from scipy.stats import norm, uniform
 
 # Configure the logger
 configure_logger("INFO")
@@ -71,6 +72,8 @@ with AspireFile(outdir / "aspire_result.h5", "w") as f:
     aspire.save_config(f, "aspire_config")
     samples.save(f, "posterior_samples")
     history.save(f, "flow_history")
+    # Save the flow
+    aspire.save_flow(f, "flow")
 
 fig = plot_comparison(
     initial_samples,

aspire_inference-0.1.0a6/examples/smc_example.py

@@ -0,0 +1,110 @@
+"""Example using sequential posterior inference with SMC.
+
+This examples is slightly contrived, using a mixture of two Gaussians in 4D
+as the target distribution. The goal is to demonstrate the ability of SMC to
+explore multi-modal distributions, even when the initial samples deviate
+significantly from the true modes.
+
+In practice, one would ideally use more informative initial samples.
+"""
+
+from pathlib import Path
+
+import numpy as np
+
+from aspire import Aspire
+from aspire.plot import plot_comparison
+from aspire.samples import Samples
+from aspire.utils import configure_logger
+
+# RNG for generating initial samples
+rng = np.random.default_rng(42)
+
+# Output directory
+outdir = Path("outdir") / "smc_example"
+outdir.mkdir(parents=True, exist_ok=True)
+
+# Configure logger to show INFO level messages
+configure_logger()
+
+# Number of dimensions
+dims = 4
+
+# Means and covariances of the two Gaussian components
+mu1 = 2 * np.ones(dims)
+mu2 = -2 * np.ones(dims)
+cov1 = 0.5 * np.eye(dims)
+cov2 = np.eye(dims)
+
+
+def log_likelihood(samples):
+    """Log-likelihood of a mixture of two Gaussians"""
+    x = samples.x
+    comp1 = (
+        -0.5 * ((x - mu1) @ np.linalg.inv(cov1) * (x - mu1)).sum(axis=-1)
+        - 0.5 * dims * np.log(2 * np.pi)
+        - 0.5 * np.linalg.slogdet(cov1)[1]
+    )
+    comp2 = (
+        -0.5 * ((x - mu2) @ np.linalg.inv(cov2) * (x - mu2)).sum(axis=-1)
+        - 0.5 * dims * np.log(2 * np.pi)
+        - 0.5 * np.linalg.slogdet(cov2)[1]
+    )
+    return np.logaddexp(comp1, comp2)  # Log-sum-exp for numerical stability
+
+
+def log_prior(samples):
+    """Standard normal prior"""
+    return -0.5 * (samples.x**2).sum(axis=-1) - dims * 0.5 * np.log(2 * np.pi)
+
+
+# Generate prior samples for comparison, these are not used in SMC
+prior_samples = Samples(rng.normal(0, 1, size=(5000, dims)))
+
+# We draw initial samples from two Gaussians centered away from the true modes
+# to demonstrate the ability of SMC to explore the posterior
+offset_1 = rng.uniform(-3, 3, size=(dims,))
+offset_2 = rng.uniform(-3, 3, size=(dims,))
+initial_samples = np.concatenate(
+    [
+        rng.normal(mu1 - offset_1, 1, size=(2500, dims)),
+        rng.normal(mu2 - offset_2, 1, size=(2500, dims)),
+    ],
+    axis=0,
+)
+initial_samples = Samples(initial_samples)
+
+# Initialize Aspire with the log-likelihood and log-prior
+aspire = Aspire(
+    log_likelihood=log_likelihood,
+    log_prior=log_prior,
+    dims=dims,
+    flow_class="NSF",  # Use Neural Spline Flow from zuko (default backend)
+)
+
+# Fit the normalizing flow to the initial samples
+fit_history = aspire.fit(initial_samples, n_epochs=30)
+# Plot loss
+fit_history.plot_loss().savefig(outdir / "loss.png")
+
+# Sample from the posterior using SMC
+samples, history = aspire.sample_posterior(
+    sampler="smc",  # Sequential Monte Carlo, this uses the default minipcn sampler
+    n_samples=500,  # Number of particles in SMC
+    n_final_samples=5000,  # Number of samples to draw from the final distribution
+    sampler_kwargs=dict(  # Keyword arguments for the specific sampler
+        n_steps=20,  # MCMC steps per SMC iteration
+    ),
+    return_history=True,  # To return the SMC history (e.g., ESS, betas)
+)
+# Plot SMC diagnostics
+history.plot().savefig(outdir / "smc_diagnostics.png")
+
+# Plot corner plot of the samples
+# Include initial samples and prior samples for comparison
+plot_comparison(
+    initial_samples,
+    prior_samples,
+    samples,
+    labels=["Initial Samples", "Prior Samples", "SMC Samples"],
+).savefig(outdir / "posterior.png")

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/pyproject.toml

@@ -50,6 +50,7 @@ blackjax = [
 test = [
     "pytest",
     "pytest-requires",
+    "pytest-cov",
 ]
 
 [project.urls]
@@ -69,3 +70,7 @@ target-version = "py39"
 # Allow fix for all enabled rules (when `--fix`) is provided.
 fixable = ["ALL"]
 extend-select = ["I"]
+
+[tool.pytest.ini_options]
+addopts = "--cov=aspire --cov-report=term-missing -ra"
+testpaths = ["tests"]

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/src/aspire/aspire.py

@@ -6,6 +6,7 @@ from typing import Any, Callable
 import h5py
 
 from .flows import get_flow_wrapper
+from .flows.base import Flow
 from .history import History
 from .samples import Samples
 from .transforms import (
@@ -48,12 +49,17 @@ class Aspire:
     xp : Callable | None
         The array backend to use. If None, the default backend will be
         used.
+    flow : Flow | None
+        The flow object, if it already exists.
+        If None, a new flow will be created.
     flow_backend : str
         The backend to use for the flow. Options are 'zuko' or 'flowjax'.
     flow_matching : bool
         Whether to use flow matching.
     eps : float
         The epsilon value to use for data transforms.
+    dtype : Any | str | None
+        The data type to use for the samples, flow and transforms.
     **kwargs
         Keyword arguments to pass to the flow.
     """
@@ -71,9 +77,11 @@ class Aspire:
         bounded_transform: str = "logit",
         device: str | None = None,
         xp: Callable | None = None,
+        flow: Flow | None = None,
         flow_backend: str = "zuko",
         flow_matching: bool = False,
         eps: float = 1e-6,
+        dtype: Any | str | None = None,
         **kwargs,
     ) -> None:
         self.log_likelihood = log_likelihood
@@ -91,14 +99,20 @@ class Aspire:
         self.flow_backend = flow_backend
         self.flow_kwargs = kwargs
         self.xp = xp
+        self.dtype = dtype
 
-        self._flow = None
+        self._flow = flow
 
     @property
     def flow(self):
         """The normalizing flow object."""
         return self._flow
 
+    @flow.setter
+    def flow(self, flow: Flow):
+        """Set the normalizing flow object."""
+        self._flow = flow
+
     @property
     def sampler(self):
         """The sampler object."""
@@ -130,6 +144,7 @@ class Aspire:
             log_prior=log_prior,
             log_q=log_q,
             xp=xp,
+            dtype=self.dtype,
         )
 
         if evaluate:
@@ -159,6 +174,7 @@ class Aspire:
             device=self.device,
             xp=xp,
             eps=self.eps,
+            dtype=self.dtype,
         )
 
         # Check if FlowClass takes `parameters` as an argument
@@ -172,6 +188,7 @@ class Aspire:
             dims=self.dims,
             device=self.device,
             data_transform=data_transform,
+            dtype=self.dtype,
             **self.flow_kwargs,
         )
 
@@ -245,6 +262,7 @@ class Aspire:
                 periodic_parameters=self.periodic_parameters,
                 xp=self.xp,
                 device=self.device,
+                dtype=self.dtype,
                 **preconditioning_kwargs,
             )
         elif preconditioning == "flow":
@@ -259,6 +277,7 @@ class Aspire:
                 bounded_to_unbounded=self.bounded_to_unbounded,
                 prior_bounds=self.prior_bounds,
                 xp=self.xp,
+                dtype=self.dtype,
                 device=self.device,
                 **preconditioning_kwargs,
             )
@@ -271,6 +290,7 @@ class Aspire:
             dims=self.dims,
             prior_flow=self.flow,
             xp=self.xp,
+            dtype=self.dtype,
             preconditioning_transform=transform,
             **kwargs,
         )
@@ -397,17 +417,17 @@ class Aspire:
             method of the sampler.
         """
         config = {
-            # "log_likelihood": self.log_likelihood,
-            # "log_prior": self.log_prior,
+            "log_likelihood": self.log_likelihood.__name__,
+            "log_prior": self.log_prior.__name__,
             "dims": self.dims,
             "parameters": self.parameters,
             "periodic_parameters": self.periodic_parameters,
             "prior_bounds": self.prior_bounds,
             "bounded_to_unbounded": self.bounded_to_unbounded,
-            # "bounded_transform": self.bounded_transform,
+            "bounded_transform": self.bounded_transform,
             "flow_matching": self.flow_matching,
-            # "device": self.device,
-            # "xp": self.xp,
+            "device": self.device,
+            "xp": self.xp.__name__ if self.xp else None,
             "flow_backend": self.flow_backend,
             "flow_kwargs": self.flow_kwargs,
             "eps": self.eps,
@@ -437,6 +457,35 @@ class Aspire:
             self.config_dict(**kwargs),
         )
 
+    def save_flow(self, h5_file: h5py.File, path="flow") -> None:
+        """Save the flow to an HDF5 file.
+
+        Parameters
+        ----------
+        h5_file : h5py.File
+            The HDF5 file to save the flow to.
+        path : str
+            The path in the HDF5 file to save the flow to.
+        """
+        if self.flow is None:
+            raise ValueError("Flow has not been initialized.")
+        self.flow.save(h5_file, path=path)
+
+    def load_flow(self, h5_file: h5py.File, path="flow") -> None:
+        """Load the flow from an HDF5 file.
+
+        Parameters
+        ----------
+        h5_file : h5py.File
+            The HDF5 file to load the flow from.
+        path : str
+            The path in the HDF5 file to load the flow from.
+        """
+        FlowClass, xp = get_flow_wrapper(
+            backend=self.flow_backend, flow_matching=self.flow_matching
+        )
+        self._flow = FlowClass.load(h5_file, path=path)
+
     def save_config_to_json(self, filename: str) -> None:
         """Save the configuration to a JSON file."""
         import json

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/src/aspire/flows/base.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 from typing import Any
 
@@ -45,3 +46,39 @@ class Flow:
 
     def inverse_rescale(self, x):
         return self.data_transform.inverse(x)
+
+    def config_dict(self):
+        """Return a dictionary of the configuration of the flow.
+
+        This can be used to recreate the flow by passing the dictionary
+        as keyword arguments to the constructor.
+
+        This is automatically populated with the arguments passed to the
+        constructor.
+
+        Returns
+        -------
+        config : dict
+            The configuration dictionary.
+        """
+        return getattr(self, "_init_args", {})
+
+    def save(self, h5_file, path="flow"):
+        raise NotImplementedError
+
+    @classmethod
+    def load(cls, h5_file, path="flow"):
+        raise NotImplementedError
+
+    def __new__(cls, *args, **kwargs):
+        # Create instance
+        obj = super().__new__(cls)
+        # Inspect the subclass's __init__ signature
+        sig = inspect.signature(cls.__init__)
+        bound = sig.bind_partial(obj, *args, **kwargs)
+        bound.apply_defaults()
+        # Save args (excluding self)
+        obj._init_args = {
+            k: v for k, v in bound.arguments.items() if k != "self"
+        }
+        return obj

aspire_inference-0.1.0a6/src/aspire/flows/jax/flows.py

@@ -0,0 +1,196 @@
+import logging
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from flowjax.train import fit_to_data
+
+from ...transforms import IdentityTransform
+from ...utils import decode_dtype, encode_dtype, resolve_dtype
+from ..base import Flow
+from .utils import get_flow
+
+logger = logging.getLogger(__name__)
+
+
+class FlowJax(Flow):
+    xp = jnp
+
+    def __init__(
+        self,
+        dims: int,
+        key=None,
+        data_transform=None,
+        dtype=None,
+        **kwargs,
+    ):
+        device = kwargs.pop("device", None)
+        if device is not None:
+            logger.warning("The device argument is not used in FlowJax. ")
+        resolved_dtype = (
+            resolve_dtype(dtype, jnp)
+            if dtype is not None
+            else jnp.dtype(jnp.float32)
+        )
+        if data_transform is None:
+            data_transform = IdentityTransform(self.xp, dtype=resolved_dtype)
+        elif getattr(data_transform, "dtype", None) is None:
+            data_transform.dtype = resolved_dtype
+        super().__init__(dims, device=device, data_transform=data_transform)
+        self.dtype = resolved_dtype
+        if key is None:
+            key = jrandom.key(0)
+            logger.warning(
+                "The key argument is None. "
+                "A random key will be used for the flow. "
+                "Results may not be reproducible."
+            )
+        self.key = key
+        self.loc = None
+        self.scale = None
+        self.key, subkey = jrandom.split(self.key)
+        self._flow = get_flow(
+            key=subkey,
+            dims=self.dims,
+            dtype=self.dtype,
+            **kwargs,
+        )
+
+    def fit(self, x, **kwargs):
+        from ...history import FlowHistory
+
+        x = jnp.asarray(x, dtype=self.dtype)
+        x_prime = jnp.asarray(self.fit_data_transform(x), dtype=self.dtype)
+        self.key, subkey = jrandom.split(self.key)
+        self._flow, losses = fit_to_data(subkey, self._flow, x_prime, **kwargs)
+        return FlowHistory(
+            training_loss=list(losses["train"]),
+            validation_loss=list(losses["val"]),
+        )
+
+    def forward(self, x, xp: Callable = jnp):
+        x = jnp.asarray(x, dtype=self.dtype)
+        x_prime, log_abs_det_jacobian = self.rescale(x)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
+        z, log_abs_det_jacobian_flow = self._flow.forward(x_prime)
+        return xp.asarray(z), xp.asarray(
+            log_abs_det_jacobian + log_abs_det_jacobian_flow
+        )
+
+    def inverse(self, z, xp: Callable = jnp):
+        z = jnp.asarray(z, dtype=self.dtype)
+        x_prime, log_abs_det_jacobian_flow = self._flow.inverse(z)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
+        x, log_abs_det_jacobian = self.inverse_rescale(x_prime)
+        return xp.asarray(x), xp.asarray(
+            log_abs_det_jacobian + log_abs_det_jacobian_flow
+        )
+
+    def log_prob(self, x, xp: Callable = jnp):
+        x = jnp.asarray(x, dtype=self.dtype)
+        x_prime, log_abs_det_jacobian = self.rescale(x)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
+        log_prob = self._flow.log_prob(x_prime)
+        return xp.asarray(log_prob + log_abs_det_jacobian)
+
+    def sample(self, n_samples: int, xp: Callable = jnp):
+        self.key, subkey = jrandom.split(self.key)
+        x_prime = self._flow.sample(subkey, (n_samples,))
+        x = self.inverse_rescale(x_prime)[0]
+        return xp.asarray(x)
+
+    def sample_and_log_prob(self, n_samples: int, xp: Callable = jnp):
+        self.key, subkey = jrandom.split(self.key)
+        x_prime = self._flow.sample(subkey, (n_samples,))
+        log_prob = self._flow.log_prob(x_prime)
+        x, log_abs_det_jacobian = self.inverse_rescale(x_prime)
+        return xp.asarray(x), xp.asarray(log_prob - log_abs_det_jacobian)
+
+    def save(self, h5_file, path="flow"):
+        import equinox as eqx
+        from array_api_compat import numpy as np
+
+        from ...utils import recursively_save_to_h5_file
+
+        grp = h5_file.require_group(path)
+
+        # ---- config ----
+        config = self.config_dict().copy()
+        config.pop("key", None)
+        config["key_data"] = jax.random.key_data(self.key)
+        dtype_value = config.get("dtype")
+        if dtype_value is None:
+            dtype_value = self.dtype
+        else:
+            dtype_value = jnp.dtype(dtype_value)
+        print(dtype_value)
+        config["dtype"] = encode_dtype(jnp, dtype_value)
+
+        data_transform = config.pop("data_transform", None)
+        if data_transform is not None:
+            data_transform.save(grp, "data_transform")
+
+        recursively_save_to_h5_file(grp, "config", config)
+
+        # ---- save arrays ----
+        arrays, _ = eqx.partition(self._flow, eqx.is_array)
+        leaves, _ = jax.tree_util.tree_flatten(arrays)
+
+        params_grp = grp.require_group("params")
+        # clear old datasets
+        for name in list(params_grp.keys()):
+            del params_grp[name]
+
+        for i, p in enumerate(leaves):
+            params_grp.create_dataset(str(i), data=np.asarray(p))
+
+    @classmethod
+    def load(cls, h5_file, path="flow"):
+        import equinox as eqx
+
+        from ...utils import load_from_h5_file
+
+        grp = h5_file[path]
+
+        # ---- config ----
+        config = load_from_h5_file(grp, "config")
+        config["dtype"] = decode_dtype(jnp, config.get("dtype"))
+        if "data_transform" in grp:
+            from ...transforms import BaseTransform
+
+            config["data_transform"] = BaseTransform.load(
+                grp,
+                "data_transform",
+                strict=False,
+            )
+
+        key_data = config.pop("key_data", None)
+        if key_data is not None:
+            config["key"] = jax.random.wrap_key_data(key_data)
+
+        kwargs = config.pop("kwargs", {})
+        config.update(kwargs)
+
+        # build object (will replace its _flow)
+        obj = cls(**config)
+
+        # ---- load arrays ----
+        params_grp = grp["params"]
+        loaded_params = [
+            jnp.array(params_grp[str(i)][:]) for i in range(len(params_grp))
+        ]
+
+        # rebuild template flow
+        kwargs.pop("device")
+        flow_template = get_flow(key=jrandom.key(0), dims=obj.dims, **kwargs)
+        arrays_template, static = eqx.partition(flow_template, eqx.is_array)
+
+        # use treedef from template
+        treedef = jax.tree_util.tree_structure(arrays_template)
+        arrays = jax.tree_util.tree_unflatten(treedef, loaded_params)
+
+        # recombine
+        obj._flow = eqx.combine(static, arrays)
+
+        return obj

{aspire_inference-0.1.0a5 → aspire_inference-0.1.0a6}/src/aspire/flows/jax/utils.py

@@ -29,8 +29,11 @@ def get_flow(
     flow_type: str | Callable = "masked_autoregressive_flow",
     bijection_type: str | flowjax.bijections.AbstractBijection | None = None,
     bijection_kwargs: dict | None = None,
+    dtype=None,
     **kwargs,
 ) -> flowjax.distributions.Transformed:
+    dtype = dtype or jnp.float32
+
     if isinstance(flow_type, str):
         flow_type = get_flow_function_class(flow_type)
 
@@ -44,7 +47,7 @@ def get_flow(
     if bijection_kwargs is None:
         bijection_kwargs = {}
 
-    base_dist = flowjax.distributions.Normal(jnp.zeros(dims))
+    base_dist = flowjax.distributions.Normal(jnp.zeros(dims, dtype=dtype))
     key, subkey = jrandom.split(key)
     return flow_type(
         subkey,