bayinx 0.3.17__tar.gz → 0.3.19__tar.gz

{bayinx-0.3.17 → bayinx-0.3.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.17
+Version: 0.3.19
 Summary: Bayesian Inference with JAX
 License-File: LICENSE
 Requires-Python: >=3.12

{bayinx-0.3.17 → bayinx-0.3.19}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.3.17"
+version = "0.3.19"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,7 +19,7 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [tool.bumpversion]
-current_version = "0.3.17"
+current_version = "0.3.19"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.3.17 → bayinx-0.3.19}/src/bayinx/core/_variational.py

@@ -63,6 +63,16 @@ class Variational(eqx.Module, Generic[M]):
         """
         pass
 
+    @eqx.filter_jit
+    def reconstruct_model(self, draw: Array) -> M:
+        # Unflatten variational draw
+        model: M = self._unflatten(draw)
+
+        # Combine with constraints
+        model: M = eqx.combine(model, self._constraints)
+
+        return model
+
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0, None))
     def eval_model(self, draws: Array, data: Any = None) -> Array:
@@ -74,10 +84,7 @@ class Variational(eqx.Module, Generic[M]):
         - `data`: Data used to evaluate the posterior(if needed).
         """
         # Unflatten variational draw
-        model: M = self._unflatten(draws)
-
-        # Combine with constraints
-        model: M = eqx.combine(model, self._constraints)
+        model: M = self.reconstruct_model(draws)
 
         # Evaluate posterior density
         return model.eval(data)
@@ -177,7 +184,7 @@ class Variational(eqx.Module, Generic[M]):
     ) -> Array:
         # Sample a single draw to evaluate shape of output
         draw: Array = self.sample(1, key)[0]
-        output: Array = func(self._unflatten(draw), data)
+        output: Array = func(self.reconstruct_model(draw), data)
 
         # Allocate space for results
         results: Array = jnp.zeros((n,) + output.shape, dtype=output.dtype)
@@ -191,10 +198,10 @@ class Variational(eqx.Module, Generic[M]):
             next, key = jr.split(key)
 
             # Draw from variational
-            draw: Array = self.sample(1, key)
+            draw: Array = self.sample(1, key)[0]
 
             # Reconstruct model
-            model: M = self._unflatten(draw)
+            model: M = self.reconstruct_model(draw)
 
             # Update results with output
             results = results.at[i].set(func(model, data))

{bayinx-0.3.17 → bayinx-0.3.19}/src/bayinx/mhx/vi/meanfield.py

@@ -1,11 +1,11 @@
-from typing import Any, Dict, Generic, Self, TypeVar
+from typing import Any, Generic, Self, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
 from jax.flatten_util import ravel_pytree
-from jaxtyping import Array, Float, Key, Scalar
+from jaxtyping import Array, Key, Scalar
 
 from bayinx.core import Model, Variational
 from bayinx.dists import normal
@@ -18,12 +18,14 @@ class MeanField(Variational, Generic[M]):
     A fully factorized Gaussian approximation to a posterior distribution.
 
     # Attributes
-    - `var_params`: The variational parameters for the approximation.
+    - `mean`: The mean of the unconstrained approximation.
+    - `log_std` The log-transformed standard deviation of the unconstrained approximation.
     """
 
-    var_params: Dict[str, Float[Array, "..."]]  # todo: just expand to attributes
+    mean: Array
+    log_std: Array
 
-    def __init__(self, model: M, init_log_std: float = 0.0):
+    def __init__(self, model: M, init_log_std: float = -5.0):
         """
         Constructs an unoptimized meanfield posterior approximation.
 
@@ -38,10 +40,8 @@ class MeanField(Variational, Generic[M]):
         params, self._unflatten = ravel_pytree(params)
 
         # Initialize variational parameters
-        self.var_params = {
-            "mean": params,
-            "log_std": jnp.full(params.size, init_log_std, params.dtype),
-        }
+        self.mean = params
+        self.log_std = jnp.full(params.size, init_log_std, params.dtype)
 
     @property
     @eqx.filter_jit
@@ -51,7 +51,12 @@ class MeanField(Variational, Generic[M]):
 
         # Specify variational parameters
         filter_spec = eqx.tree_at(
-            lambda mf: mf.var_params,
+            lambda mf: mf.mean,
+            filter_spec,
+            replace=True,
+        )
+        filter_spec = eqx.tree_at(
+            lambda mf: mf.log_std,
             filter_spec,
             replace=True,
         )
@@ -62,9 +67,9 @@ class MeanField(Variational, Generic[M]):
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws
         draws: Array = (
-            jr.normal(key=key, shape=(n, self.var_params["mean"].size))
-            * jnp.exp(self.var_params["log_std"])
-            + self.var_params["mean"]
+            jr.normal(key=key, shape=(n, self.mean.size))
+            * jnp.exp(self.log_std)
+            + self.mean
         )
 
         return draws
@@ -73,8 +78,8 @@ class MeanField(Variational, Generic[M]):
     def eval(self, draws: Array) -> Array:
         return normal.logprob(
             x=draws,
-            mu=self.var_params["mean"],
-            sigma=jnp.exp(self.var_params["log_std"]),
+            mu=self.mean,
+            sigma=jnp.exp(self.log_std),
         ).sum(axis=1)
 
     @eqx.filter_jit

{bayinx-0.3.17 → bayinx-0.3.19}/tests/test_variational.py

@@ -43,8 +43,8 @@ def test_meanfield(benchmark, var_draws):
     benchmark(benchmark_fit)
 
     # Assert parameters are roughly correct
-    assert all(abs(10.0 - vari.var_params["mean"]) < 0.1) and all(
-        abs(0.0 - vari.var_params["log_std"]) < 0.1
+    assert all(abs(10.0 - vari.mean) < 0.1) and all(
+        abs(0.0 - vari.log_std) < 0.1
     )