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

aspire/flows/__init__.py

@@ -1,5 +1,31 @@
-def get_flow_wrapper(backend: str = "zuko", flow_matching: bool = False):
-    """Get the wrapper for the flow implementation."""
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def get_flow_wrapper(
+    backend: str = "zuko", flow_matching: bool = False
+) -> tuple[type, Any]:
+    """Get the wrapper for the flow implementation.
+
+    Parameters
+    ----------
+    backend : str
+        The backend to use. Options are "zuko" (PyTorch), "flowjax" (JAX), or
+        any other registered flow class via entry points. Default is "zuko".
+    flow_matching : bool, optional
+        Whether to use flow matching variant of the flow. Default is False.
+
+    Returns
+    -------
+    FlowClass : type
+        The flow class corresponding to the specified backend.
+    xp : module
+        The array API module corresponding to the specified backend.
+    """
+    from importlib.metadata import entry_points
+
     if backend == "zuko":
         import array_api_compat.torch as torch_api
 
@@ -20,11 +46,12 @@ def get_flow_wrapper(backend: str = "zuko", flow_matching: bool = False):
             )
         return FlowJax, jnp
     else:
-        from importlib.metadata import entry_points
-
+        if flow_matching:
+            logger.warning(
+                "Flow matching option is ignored for external backends."
+            )
         eps = {
-            ep.name.lower(): ep
-            for ep in entry_points().get("aspire.flows", [])
+            ep.name.lower(): ep for ep in entry_points(group="aspire.flows")
         }
         if backend in eps:
             FlowClass = eps[backend].load()
@@ -35,6 +62,7 @@ def get_flow_wrapper(backend: str = "zuko", flow_matching: bool = False):
                 )
             return FlowClass, xp
         else:
+            known_backends = ["zuko", "flowjax"] + list(eps.keys())
             raise ValueError(
-                f"Unknown flow class: {backend}. Available classes: {list(eps.keys())}"
+                f"Unknown backend '{backend}'. Available backends: {known_backends}"
             )

aspire/samples.py

@@ -14,6 +14,7 @@ from array_api_compat import (
 )
 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,
@@ -161,7 +162,24 @@ class BaseSamples:
 
         return pd.DataFrame(self.to_dict(flat=flat))
 
-    def plot_corner(self, parameters: list[str] | None = None, **kwargs):
+    def plot_corner(
+        self,
+        parameters: list[str] | None = None,
+        fig: Figure | None = None,
+        **kwargs,
+    ):
+        """Plot a corner plot of the samples.
+
+        Parameters
+        ----------
+        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
+            are deep-copied before use.
+        """
         import corner
 
         kwargs = copy.deepcopy(kwargs)
@@ -173,7 +191,7 @@ class BaseSamples:
             x = self.x[:, indices] if self.x.ndim > 1 else self.x[indices]
         else:
             x = self.x
-        fig = corner.corner(to_numpy(x), **kwargs)
+        fig = corner.corner(to_numpy(x), fig=fig, **kwargs)
         return fig
 
     def __str__(self):

aspire/transforms.py

@@ -684,7 +684,7 @@ class FlowPreconditioningTransform(BaseTransform):
             self.flow_kwargs.setdefault("dtype", dtype)
         self.fit_kwargs = dict(fit_kwargs or {})
 
-        FlowClass = get_flow_wrapper(
+        FlowClass, xp = get_flow_wrapper(
             backend=flow_backend, flow_matching=flow_matching
         )
         transform = CompositeTransform(
@@ -695,7 +695,7 @@ class FlowPreconditioningTransform(BaseTransform):
             bounded_transform=bounded_transform,
             affine_transform=affine_transform,
             device=device,
-            xp=FlowClass.xp,
+            xp=xp,
             eps=eps,
             dtype=dtype,
         )

{aspire_inference-0.1.0a6.dist-info → aspire_inference-0.1.0a8.dist-info}/METADATA

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a6
+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
@@ -38,7 +39,7 @@ 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 code.
+aspire is a framework for reusing existing posterior samples to obtain new results at a reduced cost.
 
 ## Installation
 
@@ -50,3 +51,16 @@ 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.0a6.dist-info → aspire_inference-0.1.0a8.dist-info}/RECORD

@@ -2,10 +2,10 @@ aspire/__init__.py,sha256=45R0xWaLg0aJEPK5zoTK0aIek0KOwpHwQWS1jLCDhIE,365
 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=lUn3cQdnN5gSHREuamZZteecdDIuasaWfcDXdVAvfZA,18662
-aspire/transforms.py,sha256=XMbf5MxK49elQeKDsmFraHN-0JeO1AciljdTk7k2ujk,24928
+aspire/samples.py,sha256=nVQULOr19lQVyGhitI8EgDdCDGx9sExPomQBJrV4rTc,19237
+aspire/transforms.py,sha256=W2JL045zkczvdBVXXeSn2c8aKyX8ASxFKgWU2cJHufk,24922
 aspire/utils.py,sha256=pj8O0chqfP6VS8bpW0wCw8W0P5JNQKvWRz1Rg9AYIhg,22525
-aspire/flows/__init__.py,sha256=3gGXF4HziMlZSmcEdJ_uHtrP-QEC6RXvylm4vtM-Xnk,1306
+aspire/flows/__init__.py,sha256=GUZToPVNJoTwULpbeW10UijfQukNrILoAQ_ubeq7G3w,2110
 aspire/flows/base.py,sha256=5UWKAiXDXLJ6Sg6a380ajLrGFaZSQyOnFEihQiiA4ko,2237
 aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
 aspire/flows/jax/flows.py,sha256=1HnVgQ1GUXNcvxiZqEV19H2QI9Th5bWX_QbNfGaUhuA,6625
@@ -21,8 +21,8 @@ aspire/samplers/smc/base.py,sha256=66f_ORUvcKRqMIW35qjhUc-c0PFuY87lJa91MpSaTZI,1
 aspire/samplers/smc/blackjax.py,sha256=4L4kgRKlaWl-knTWXXzdJTh-zZBh5BTpy5GaLDzT8Sc,11803
 aspire/samplers/smc/emcee.py,sha256=Wm0vvAlCcRhJMBt7_fU2ZnjDb8SN8jgUOTXLzNstRpA,2516
 aspire/samplers/smc/minipcn.py,sha256=ju1gcgyKHjodLEACPdL3eXA9ai8ZJ9_LwitD_Gmf1Rc,2765
-aspire_inference-0.1.0a6.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
-aspire_inference-0.1.0a6.dist-info/METADATA,sha256=--Q4vjeyHU7-TcgV3HicQ9wxMuO_Vi1pba5wDUK2oD0,1617
-aspire_inference-0.1.0a6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aspire_inference-0.1.0a6.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
-aspire_inference-0.1.0a6.dist-info/RECORD,,
+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,,