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

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/samplers/base.py

@@ -4,7 +4,7 @@ from typing import Any, Callable
 from ..flows.base import Flow
 from ..samples import Samples
 from ..transforms import IdentityTransform
-from ..utils import track_calls
+from ..utils import asarray, track_calls
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +56,11 @@ class Sampler:
 
     def fit_preconditioning_transform(self, x):
         """Fit the data transform to the data."""
-        x = self.preconditioning_transform.xp.asarray(x)
+        x = asarray(
+            x,
+            xp=self.preconditioning_transform.xp,
+            dtype=self.preconditioning_transform.dtype,
+        )
         return self.preconditioning_transform.fit(x)
 
     @track_calls

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/samplers/smc/minipcn.py

@@ -3,7 +3,11 @@ from functools import partial
 import numpy as np
 
 from ...samples import SMCSamples
-from ...utils import to_numpy, track_calls
+from ...utils import (
+    asarray,
+    determine_backend_name,
+    track_calls,
+)
 from .base import NumpySMCSampler
 
 
@@ -13,7 +17,7 @@ class MiniPCNSMC(NumpySMCSampler):
     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(
@@ -29,11 +33,14 @@ class MiniPCNSMC(NumpySMCSampler):
         sampler_kwargs: dict | None = None,
         rng: np.random.Generator | 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,
@@ -58,9 +65,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

@@ -425,19 +425,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/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/utils.py

@@ -12,7 +12,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 +34,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 +251,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 +261,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:

aspire_inference-0.1.0a10.dist-info/METADATA

@@ -0,0 +1,130 @@
+Metadata-Version: 2.4
+Name: aspire-inference
+Version: 0.1.0a10
+Summary: Accelerate Sequential Posterior Inference via REuse
+Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/mj-will/aspire
+Project-URL: Documentation, https://aspire.readthedocs.io/
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: array-api-compat
+Requires-Dist: wrapt
+Requires-Dist: h5py
+Provides-Extra: scipy
+Requires-Dist: scipy; extra == "scipy"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: jaxlib; extra == "jax"
+Requires-Dist: flowjax; extra == "jax"
+Provides-Extra: torch
+Requires-Dist: torch; extra == "torch"
+Requires-Dist: zuko; extra == "torch"
+Requires-Dist: tqdm; extra == "torch"
+Provides-Extra: minipcn
+Requires-Dist: minipcn[array-api]>=0.2.0a3; extra == "minipcn"
+Requires-Dist: orng; extra == "minipcn"
+Provides-Extra: emcee
+Requires-Dist: emcee; extra == "emcee"
+Provides-Extra: blackjax
+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
+
+[![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`. By default, you need to install
+one of the backends for the normalizing flows, either `torch` or `jax`.
+We also recommend installing `minipcn` if using the `smc` sampler:
+
+
+**Torch**
+
+We recommend installing `torch` manually to ensure correct CPU/CUDA versions are
+ installed. See the [PyTorch installation instructions](https://pytorch.org/)
+ for more details.
+
+```
+pip install aspire-inference[torch,minipcn]
+```
+
+**Jax**:
+
+We recommend install `jax` manually to ensure the correct GPU/CUDA versions
+are installed. See the [jax documentation for details](https://docs.jax.dev/en/latest/installation.html)
+
+```
+pip install aspire-inference[jax,minipcn]
+```
+
+**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.dist-info → aspire_inference-0.1.0a10.dist-info}/RECORD

@@ -1,10 +1,10 @@
-aspire/__init__.py,sha256=45R0xWaLg0aJEPK5zoTK0aIek0KOwpHwQWS1jLCDhIE,365
+aspire/__init__.py,sha256=B2IETvlpB0oBh57prRYLdi8jB5yFGw8qVviGdf1NcnE,409
 aspire/aspire.py,sha256=M5o-QxLthE_5daa1trgUfWxPz-g4rmpEUKimKosw4lw,17400
 aspire/history.py,sha256=l_j-riZKbTWK7Wz9zvvD_mTk9psNCKItiveYhr_pYv8,4313
 aspire/plot.py,sha256=oXwUDOb_953_ADm2KLk41JIfpE3JeiiQiSYKvUVwLqw,1423
-aspire/samples.py,sha256=nVQULOr19lQVyGhitI8EgDdCDGx9sExPomQBJrV4rTc,19237
-aspire/transforms.py,sha256=W2JL045zkczvdBVXXeSn2c8aKyX8ASxFKgWU2cJHufk,24922
-aspire/utils.py,sha256=pj8O0chqfP6VS8bpW0wCw8W0P5JNQKvWRz1Rg9AYIhg,22525
+aspire/samples.py,sha256=z5x5hpWuVFH1hYhltmROAe8pbWxGD2UvHi3vcc132dg,19399
+aspire/transforms.py,sha256=CHrfPQHEyHQ9I0WWiAgUplWwxYypZsD7uCYIHUbSFtY,24974
+aspire/utils.py,sha256=eKGmchpuoNL15Xbu-AGoeZ00PcQEykEQiDZMnnRyV6A,24234
 aspire/flows/__init__.py,sha256=GUZToPVNJoTwULpbeW10UijfQukNrILoAQ_ubeq7G3w,2110
 aspire/flows/base.py,sha256=5UWKAiXDXLJ6Sg6a380ajLrGFaZSQyOnFEihQiiA4ko,2237
 aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
@@ -13,16 +13,16 @@ aspire/flows/jax/utils.py,sha256=5T6UrgpARG9VywC9qmTl45LjyZWuEdkW3XUladE6xJE,151
 aspire/flows/torch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aspire/flows/torch/flows.py,sha256=0_YkiMT49QolyQnEFsh28tfKLnURVF0Z6aTnaWLIUDI,11672
 aspire/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/base.py,sha256=8slvgOBnacUrHXCVDAqo-3IZ_LB7-dS8wdMP55MI43Y,2907
+aspire/samplers/base.py,sha256=VEHawyVA33jXHMo63p5hBHkp9k2qxU_bOxh5iaZSXew,3011
 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=4L4kgRKlaWl-knTWXXzdJTh-zZBh5BTpy5GaLDzT8Sc,11803
+aspire/samplers/smc/blackjax.py,sha256=IcTguAETiPmgFofmVW2GN40P5HBIxkmyd2VR8AU8f4k,12115
 aspire/samplers/smc/emcee.py,sha256=Wm0vvAlCcRhJMBt7_fU2ZnjDb8SN8jgUOTXLzNstRpA,2516
-aspire/samplers/smc/minipcn.py,sha256=ju1gcgyKHjodLEACPdL3eXA9ai8ZJ9_LwitD_Gmf1Rc,2765
-aspire_inference-0.1.0a8.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
-aspire_inference-0.1.0a8.dist-info/METADATA,sha256=Or73qfnx3KsPNSeaNTJ70t56GME2r9qACM03aSW40t8,1965
-aspire_inference-0.1.0a8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aspire_inference-0.1.0a8.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
-aspire_inference-0.1.0a8.dist-info/RECORD,,
+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_inference-0.1.0a8.dist-info/METADATA

@@ -1,66 +0,0 @@
-Metadata-Version: 2.4
-Name: aspire-inference
-Version: 0.1.0a8
-Summary: Accelerate Sequential Posterior Inference via REuse
-Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
-License: MIT
-Project-URL: Homepage, https://github.com/mj-will/aspire
-Project-URL: Documentation, https://aspire.readthedocs.io/
-Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: matplotlib
-Requires-Dist: numpy
-Requires-Dist: array-api-compat
-Requires-Dist: wrapt
-Requires-Dist: h5py
-Provides-Extra: scipy
-Requires-Dist: scipy; extra == "scipy"
-Provides-Extra: jax
-Requires-Dist: jax; extra == "jax"
-Requires-Dist: jaxlib; extra == "jax"
-Requires-Dist: flowjax; extra == "jax"
-Provides-Extra: torch
-Requires-Dist: torch; extra == "torch"
-Requires-Dist: zuko; extra == "torch"
-Requires-Dist: tqdm; extra == "torch"
-Provides-Extra: minipcn
-Requires-Dist: minipcn; extra == "minipcn"
-Provides-Extra: emcee
-Requires-Dist: emcee; extra == "emcee"
-Provides-Extra: blackjax
-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 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