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

aspire/aspire.py

@@ -20,6 +20,7 @@ from .transforms import (
 )
 from .utils import (
     AspireFile,
+    function_id,
     load_from_h5_file,
     recursively_save_to_h5_file,
     resolve_xp,
@@ -353,6 +354,7 @@ class Aspire:
             xp=self.xp,
             dtype=self.dtype,
             preconditioning_transform=transform,
+            parameters=self.parameters,
             **kwargs,
         )
         return sampler
@@ -687,8 +689,8 @@ class Aspire:
             method of the sampler.
         """
         config = {
-            "log_likelihood": self.log_likelihood.__name__,
-            "log_prior": self.log_prior.__name__,
+            "log_likelihood": function_id(self.log_likelihood),
+            "log_prior": function_id(self.log_prior),
             "dims": self.dims,
             "parameters": self.parameters,
             "periodic_parameters": self.periodic_parameters,

aspire/samplers/mcmc.py

@@ -13,14 +13,13 @@ class MCMCSampler(Sampler):
         n_samples_drawn = 0
         samples = None
         while n_samples_drawn < n_samples:
-            n_to_draw = n_samples - n_samples_drawn
-            x, log_q = self.prior_flow.sample_and_log_prob(n_to_draw)
+            x, log_q = self.prior_flow.sample_and_log_prob(n_samples)
             new_samples = Samples(x, xp=self.xp, log_q=log_q, dtype=self.dtype)
             new_samples.log_prior = new_samples.array_to_namespace(
                 self.log_prior(new_samples)
             )
             valid = self.xp.isfinite(new_samples.log_prior)
-            n_valid = self.xp.sum(valid)
+            n_valid = int(self.xp.sum(valid))
             if n_valid > 0:
                 if samples is None:
                     samples = new_samples[valid]

aspire/samplers/smc/base.py

@@ -95,6 +95,7 @@ class SMCSampler(MCMCSampler):
         beta: float,
         beta_step: float,
         min_step: float,
+        beta_tolerance: float = 1e-6,
     ) -> tuple[float, float]:
         """Determine the next beta value.
 
@@ -108,6 +109,8 @@ class SMCSampler(MCMCSampler):
             The fixed beta step size if not adaptive.
         min_step : float
             The minimum beta step size.
+        beta_tolerance : float
+            Tolerance when checking for beta convergence.
 
         Returns
         -------
@@ -124,14 +127,13 @@ class SMCSampler(MCMCSampler):
             beta_prev = beta
             beta_min = beta_prev
             beta_max = 1.0
-            tol = 1e-5
             eff_beta_max = effective_sample_size(
                 samples.log_weights(beta_max)
             ) / len(samples)
             if eff_beta_max >= self.current_target_efficiency(beta_prev):
                 beta_min = 1.0
             target_eff = self.current_target_efficiency(beta_prev)
-            while beta_max - beta_min > tol:
+            while beta_max - beta_min > beta_tolerance:
                 beta_try = 0.5 * (beta_max + beta_min)
                 eff = effective_sample_size(
                     samples.log_weights(beta_try)
@@ -163,6 +165,7 @@ class SMCSampler(MCMCSampler):
         checkpoint_every: int | None = None,
         checkpoint_file_path: str | None = None,
         resume_from: str | bytes | dict | None = None,
+        beta_tolerance: float = 1e-6,
     ) -> SMCSamples:
         resumed = resume_from is not None
         if resumed:
@@ -249,6 +252,7 @@ class SMCSampler(MCMCSampler):
                     beta,
                     beta_step,
                     min_step,
+                    beta_tolerance=beta_tolerance,
                 )
                 self.history.eff_target.append(
                     self.current_target_efficiency(beta)

aspire/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import inspect
 import logging
 import pickle
@@ -911,3 +912,23 @@ def track_calls(wrapped=None):
         return wrapped_func(*args, **kwargs)
 
     return wrapper(wrapped) if wrapped else wrapper
+
+
+def function_id(fn: Any) -> str:
+    """Get a unique identifier for a function.
+
+    Parameters
+    ----------
+    fn : Any
+        The function to get the identifier for.
+
+    Returns
+    -------
+    str
+        The unique identifier for the function.
+    """
+    if isinstance(fn, functools.partial):
+        base = fn.func
+    else:
+        base = fn
+    return f"{base.__module__}:{getattr(base, '__qualname__', type(base).__name__)}"

{aspire_inference-0.1.0a11.dist-info → aspire_inference-0.1.0a13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a11
+Version: 0.1.0a13
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -19,7 +19,8 @@ Provides-Extra: scipy
 Requires-Dist: scipy; extra == "scipy"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
-Requires-Dist: jaxlib; extra == "jax"
+Requires-Dist: jaxlib<0.8.2; python_version < "3.12" and extra == "jax"
+Requires-Dist: jaxlib; python_version >= "3.12" and extra == "jax"
 Requires-Dist: flowjax; extra == "jax"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
@@ -32,6 +33,7 @@ Provides-Extra: emcee
 Requires-Dist: emcee; extra == "emcee"
 Provides-Extra: blackjax
 Requires-Dist: blackjax; extra == "blackjax"
+Requires-Dist: fastprogress<1.1.0; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"

{aspire_inference-0.1.0a11.dist-info → aspire_inference-0.1.0a13.dist-info}/RECORD

@@ -1,10 +1,10 @@
 aspire/__init__.py,sha256=B2IETvlpB0oBh57prRYLdi8jB5yFGw8qVviGdf1NcnE,409
-aspire/aspire.py,sha256=lr0bD5GDWdlAGfODGzj4BoELKUF5HYAMb8yYGvRR_y0,30860
+aspire/aspire.py,sha256=EelzGme70nrR_iEeT_HTB4feloYXzW2WADZUwOcREpo,30925
 aspire/history.py,sha256=l_j-riZKbTWK7Wz9zvvD_mTk9psNCKItiveYhr_pYv8,4313
 aspire/plot.py,sha256=oXwUDOb_953_ADm2KLk41JIfpE3JeiiQiSYKvUVwLqw,1423
 aspire/samples.py,sha256=v7y8DkirUCHOJbCE-o9y2K7xzU2HicIo_O0CdFhLgXE,19478
 aspire/transforms.py,sha256=CHrfPQHEyHQ9I0WWiAgUplWwxYypZsD7uCYIHUbSFtY,24974
-aspire/utils.py,sha256=87avRkTce9QvELpIcqlKauSoUZSYi1fqe1asC97TzqA,26947
+aspire/utils.py,sha256=sIONKn3gT7i3hVdlK9bRWy_I79rdk0QPkXTA4O1FlCI,27405
 aspire/flows/__init__.py,sha256=GUZToPVNJoTwULpbeW10UijfQukNrILoAQ_ubeq7G3w,2110
 aspire/flows/base.py,sha256=5UWKAiXDXLJ6Sg6a380ajLrGFaZSQyOnFEihQiiA4ko,2237
 aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
@@ -15,14 +15,14 @@ aspire/flows/torch/flows.py,sha256=QcQOcFZEsLWHPwbQUFGOFdfEslyc59Vf_UEsS0xAGPo,1
 aspire/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aspire/samplers/base.py,sha256=ygqrvqSedWSb0cz8DQ_MHokOOxi6aBRdHxf_qoEPwUE,8243
 aspire/samplers/importance.py,sha256=opn_jY-V8snUz0JztLBtnaTT3WfrZ5OSpHBV5WAuM3M,676
-aspire/samplers/mcmc.py,sha256=ihHgzqvSoy1oxdFBjyqNUbCuRX7CqWjlshCUZcgEL5E,5151
+aspire/samplers/mcmc.py,sha256=vLLuXUJfNH9QpkRJTLdJB5RylKRy-3cHNPe_eyi1WQE,5104
 aspire/samplers/smc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/smc/base.py,sha256=40A9yVuKS1F8cPzbfVQ9rNk3y07mnkfbuyRDIh_fy5A,14122
+aspire/samplers/smc/base.py,sha256=-WTUWLsODsbOQcXjnVQW1-k75xbQ6H31vpCBP8EwQEE,14326
 aspire/samplers/smc/blackjax.py,sha256=2riWDSRmpL5lGmnhNtdieiRs0oYC6XZA2X-nVlQaqpE,12490
 aspire/samplers/smc/emcee.py,sha256=4CI9GvH69FCoLiFBbKKYwYocYyiM95IijC5EvrcAmUo,2891
 aspire/samplers/smc/minipcn.py,sha256=IJ5466VvARd4qZCWXXl-l3BPaKW1AgcwmbP3ISL2bto,3368
-aspire_inference-0.1.0a11.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
-aspire_inference-0.1.0a11.dist-info/METADATA,sha256=lLd2d5HR-t942wKLyYbdJ1DL9CKl7tkpBul_vX8DU4M,3869
-aspire_inference-0.1.0a11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aspire_inference-0.1.0a11.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
-aspire_inference-0.1.0a11.dist-info/RECORD,,
+aspire_inference-0.1.0a13.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
+aspire_inference-0.1.0a13.dist-info/METADATA,sha256=l0CpbJt7o8zFQGRXM6lKxDQXsndXk2eaAZBPejfwkUQ,4025
+aspire_inference-0.1.0a13.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+aspire_inference-0.1.0a13.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
+aspire_inference-0.1.0a13.dist-info/RECORD,,

{aspire_inference-0.1.0a11.dist-info → aspire_inference-0.1.0a13.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any