bayinx 0.3.16__py3-none-any.whl → 0.3.18__py3-none-any.whl

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)
@@ -176,8 +183,8 @@ class Variational(eqx.Module, Generic[M]):
         key: Key = jr.PRNGKey(0),
     ) -> Array:
         # Sample a single draw to evaluate shape of output
-        draw: Array = self.sample(1, key)
-        output: Array = func(self._unflatten(draw), data)
+        draw: Array = self.sample(1, key)[0]
+        output: Array = func(self.reconstruct_model(draw), data)
 
         # Allocate space for results
         results: Array = jnp.zeros((n,) + output.shape, dtype=output.dtype)
@@ -194,7 +201,7 @@ class Variational(eqx.Module, Generic[M]):
             draw: Array = self.sample(1, key)
 
             # 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/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.16.dist-info → bayinx-0.3.18.dist-info}/METADATA

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

{bayinx-0.3.16.dist-info → bayinx-0.3.18.dist-info}/RECORD

@@ -8,7 +8,7 @@ bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
 bayinx/core/_optimization.py,sha256=mmeVUqfFARz8F7q4LRl-uEwVWzekmzh-9o7PnuvsHZk,2651
 bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-bayinx/core/_variational.py,sha256=MbOgWs0-mkEwIYsh1Hhn7anWkQ7fAGjkr6VJybYAuR0,6185
+bayinx/core/_variational.py,sha256=NlCvBHsfP-7HeNimhD4x9JLT5VWeiCBUADr3bx6tt84,6381
 bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=HtB60LUQdPj4yDAHme2jsHNmLfrAKWsSZnDYkxAGaOI,1548
@@ -23,7 +23,7 @@ bayinx/dists/censored/posnormal/r.py,sha256=wMDt2Am1TD376ms8B-o6PFCJZXmUJd2-aBC-
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
-bayinx/mhx/vi/meanfield.py,sha256=wM0v6Q2m0pPyEdOaT8DvFqdRLsHijr7AYNQKAcsZtRQ,3881
+bayinx/mhx/vi/meanfield.py,sha256=iX4AeDG9jrLZd6d9NimuJ3O5zaoBXsD03JbgPgxVrfY,3917
 bayinx/mhx/vi/normalizing_flow.py,sha256=vzLu5H1G1-pBqhgHWmIZkUTyPE1DxC9vBwpiZeIyu1I,4712
 bayinx/mhx/vi/standard.py,sha256=LYgglaGQMGmXpzFR4eMJnXkl2PhBeggbXMvO5zJpf2c,1578
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
@@ -31,7 +31,7 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvAD
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.16.dist-info/METADATA,sha256=9oLNZAc1YvN9FpWiuFFlztL-OkhclOYPBQFkuiHmUT8,3087
-bayinx-0.3.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.16.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.16.dist-info/RECORD,,
+bayinx-0.3.18.dist-info/METADATA,sha256=osJ-IhnQs0OuZoA6AGk0GKSwibISSNaOeyLg5F0T4w4,3087
+bayinx-0.3.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.18.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.18.dist-info/RECORD,,