aspire-inference 0.1.0a8__tar.gz → 0.1.0a9__tar.gz

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a8
+Version: 0.1.0a9
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -39,6 +39,12 @@ Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse
 
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.15658747.svg)](https://doi.org/10.5281/zenodo.15658747)
+[![PyPI](https://img.shields.io/pypi/v/aspire-inference)](https://pypi.org/project/aspire-inference/)
+[![Documentation Status](https://readthedocs.org/projects/aspire/badge/?version=latest)](https://aspire.readthedocs.io/en/latest/?badge=latest)
+![tests](https://github.com/mj-will/aspire/actions/workflows/tests.yml/badge.svg)
+
+
 aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
@@ -52,6 +58,45 @@ pip install aspire-inference
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
 the package can be imported and used as `aspire`.
 
+## Quickstart
+
+```python
+import numpy as np
+from aspire import Aspire, Samples
+
+# Define a log-likelihood and log-prior
+def log_likelihood(samples):
+    x = samples.x
+    return -0.5 * np.sum(x**2, axis=-1)
+
+def log_prior(samples):
+    return -0.5 * np.sum(samples.x**2, axis=-1)
+
+# Create the initial samples
+init = Samples(np.random.normal(size=(2_000, 4)))
+
+# Define the aspire object
+aspire = Aspire(
+    log_likelihood=log_likelihood,
+    log_prior=log_prior,
+    dims=4,
+    parameters=[f"x{i}" for i in range(4)],
+)
+
+# Fit the normalizing flow
+aspire.fit(init, n_epochs=20)
+
+# Sample the posterior
+posterior = aspire.sample_posterior(
+    sampler="smc",
+    n_samples=500,
+    sampler_kwargs=dict(n_steps=100),
+)
+
+# Plot the posterior distribution
+posterior.plot_corner()
+```
+
 ## Documentation
 
 See the [documentation on ReadTheDocs][docs].

aspire_inference-0.1.0a9/README.md

@@ -0,0 +1,72 @@
+# aspire: Accelerated Sequential Posterior Inference via REuse
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.15658747.svg)](https://doi.org/10.5281/zenodo.15658747)
+[![PyPI](https://img.shields.io/pypi/v/aspire-inference)](https://pypi.org/project/aspire-inference/)
+[![Documentation Status](https://readthedocs.org/projects/aspire/badge/?version=latest)](https://aspire.readthedocs.io/en/latest/?badge=latest)
+![tests](https://github.com/mj-will/aspire/actions/workflows/tests.yml/badge.svg)
+
+
+aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
+
+## Installation
+
+aspire can be installed from PyPI using `pip`
+
+```
+pip install aspire-inference
+```
+
+**Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
+the package can be imported and used as `aspire`.
+
+## Quickstart
+
+```python
+import numpy as np
+from aspire import Aspire, Samples
+
+# Define a log-likelihood and log-prior
+def log_likelihood(samples):
+    x = samples.x
+    return -0.5 * np.sum(x**2, axis=-1)
+
+def log_prior(samples):
+    return -0.5 * np.sum(samples.x**2, axis=-1)
+
+# Create the initial samples
+init = Samples(np.random.normal(size=(2_000, 4)))
+
+# Define the aspire object
+aspire = Aspire(
+    log_likelihood=log_likelihood,
+    log_prior=log_prior,
+    dims=4,
+    parameters=[f"x{i}" for i in range(4)],
+)
+
+# Fit the normalizing flow
+aspire.fit(init, n_epochs=20)
+
+# Sample the posterior
+posterior = aspire.sample_posterior(
+    sampler="smc",
+    n_samples=500,
+    sampler_kwargs=dict(n_steps=100),
+)
+
+# Plot the posterior distribution
+posterior.plot_corner()
+```
+
+## Documentation
+
+See the [documentation on ReadTheDocs][docs].
+
+## Citation
+
+If you use `aspire` in your work please cite the [DOI][DOI] and [paper][paper].
+
+
+[docs]: https://aspire.readthedocs.io/
+[DOI]: https://doi.org/10.5281/zenodo.15658747
+[paper]: https://arxiv.org/abs/2511.04218

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/aspire_inference.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a8
+Version: 0.1.0a9
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -39,6 +39,12 @@ Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse
 
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.15658747.svg)](https://doi.org/10.5281/zenodo.15658747)
+[![PyPI](https://img.shields.io/pypi/v/aspire-inference)](https://pypi.org/project/aspire-inference/)
+[![Documentation Status](https://readthedocs.org/projects/aspire/badge/?version=latest)](https://aspire.readthedocs.io/en/latest/?badge=latest)
+![tests](https://github.com/mj-will/aspire/actions/workflows/tests.yml/badge.svg)
+
+
 aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
@@ -52,6 +58,45 @@ pip install aspire-inference
 **Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
 the package can be imported and used as `aspire`.
 
+## Quickstart
+
+```python
+import numpy as np
+from aspire import Aspire, Samples
+
+# Define a log-likelihood and log-prior
+def log_likelihood(samples):
+    x = samples.x
+    return -0.5 * np.sum(x**2, axis=-1)
+
+def log_prior(samples):
+    return -0.5 * np.sum(samples.x**2, axis=-1)
+
+# Create the initial samples
+init = Samples(np.random.normal(size=(2_000, 4)))
+
+# Define the aspire object
+aspire = Aspire(
+    log_likelihood=log_likelihood,
+    log_prior=log_prior,
+    dims=4,
+    parameters=[f"x{i}" for i in range(4)],
+)
+
+# Fit the normalizing flow
+aspire.fit(init, n_epochs=20)
+
+# Sample the posterior
+posterior = aspire.sample_posterior(
+    sampler="smc",
+    n_samples=500,
+    sampler_kwargs=dict(n_steps=100),
+)
+
+# Plot the posterior distribution
+posterior.plot_corner()
+```
+
 ## Documentation
 
 See the [documentation on ReadTheDocs][docs].

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/aspire_inference.egg-info/SOURCES.txt

@@ -13,15 +13,17 @@ aspire_inference.egg-info/dependency_links.txt
 aspire_inference.egg-info/requires.txt
 aspire_inference.egg-info/top_level.txt
 docs/Makefile
-docs/api.rst
 docs/conf.py
 docs/entry_points.rst
 docs/examples.rst
 docs/index.rst
 docs/installation.rst
+docs/multiprocessing.rst
+docs/recipes.rst
 docs/requirements.txt
 docs/user_guide.rst
 examples/basic_example.py
+examples/blackjax_smc_example.py
 examples/smc_example.py
 src/aspire/__init__.py
 src/aspire/aspire.py

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/docs/conf.py

@@ -29,6 +29,7 @@ extensions = [
     "sphinx.ext.autosummary",
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
+    "autoapi.extension",
 ]
 
 autodoc_typehints = "description"
@@ -41,6 +42,19 @@ napoleon_preprocess_types = True
 templates_path = ["_templates"]
 exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
+# -- Configure autoapi -------------------------------------------------------
+
+autoapi_type = "python"
+autoapi_dirs = ["../src/aspire/"]
+autoapi_add_toctree_entry = True
+autoapi_options = [
+    "members",
+    "imported-members",
+    "show-inheritance",
+    "show-module-summary",
+    "undoc-members",
+]
+
 # -- Options for HTML output -------------------------------------------------
 
 html_theme = "sphinx_book_theme"

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/docs/index.rst

@@ -22,8 +22,7 @@ Quick start
 .. code-block:: python
 
     import numpy as np
-    from aspire import Aspire
-    from aspire.samples import Samples
+    from aspire import Aspire, Samples
 
     def log_likelihood(samples):
         x = samples.x
@@ -58,6 +57,7 @@ examples, and the complete API reference.
 
     installation
     user_guide
+    recipes
+    multiprocessing
     examples
     entry_points
-    api

aspire_inference-0.1.0a9/docs/multiprocessing.rst

@@ -0,0 +1,70 @@
+Multiprocessing
+===============
+
+Use :meth:`aspire.Aspire.enable_pool` to run your likelihood (and optionally
+prior) in parallel across a :class:`multiprocessing.Pool`. The helper swaps the
+``map_fn`` argument expected by your log-likelihood / log-prior for
+``pool.map`` while the context is active, then restores the original methods.
+
+Prepare a map-aware likelihood
+------------------------------
+
+Your likelihood must accept ``map_fn``. A minimal
+pattern:
+
+.. code-block:: python
+
+    import numpy as np
+
+
+    def _global_log_likelihood(x):
+        # Expensive likelihood computation for a single sample `x`
+        return -np.sum(x**2)  # Example likelihood
+
+    def log_likelihood(samples, map_fn=map):
+        logl = -np.inf * np.ones(len(samples.x))
+        if samples.log_prior is None:
+            raise RuntimeError("log-prior has not been evaluated!")
+        mask = np.isfinite(samples.log_prior, dtype=bool)
+        x = np.asarray(samples.x[mask, :], dtype=float)
+        logl[mask] = np.fromiter(
+            map_fn(_global_log_likelihood, x),
+            dtype=float,
+        )
+        return logl
+
+Swap in a multiprocessing pool
+------------------------------
+
+Wrap your sampling call inside ``enable_pool`` to parallelize the map step:
+
+.. code-block:: python
+
+    import multiprocessing as mp
+    from aspire import Aspire
+
+    aspire = Aspire(
+        log_likelihood=log_likelihood,
+        log_prior=log_prior,   # must also accept map_fn if parallelize_prior=True
+        dims=4,
+        parameters=["a", "b", "c", "d"],
+    )
+
+    with mp.Pool() as pool, aspire.enable_pool(pool):
+        samples, history = aspire.sample_posterior(
+            sampler="smc",
+            n_samples=1_000,
+            return_history=True,
+        )
+
+Notes
+-----
+
+- By default only the likelihood is parallelized; set
+  ``aspire.enable_pool(pool, parallelize_prior=True)`` if your prior also
+  accepts ``map_fn``.
+- ``enable_pool`` closes the pool on exit unless you pass ``close_pool=False``.
+- The context manager itself is implemented by
+  :class:`aspire.utils.PoolHandler`; if you need finer control (for example,
+  reusing the same pool across multiple ``Aspire`` instances) you can
+  instantiate it directly.

aspire_inference-0.1.0a9/docs/recipes.rst

@@ -0,0 +1,70 @@
+Practical recipes
+=================
+
+Checking the prior when evaluating the likelihood
+-------------------------------------------------
+
+By default, Aspire samplers always evaluate the log-prior before the
+log-likelihood. This allows users to check the prior support and skip
+likelihood evaluations for samples that lie outside the prior bounds.
+
+.. code-block:: python
+
+    import aspire
+    import numpy as np
+
+
+    def log_likelihood(samples: aspire.Samples) -> np.ndarray:
+        if samples.log_prior is None:
+            raise RuntimeError("log-prior has not been evaluated!")
+        # Return -inf for samples with invalid prior
+        logl = np.full(samples.n_samples, -np.inf, dtype=float)
+        # Only evaluate the likelihood where the prior is finite
+        mask = np.isfinite(samples.log_prior, dtype=bool)
+        # Valid samples
+        x = samples.x[mask, :]
+        logl[mask] = -np.sum(x**2, axis=1)  # Example likelihood
+        return logl
+
+
+Checking the flow distribution
+------------------------------
+
+It can be useful to inspect the flow-based proposal distribution before sampling
+from the posterior. You can do this by drawing samples from the flow after fitting
+and comparing them to the initial samples:
+
+
+.. code-block:: python
+
+    from aspire import Aspire, Samples
+    from aspire.plot import plot_comparison
+
+    # Define the initial samples
+    initial_samples = Samples(...)
+
+    # Define the Aspire instance
+    aspire = Aspire(
+        log_likelihood=log_likelihood,
+        log_prior=log_prior,
+        ...
+    )
+
+    # Fit the flow to the initial samples
+    fit_history = aspire.fit(initial_samples)
+
+    # Draw samples from the flow
+    flow_samples = aspire.sample_flow(10_000)
+
+    # Plot a comparison between initial samples and flow samples
+    fig = plot_comparison(
+        initial_samples,
+        flow_samples,
+        per_samples_kwargs=[
+            dict(include_weights=False, color="C0"),
+            dict(include_weights=False, color="C1"),
+        ],
+        labels=["Initial samples", "Flow samples"],
+    )
+    # Save or show the figure
+    fig.savefig("flow_comparison.png")

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/docs/requirements.txt

@@ -1,2 +1,3 @@
 sphinx>=7.2
 sphinx-book-theme>=1.1
+sphinx-autoapi>=3.2

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/docs/user_guide.rst

@@ -44,8 +44,8 @@ switch namespaces or merge multiple runs with
 Flows and transforms
 --------------------
 
-Aspire can work with any proposal that implements ``sample_and_log_prob`` and
-``log_prob``; normalising flows remain the default. Flows are defined via
+Aspire can work with any flow that implements ``sample_and_log_prob`` and
+``log_prob``. Flows are defined via
 :class:`aspire.flows.base.Flow` and instantiated by
 :meth:`aspire.Aspire.init_flow`. By default Aspire uses the ``zuko``
 implementation of Masked Autoregressive Flows on top of PyTorch. The flow is
@@ -61,6 +61,55 @@ estimator (requires the `zuko` backend).
 External flow implementations can be plugged in via the
 ``aspire.flows`` entry point group. See :ref:`custom_flows` for details.
 
+Transform mechanics
+~~~~~~~~~~~~~~~~~~~
+
+Aspire keeps a clear separation between your native parameters and the space
+where flows or kernels operate:
+
+* :class:`aspire.transforms.FlowTransform` is attached to every flow created by
+  :meth:`aspire.Aspire.init_flow`. By default, it maps bounded parameters to the real line (``probit`` or
+  ``logit``), and recentres / rescales dimensions with an affine
+  transform learned from the training samples. Log-Jacobian terms are tracked so
+  calls to ``log_prob`` or ``sample_and_log_prob`` remain properly normalised.
+  ``bounded_to_unbounded`` and ``affine_transform`` can be specified when creating
+  the Aspire instance to control this behaviour.
+* The same components are exposed via :class:`aspire.transforms.CompositeTransform`
+  if you want to opt out of the bounded-to-unbounded step or the affine
+  whitening when building custom transports.
+
+Preconditioning inside samplers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+SMC and MCMC samplers also work in a transformed space. They fit the chosen
+``preconditioning`` transform to the initial particles, perform moves there, and
+then call ``inverse(...)`` (including the log-Jacobian) whenever the likelihood
+or prior is evaluated. Configure it via
+:meth:`aspire.Aspire.sample_posterior`:
+
+* ``"default"`` / ``"standard"`` uses :class:`aspire.transforms.CompositeTransform`
+  with bounded-to-unbounded and affine scaling turned off by default; periodic
+  wrapping still applies. To whiten dimensions or map bounds to the real line,
+  pass ``preconditioning_kwargs={"affine_transform": True, "bounded_to_unbounded": True}``.
+* ``"flow"`` fits a lightweight :class:`aspire.transforms.FlowPreconditioningTransform`
+  to the current particles and treats it as a transport map during SMC/MCMC
+  updates. This reuses the same bounded / periodic handling while providing a
+  richer geometry for the kernels.
+* ``None`` leaves the sampler in the original parameterisation with an identity
+  transform. The importance sampler defaults to this; other samplers default to
+  ``"standard"`` so periodic parameters are at least kept consistent with their
+  bounds.
+
+
+.. note::
+
+    By default, the preconditioning transform does not include bounded-to-unbounded
+    steps. This means your log-prior and log-likelihood must handle points that
+    lie outside the specified bounds (e.g. by returning ``-inf``). If you want
+    the sampler to automatically map bounded parameters to an unconstrained
+    space, enable the ``bounded_to_unbounded`` option in
+    ``preconditioning_kwargs``.
+
 Sampling strategies
 -------------------
 
@@ -101,11 +150,6 @@ Sequential Monte Carlo
     Replaces the internal MCMC move with the ``emcee`` ensemble sampler,
     providing a gradient-free option that still benefits from SMC tempering.
 
-You can plug in custom preconditioning by setting ``preconditioning`` to
-``"standard"`` (affine normalisation based on current samples), ``"flow"``
-(use the fitted flow as a transport map), or ``None`` to disable additional
-transforms.
-
 History, diagnostics, and persistence
 -------------------------------------
 

aspire_inference-0.1.0a9/examples/blackjax_smc_example.py

@@ -0,0 +1,158 @@
+"""Example using sequential posterior inference with SMC.
+
+This example uses JAX for computations and BlackJAX for the MCMC sampling
+in SMC step.
+"""
+
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+
+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
+key = jax.random.key(42)
+
+# Output directory
+outdir = Path("outdir") / "blackjax_smc_example"
+outdir.mkdir(parents=True, exist_ok=True)
+
+# Configure logger to show INFO level messages
+configure_logger()
+
+# Number of dimensions
+dims = 4
+# Parameter names
+parameters = [f"x{i}" for i in range(dims)]
+# Prior bounds
+prior_bounds = {param: (-5, 5) for param in parameters}
+
+# Means and covariances of the two Gaussian components
+mu1 = 2 * jnp.ones(dims)
+mu2 = -2 * jnp.ones(dims)
+cov1 = 0.5 * jnp.eye(dims)
+cov2 = jnp.eye(dims)
+
+
+def log_likelihood(samples):
+    """Log-likelihood of a mixture of two Gaussians"""
+    x = samples.x
+    comp1 = (
+        -0.5 * ((x - mu1) @ jnp.linalg.inv(cov1) * (x - mu1)).sum(axis=-1)
+        - 0.5 * dims * jnp.log(2 * jnp.pi)
+        - 0.5 * jnp.linalg.slogdet(cov1)[1]
+    )
+    comp2 = (
+        -0.5 * ((x - mu2) @ jnp.linalg.inv(cov2) * (x - mu2)).sum(axis=-1)
+        - 0.5 * dims * jnp.log(2 * jnp.pi)
+        - 0.5 * jnp.linalg.slogdet(cov2)[1]
+    )
+    return jnp.logaddexp(comp1, comp2)  # Log-sum-exp for numerical stability
+
+
+def log_prior(samples):
+    """Uniform prior between -5 and 5 in each dimension"""
+    x = samples.x
+    in_bounds = jnp.all((x >= -5) & (x <= 5), axis=-1)
+    logp = jnp.where(in_bounds, -dims * jnp.log(10), -jnp.inf)
+    return logp
+
+
+# Generate prior samples for comparison, these are not used in SMC
+key, prior_key = jax.random.split(key)
+prior_samples = Samples(
+    jax.random.uniform(prior_key, shape=(5000, dims), minval=-5, maxval=5),
+    parameters=parameters,
+)
+
+# True posterior samples for comparison
+key, post_key0, post_key1 = jax.random.split(key, 3)
+true_posterior_samples = Samples(
+    jnp.concatenate(
+        [
+            jax.random.multivariate_normal(
+                post_key0, mu1, cov1, shape=(2500,)
+            ),
+            jax.random.multivariate_normal(
+                post_key1, mu2, cov2, shape=(2500,)
+            ),
+        ],
+        axis=0,
+    ),
+    parameters=parameters,
+)
+
+# We draw initial samples from two Gaussians centered away from the true modes
+# to demonstrate the ability of SMC to explore the posterior
+key, offset_key1, offset_key2 = jax.random.split(key, 3)
+offset_1 = jax.random.uniform(offset_key1, shape=(dims,), minval=-3, maxval=3)
+offset_2 = jax.random.uniform(offset_key2, shape=(dims,), minval=-3, maxval=3)
+key, init_key1, init_key2 = jax.random.split(key, 3)
+initial_samples = jnp.concatenate(
+    [
+        jax.random.normal(init_key1, shape=(2500, dims)) + (mu1 - offset_1),
+        jax.random.normal(init_key2, shape=(2500, dims)) + (mu2 - offset_2),
+    ],
+    axis=0,
+)
+initial_samples = Samples(initial_samples, parameters=parameters)
+
+# Initialize Aspire with the log-likelihood and log-prior
+key, aspire_key = jax.random.split(key)
+aspire = Aspire(
+    log_likelihood=log_likelihood,
+    log_prior=log_prior,
+    dims=dims,
+    flow_backend="flowjax",  # Use Flowjax as the normalizing flow backend
+    prior_bounds=prior_bounds,  # Specify prior bounds
+    parameters=parameters,
+    key=aspire_key,
+)
+
+# Fit the normalizing flow to the initial samples
+fit_history = aspire.fit(initial_samples, max_epochs=30)
+
+# Plot loss
+fit_history.plot_loss().savefig(outdir / "loss.png")
+
+# Sample from the posterior using SMC
+# We use BlackJAX's NUTS as the MCMC kernel within SMC
+# We enable the bounded to unbounded transform in the preconditioning to avoid
+# issues with NUTS on bounded spaces
+samples, history = aspire.sample_posterior(
+    sampler="blackjax_smc",  # use the BlackJAX SMC sampler
+    n_samples=500,  # Number of particles in SMC
+    n_final_samples=5000,  # Number of samples to draw from the final distribution
+    adaptive=True,
+    target_efficiency=0.8,
+    sampler_kwargs=dict(  # Keyword arguments for the specific sampler
+        algorithm="nuts",  # Use NUTS within SMC
+        step_size=0.1,  # Step size for NUTS, this will need tuning
+        n_steps=20,  # Number of leapfrog steps for NUTS
+    ),
+    preconditioning_kwargs=dict(
+        affine_transform=True,  # Use affine transform preconditioning
+        bounded_to_unbounded=True,  # Transform bounded parameters to unbounded space
+    ),
+    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,
+    true_posterior_samples,
+    samples,
+    labels=["Initial Samples", "True Posterior Samples", "SMC Samples"],
+    per_samples_kwargs=[
+        {"color": "grey"},
+        {"color": "k"},
+        {"color": "C1"},
+    ],
+).savefig(outdir / "posterior.png")

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/src/aspire/__init__.py

@@ -6,6 +6,7 @@ import logging
 from importlib.metadata import PackageNotFoundError, version
 
 from .aspire import Aspire
+from .samples import Samples
 
 try:
     __version__ = version("aspire")
@@ -16,4 +17,5 @@ logging.getLogger(__name__).addHandler(logging.NullHandler())
 
 __all__ = [
     "Aspire",
+    "Samples",
 ]

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/src/aspire/samplers/smc/blackjax.py

@@ -220,7 +220,9 @@ class BlackJAXSMC(SMCSampler):
 
         elif algorithm == "nuts":
             # Initialize step size and mass matrix if not provided
-            inverse_mass_matrix = self.sampler_kwargs["inverse_mass_matrix"]
+            inverse_mass_matrix = self.sampler_kwargs.get(
+                "inverse_mass_matrix"
+            )
             if inverse_mass_matrix is None:
                 inverse_mass_matrix = jax.numpy.eye(self.dims)
 
@@ -258,8 +260,13 @@ class BlackJAXSMC(SMCSampler):
             z_final = final_states.position
 
             # Calculate acceptance rates
-            acceptance_rates = jax.numpy.mean(all_infos.is_accepted, axis=1)
-            mean_acceptance = jax.numpy.mean(acceptance_rates)
+            try:
+                acceptance_rates = jax.numpy.mean(
+                    all_infos.is_accepted, axis=1
+                )
+                mean_acceptance = jax.numpy.mean(acceptance_rates)
+            except AttributeError:
+                mean_acceptance = np.nan
 
         elif algorithm == "hmc":
             # Initialize HMC sampler
@@ -294,8 +301,13 @@ class BlackJAXSMC(SMCSampler):
 
             final_states, all_infos = jax.vmap(init_and_sample)(keys, z_jax)
             z_final = final_states.position
-            acceptance_rates = jax.numpy.mean(all_infos.is_accepted, axis=1)
-            mean_acceptance = jax.numpy.mean(acceptance_rates)
+            try:
+                acceptance_rates = jax.numpy.mean(
+                    all_infos.is_accepted, axis=1
+                )
+                mean_acceptance = jax.numpy.mean(acceptance_rates)
+            except AttributeError:
+                mean_acceptance = np.nan
 
         else:
             raise ValueError(f"Unsupported algorithm: {algorithm}")

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/src/aspire/transforms.py

@@ -332,6 +332,7 @@ class CompositeTransform(BaseTransform):
             prior_bounds=self.prior_bounds,
             bounded_to_unbounded=self.bounded_to_unbounded,
             bounded_transform=self.bounded_transform,
+            affine_transform=self.affine_transform,
             device=self.device,
             xp=xp or self.xp,
             eps=self.eps,

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/tests/integration_tests/conftest.py

@@ -116,7 +116,8 @@ def log_prior(dims, xp):
         "smc",
         "emcee",
         "minipcn",
-        "blackjax_smc",
+        "blackjax_smc_rwmh",
+        "blackjax_smc_nuts",
     ]
 )
 def sampler_config(request):
@@ -126,7 +127,7 @@ def sampler_config(request):
         return SamplerConfig(
             sampler="emcee",
             sampler_kwargs={
-                "nsteps": 500,
+                "nsteps": 100,
             },
         )
     elif request.param == "minipcn":
@@ -157,7 +158,7 @@ def sampler_config(request):
                 },
             },
         )
-    elif request.param == "blackjax_smc":
+    elif request.param == "blackjax_smc_rwmh":
         return SamplerConfig(
             sampler="blackjax_smc",
             sampler_kwargs={
@@ -165,7 +166,19 @@ def sampler_config(request):
                 "sampler_kwargs": {
                     "algorithm": "rwmh",
                     "sigma": 0.1,
-                    "n_steps": 500,
+                    "n_steps": 10,
+                },
+            },
+        )
+    elif request.param == "blackjax_smc_nuts":
+        return SamplerConfig(
+            sampler="blackjax_smc",
+            sampler_kwargs={
+                "adaptive": True,
+                "sampler_kwargs": {
+                    "algorithm": "nuts",
+                    "step_size": 0.1,
+                    "n_steps": 10,
                 },
             },
         )

{aspire_inference-0.1.0a8 → aspire_inference-0.1.0a9}/tests/integration_tests/test_integration.py

@@ -21,7 +21,7 @@ def test_integration_zuko(
     dtype,
     tmp_path,
 ):
-    if sampler_config.sampler == "blackjax_smc":
+    if "blackjax_smc" in sampler_config.sampler:
         pytest.xfail(reason="BlackJAX requires JAX arrays.")
 
     aspire = Aspire(
@@ -63,7 +63,7 @@ def test_integration_flowjax(
 ):
     import jax
 
-    if sampler_config.sampler == "blackjax_smc":
+    if "blackjax_smc" in sampler_config.sampler:
         if samples_backend != "jax":
             pytest.xfail(reason="BlackJAX requires JAX arrays.")
         if dtype == "float32":

aspire_inference-0.1.0a8/README.md

@@ -1,27 +0,0 @@
-# aspire: Accelerated Sequential Posterior Inference via REuse
-
-aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
-
-## Installation
-
-aspire can be installed from PyPI using `pip`
-
-```
-pip install aspire-inference
-```
-
-**Important:** the name of `aspire` on PyPI is `aspire-inference` but once installed
-the package can be imported and used as `aspire`.
-
-## Documentation
-
-See the [documentation on ReadTheDocs][docs].
-
-## Citation
-
-If you use `aspire` in your work please cite the [DOI][DOI] and [paper][paper].
-
-
-[docs]: https://aspire.readthedocs.io/
-[DOI]: https://doi.org/10.5281/zenodo.15658747
-[paper]: https://arxiv.org/abs/2511.04218

aspire_inference-0.1.0a8/docs/api.rst

@@ -1,28 +0,0 @@
-API Reference
-=============
-
-This section documents the main public classes.
-
-Aspire interface
-----------------
-
-.. autoclass:: aspire.Aspire
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-Samples utilities
------------------
-
-.. automodule:: aspire.samples
-   :members: Samples, SMCSamples
-   :undoc-members:
-   :show-inheritance:
-
-History objects
----------------
-
-.. automodule:: aspire.history
-   :members: History, FlowHistory, SMCHistory
-   :undoc-members:
-   :show-inheritance: