bayinx 0.2.11__tar.gz → 0.2.13__tar.gz

{bayinx-0.2.11 → bayinx-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.2.11
+Version: 0.2.13
 Summary: Bayesian Inference with JAX
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12

{bayinx-0.2.11 → bayinx-0.2.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.11"
+version = "0.2.13"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,4 +19,4 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [dependency-groups]
-dev = ["pytest>=8.3.5", "pytest-benchmark>=5.1.0"]
+dev = ["matplotlib>=3.10.1", "pytest>=8.3.5", "pytest-benchmark>=5.1.0"]

{bayinx-0.2.11 → bayinx-0.2.13}/src/bayinx/core/variational.py

@@ -115,7 +115,7 @@ class Variational(eqx.Module):
 
             # Initialize optimizer
             optim: GradientTransformation = opx.chain(
-                opx.scale(-1.0), opx.nadamw(schedule,weight_decay=weight_decay)
+                opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
             )
             opt_state: OptState = optim.init(dyn)
 

{bayinx-0.2.11 → bayinx-0.2.13}/src/bayinx/dists/normal.py

@@ -1,17 +1,12 @@
-# MARK: Imports ----
 import jax.lax as _lax
+from jaxtyping import Array, ArrayLike, Float, Real
 
-## Typing
-from jaxtyping import Array, Real
-
-# MARK: Constants
 _PI = 3.141592653589793
 
 
-# MARK: Functions ----
 def prob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Normal distribution.
 
@@ -30,8 +25,8 @@ def prob(
 
 
 def logprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Normal distribution.
 
@@ -48,8 +43,8 @@ def logprob(
 
 
 def uprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Normal distribution.
 
@@ -66,8 +61,8 @@ def uprob(
 
 
 def ulogprob(
-    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
-) -> Real[Array, "..."]:
+    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
 

{bayinx-0.2.11 → bayinx-0.2.13}/src/bayinx/mhx/vi/flows/planar.py

@@ -3,7 +3,6 @@ from typing import Callable, Dict, Tuple
 
 import equinox as eqx
 import jax
-import jax.nn as jnn
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import Array, Float, Scalar
@@ -31,25 +30,12 @@ class Planar(Flow):
         - `dim`: The dimension of the parameter space.
         """
         self.params = {
-            "u": jr.normal(key, (dim,)),
-            "w": jr.normal(key, (dim,)),
-            "b": jr.normal(key, (1,)),
+            "u": jnp.ones(dim),
+            "w": jnp.ones(dim),
+            "b": jnp.zeros(1),
         }
         self.constraints = {}
 
-    def transform_pars(self):
-        params = self.params
-
-        u = params['u']
-        w = params['w']
-        b = params['b']
-
-        m = jnn.softplus(w.dot(u)) - 1.0
-
-        u = u + (m - w.dot(u)) * w / (w**2).sum()
-
-        return {'u': u, 'w': w, 'b': b}
-
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:

{bayinx-0.2.11 → bayinx-0.2.13}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -1,6 +1,8 @@
+from functools import partial
 from typing import Any, Callable, Self, Tuple
 
 import equinox as eqx
+import jax
 import jax.flatten_util as jfu
 import jax.numpy as jnp
 import jax.random as jr
@@ -59,6 +61,7 @@ class NormalizingFlow(Variational):
         return draws
 
     @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
     def eval(self, draws: Array) -> Array:
         # Evaluate base density
         variational_evals: Array = self.base.eval(draws)

{bayinx-0.2.11 → bayinx-0.2.13}/tests/test_variational.py

@@ -8,7 +8,7 @@ from jaxtyping import Array
 from bayinx import Model
 from bayinx.dists import normal
 from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
-from bayinx.mhx.vi.flows import FullAffine, Planar
+from bayinx.mhx.vi.flows import FullAffine, Planar, Radial
 
 
 # Tests ----
@@ -44,7 +44,7 @@ def test_meanfield(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(20000)
 
     # Assert parameters are roughly correct
     assert all(abs(10.0 - vari.var_params["mean"]) < 0.1) and all(
@@ -84,7 +84,7 @@ def test_affine(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(10000)
+    vari = vari.fit(20000)
 
     params = vari.flows[0].constrain_pars()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
@@ -95,37 +95,30 @@ def test_affine(benchmark, var_draws):
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_flows(benchmark, var_draws):
     # Construct model definition
-    class Banana(Model):
+    class NormalDist(Model):
         params: Dict[str, Array]
         constraints: Dict[str, Callable[[Array], Array]]
 
         def __init__(self):
-            self.params = {
-                'x': jnp.array(0.0),
-                'y': jnp.array(0.0)
-            }
+            self.params = {"mu": jnp.array([0.0, 0.0])}
             self.constraints = {}
 
-        def eval(self, data = None):
-            params: Dict[str, Array] = self.params
-            # Extract parameters
-            x: Array = params['x']
-            y: Array = params['y']
-
-            # Initialize target density
-            target = jnp.array(0.0)
-
-            target += normal.logprob(x, mu = jnp.array(0.0), sigma = jnp.array(1.0))
-            target += normal.logprob(y, mu = x**2 + x, sigma = jnp.array(1.0))
+        @eqx.filter_jit
+        def eval(self, data: dict):
+            # Get constrained parameters
+            params = self.constrain_pars()
 
-            return target
+            # Evaluate mu ~ N(10,1)
+            return jnp.sum(
+                normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
+            )
 
     # Construct model
-    model = Banana()
+    model = NormalDist()
 
     # Construct normalizing flow variational
     vari = NormalizingFlow(
-        Standard(model), [FullAffine(2), Planar(2)], model
+        Standard(model), [FullAffine(2), Planar(2), Radial(2)], model
     )
 
     # Optimize variational distribution
@@ -133,7 +126,7 @@ def test_flows(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(100)
+    vari = vari.fit(20000)
 
     mean = vari.sample(1000).mean(0)
     var = vari.sample(1000).var(0)