aspire-inference 0.1.0a11__tar.gz → 0.1.0a13__tar.gz

{aspire_inference-0.1.0a11 → aspire_inference-0.1.0a13}/PKG-INFO

@@ -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 → aspire_inference-0.1.0a13}/aspire_inference.egg-info/PKG-INFO

@@ -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 → aspire_inference-0.1.0a13}/aspire_inference.egg-info/requires.txt

@@ -6,15 +6,21 @@ h5py
 
 [blackjax]
 blackjax
+fastprogress<1.1.0
 
 [emcee]
 emcee
 
 [jax]
 jax
-jaxlib
 flowjax
 
+[jax:python_version < "3.12"]
+jaxlib<0.8.2
+
+[jax:python_version >= "3.12"]
+jaxlib
+
 [minipcn]
 minipcn[array-api]>=0.2.0a3
 orng

{aspire_inference-0.1.0a11 → aspire_inference-0.1.0a13}/pyproject.toml

@@ -30,8 +30,9 @@ scipy = [
 ]
 jax = [
     "jax",
-    "jaxlib",
-    "flowjax"
+    "jaxlib<0.8.2; python_version < '3.12'",
+    "jaxlib; python_version >= '3.12'",
+    "flowjax",
 ]
 torch = [
     "torch",
@@ -47,6 +48,7 @@ emcee = [
 ]
 blackjax = [
     "blackjax",
+    "fastprogress<1.1.0",
 ]
 test = [
     "pytest",

{aspire_inference-0.1.0a11 → aspire_inference-0.1.0a13}/src/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_inference-0.1.0a11 → aspire_inference-0.1.0a13}/src/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_inference-0.1.0a11 → aspire_inference-0.1.0a13}/src/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_inference-0.1.0a11 → aspire_inference-0.1.0a13}/src/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 → aspire_inference-0.1.0a13}/tests/test_utils.py

@@ -1,3 +1,4 @@
+import functools
 import pickle
 
 import array_api_compat.numpy as np_xp
@@ -6,7 +7,7 @@ import h5py
 import jax.numpy as jnp
 import pytest
 
-from aspire.utils import convert_dtype, dump_state, resolve_dtype
+from aspire.utils import convert_dtype, dump_state, function_id, resolve_dtype
 
 
 def _dtype_name(dtype):
@@ -85,3 +86,18 @@ def test_dump_state_round_trip(tmp_path):
         stored = fp["checkpoints"]["state"][...]
     restored = pickle.loads(stored.tobytes())
     assert restored == state
+
+
+# Define a simple function for testing
+def _foo(x):
+    return x
+
+
+@pytest.mark.parametrize(
+    "fn", [lambda x: x, functools.partial(lambda x, y: x, y=1), _foo]
+)
+def test_function_id(fn):
+    fn_id = function_id(fn)
+    assert isinstance(fn_id, str)
+    # Calling again should give the same result
+    assert fn_id == function_id(fn)