bayinx 0.2.22__tar.gz → 0.2.24__tar.gz

{bayinx-0.2.22 → bayinx-0.2.24}/PKG-INFO

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

{bayinx-0.2.22 → bayinx-0.2.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.22"
+version = "0.2.24"
 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.2.22"
+current_version = "0.2.24"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

bayinx-0.2.24/src/bayinx/core/constraints.py

@@ -0,0 +1,61 @@
+from abc import abstractmethod
+from typing import Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Scalar, ScalarLike
+
+
+class Constraint(eqx.Module):
+    """
+    Abstract base class for defining parameter constraints.
+
+    Subclasses should implement the `constrain` method to apply the
+    transformation and compute the ladj adjustment.
+    """
+    @abstractmethod
+    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
+        """
+        Applies the constraining transformation to an unconstrained input
+        and computes the log absolute determinant of the Jacobian (ladj)
+        of this transformation.
+
+        # Parameters
+        - `x`: The unconstrained JAX Array-like input.
+
+        # Returns
+        A tuple containing:
+            - The constrained JAX Array.
+            - A scalar JAX Array representing the ladj of the transformation.
+        """
+        pass
+
+
+class LowerBound(Constraint):
+    """
+    Enforces a lower bound on the parameter.
+    """
+    lb: ScalarLike
+
+    def __init__(self, lb: ScalarLike):
+        self.lb = lb
+
+    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
+        """
+        Applies the lower bound constraint and computes the ladj.
+
+        # Parameters
+        - `x`: The unconstrained JAX Array-like input.
+
+        # Parameters
+        A tuple containing:
+            - The constrained JAX Array (x > self.lb).
+            - A scalar JAX Array representing the ladj of the transformation.
+        """
+        # Compute transformation adjustment
+        ladj = jnp.sum(x)
+
+        # Compute transformation
+        x = jnp.exp(x) + self.lb
+
+        return x, ladj

bayinx-0.2.24/src/bayinx/core/model.py

@@ -0,0 +1,75 @@
+from abc import abstractmethod
+from typing import Any, Dict, Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from jaxtyping import Array, Scalar
+
+from bayinx.core.constraints import Constraint
+from bayinx.core.utils import __MyMeta
+
+
+class Model(eqx.Module, metaclass=__MyMeta):
+    """
+    A superclass used to define probabilistic models.
+
+    # Attributes
+    - `params`: A dictionary of JAX Arrays representing parameters of the model.
+    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    """
+
+    params: Dict[str, Array]
+    constraints: Dict[str, Constraint]
+
+    @abstractmethod
+    def eval(self, data: Any) -> Scalar:
+        pass
+
+    # Default filter specification
+    def filter_spec(self):
+        """
+        Generates a filter specification to subset relevant parameters for the model.
+        """
+        # Generate empty specification
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        # Specify JAX Array parameters
+        filter_spec = eqx.tree_at(
+            lambda model: model.params,
+            filter_spec,
+            replace=jtu.tree_map(eqx.is_array, self.params),
+        )
+
+        return filter_spec
+
+    # Add constrain method
+    @eqx.filter_jit
+    def constrain_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+        """
+        Constrain `params` to the appropriate domain.
+
+        # Returns
+        A dictionary of transformed JAX Arrays representing the constrained parameters and the adjustment to the posterior density.
+        """
+        t_params: Dict[str, Array] = self.params
+        target: Scalar = jnp.array(0.0)
+
+        for par, map in self.constraints.items():
+            # Apply transformation
+            t_params[par], ladj = map.constrain(t_params[par])
+
+            # Adjust posterior density
+            target -= ladj
+
+        return t_params, target
+
+
+    def transform_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+        """
+        Apply a custom transformation to `params` if needed.
+
+        # Returns
+        A dictionary of transformed JAX Arrays representing the transformed parameters.
+        """
+        return self.constrain_pars()

bayinx-0.2.24/src/bayinx/core/variational.py

@@ -0,0 +1,162 @@
+from abc import abstractmethod
+from functools import partial
+from typing import Any, Callable, Self, Tuple
+
+import equinox as eqx
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import optax as opx
+from jaxtyping import Array, Float, Key, PyTree, Scalar
+from optax import GradientTransformation, OptState, Schedule
+
+from bayinx.core import Model
+
+
+class Variational(eqx.Module):
+    """
+    A superclass used to define variational methods.
+
+    # Attributes
+    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
+    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
+    """
+
+    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _constraints: Model
+
+    @abstractmethod
+    def sample(self, n: int, key: Key) -> Array:
+        """
+        Sample from the variational distribution.
+        """
+        pass
+
+    @abstractmethod
+    def eval(self, draws: Array) -> Array:
+        """
+        Evaluate the variational distribution at `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def elbo(self, n: int, key: Key, data: Any = None) -> Array:
+        """
+        Evaluate the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def elbo_grad(self, n: int, key: Key, data: Any = None) -> PyTree:
+        """
+        Evaluate the gradient of the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def filter_spec(self):
+        """
+        Filter specification for dynamic and static components of the `Variational`.
+        """
+        pass
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0, None))
+    def eval_model(self, draws: Array, data: Any = None) -> Array:
+        """
+        Reconstruct models from variational draws and evaluate their posterior density.
+
+        # Parameters
+        - `draws`: A set of variational draws.
+        - `data`: Data used to evaluate the posterior(if needed).
+        """
+        # Unflatten variational draw
+        model: Model = self._unflatten(draws)
+
+        # Combine with constraints
+        model: Model = eqx.combine(model, self._constraints)
+
+        # Evaluate posterior density
+        return model.eval(data)
+
+    @eqx.filter_jit
+    def fit(
+        self,
+        max_iters: int,
+        data: Any = None,
+        learning_rate: float = 1,
+        weight_decay: float = 1e-4,
+        tolerance: float = 1e-4,
+        var_draws: int = 1,
+        key: Key = jr.PRNGKey(0),
+    ) -> Self:
+        """
+        Optimize the variational distribution.
+
+        # Parameters
+        - `max_iters`: Maximum number of iterations for the optimization loop.
+        - `data`: Data to evaluate the posterior density with(if available).
+        - `learning_rate`: Initial learning rate for optimization.
+        - `tolerance`: Relative tolerance of ELBO decrease for stopping early.
+        - `var_draws`: Number of variational draws to draw each iteration.
+        - `key`: A PRNG key.
+        """
+        # Partition variational
+        dyn, static = eqx.partition(self, self.filter_spec())
+
+        # Construct scheduler
+        schedule: Schedule = opx.cosine_decay_schedule(
+            init_value=learning_rate, decay_steps=max_iters
+        )
+
+        # Initialize optimizer
+        optim: GradientTransformation = opx.chain(
+            opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
+        )
+        opt_state: OptState = optim.init(dyn)
+
+        # Optimization loop helper functions
+        @eqx.filter_jit
+        def condition(state: Tuple[Self, OptState, Scalar, Key]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
+
+            return i < max_iters
+
+        @eqx.filter_jit
+        def body(state: Tuple[Self, OptState, Scalar, Key]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
+
+            # Update iteration
+            i = i + 1
+
+            # Update PRNG key
+            key, _ = jr.split(key)
+
+            # Combine variational
+            vari = eqx.combine(dyn, static)
+
+            # Compute gradient of the ELBO
+            updates: PyTree = vari.elbo_grad(var_draws, key, data)
+
+            # Compute updates
+            updates, opt_state = optim.update(
+                updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
+            )
+
+            # Update variational distribution
+            dyn = eqx.apply_updates(dyn, updates)
+
+            return dyn, opt_state, i, key
+
+        # Run optimization loop
+        dyn = lax.while_loop(
+            cond_fun=condition,
+            body_fun=body,
+            init_val=(dyn, opt_state, jnp.array(0, jnp.uint32), key),
+        )[0]
+
+        # Return optimized variational
+        return eqx.combine(dyn, static)

{bayinx-0.2.22 → bayinx-0.2.24}/src/bayinx/mhx/vi/flows/planar.py

@@ -30,7 +30,7 @@ class Planar(Flow):
         - `dim`: The dimension of the parameter space.
         """
         self.params = {
-            "u": jnp.ones(dim),
+            "u": jnp.zeros(dim),
             "w": jnp.ones(dim),
             "b": jnp.zeros(1),
         }

{bayinx-0.2.22 → bayinx-0.2.24}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -144,6 +144,7 @@ class NormalizingFlow(Variational):
             draws: Array = self.base.sample(n, key)
 
             posterior_evals, variational_evals = self.__eval(draws, data)
+
             # Evaluate ELBO
             return jnp.mean(posterior_evals - variational_evals)
 

{bayinx-0.2.22 → bayinx-0.2.24}/tests/test_variational.py

@@ -26,12 +26,13 @@ def test_meanfield(benchmark, var_draws):
         @eqx.filter_jit
         def eval(self, data: dict):
             # Get constrained parameters
-            params = self.constrain_pars()
+            params, target = self.constrain_pars()
 
             # Evaluate mu ~ N(10,1)
-            return jnp.sum(
-                normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
-            )
+            target += normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0)).sum()
+
+            # Evaluate mu ~ N(10,1)
+            return target
 
     # Construct model
     model = NormalDist()
@@ -44,7 +45,7 @@ def test_meanfield(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(20000)
+    vari = vari.fit(20000,var_draws=var_draws)
 
     # Assert parameters are roughly correct
     assert all(abs(10.0 - vari.var_params["mean"]) < 0.1) and all(
@@ -66,12 +67,13 @@ def test_affine(benchmark, var_draws):
         @eqx.filter_jit
         def eval(self, data: dict):
             # Get constrained parameters
-            params = self.constrain_pars()
+            params, target = self.constrain_pars()
 
             # Evaluate mu ~ N(10,1)
-            return jnp.sum(
-                normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
-            )
+            target += normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0)).sum()
+
+            # Evaluate mu ~ N(10,1)
+            return target
 
     # Construct model
     model = NormalDist()
@@ -84,7 +86,7 @@ def test_affine(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(20000)
+    vari = vari.fit(20000,var_draws=var_draws)
 
     params = vari.flows[0].constrain_pars()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
@@ -106,12 +108,12 @@ def test_flows(benchmark, var_draws):
         @eqx.filter_jit
         def eval(self, data: dict):
             # Get constrained parameters
-            params = self.constrain_pars()
+            params, target = self.constrain_pars()
 
             # Evaluate mu ~ N(10,1)
-            return jnp.sum(
-                normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
-            )
+            target += normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0)).sum()
+
+            return target
 
     # Construct model
     model = NormalDist()
@@ -126,7 +128,7 @@ def test_flows(benchmark, var_draws):
         vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
-    vari = vari.fit(20000)
+    vari = vari.fit(20000,var_draws=var_draws)
 
     mean = vari.sample(1000).mean(0)
     var = vari.sample(1000).var(0)

bayinx-0.2.22/src/bayinx/core/model.py

@@ -1,74 +0,0 @@
-from abc import abstractmethod
-from typing import Any, Callable, Dict
-
-import equinox as eqx
-import jax.tree_util as jtu
-from jaxtyping import Array, Scalar
-
-from bayinx.core.utils import __MyMeta
-
-
-class Model(eqx.Module, metaclass=__MyMeta):
-    """
-    A superclass used to define probabilistic models.
-
-    # Attributes
-    - `params`: A dictionary of JAX Arrays representing parameters of the model.
-    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
-    """
-
-    params: Dict[str, Array]
-    constraints: Dict[str, Callable[[Array], Array]]
-
-    @abstractmethod
-    def eval(self, data: Any) -> Scalar:
-        pass
-
-    # Default filter specification
-    def filter_spec(self):
-        """
-        Generates a filter specification to subset relevant parameters for the model.
-        """
-        # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        # Specify JAX Array parameters
-        filter_spec = eqx.tree_at(
-            lambda model: model.params,
-            filter_spec,
-            replace=jtu.tree_map(eqx.is_array, self.params),
-        )
-
-        return filter_spec
-
-    def __init_subclass__(cls):
-        # Add constrain method
-        def constrain_pars(self: Model) -> Dict[str, Array]:
-            """
-            Constrain `params` to the appropriate domain.
-
-            # Returns
-            A dictionary of transformed JAX Arrays representing the constrained parameters.
-            """
-            t_params = self.params
-
-            for par, map in self.constraints.items():
-                t_params[par] = map(t_params[par])
-
-            return t_params
-
-        cls.constrain_pars = eqx.filter_jit(constrain_pars)
-
-        # Add transform_pars method if not present
-        if not callable(getattr(cls, "transform_pars", None)):
-
-            def transform_pars(self: Model) -> Dict[str, Array]:
-                """
-                Apply a custom transformation to `params` if needed.
-
-                # Returns
-                A dictionary of transformed JAX Arrays representing the transformed parameters.
-                """
-                return self.constrain_pars()
-
-            cls.transform_pars = eqx.filter_jit(transform_pars)

bayinx-0.2.22/src/bayinx/core/variational.py

@@ -1,167 +0,0 @@
-from abc import abstractmethod
-from typing import Any, Callable, Self, Tuple
-
-import equinox as eqx
-import jax
-import jax.lax as lax
-import jax.numpy as jnp
-import jax.random as jr
-import optax as opx
-from jaxtyping import Array, Float, Key, PyTree, Scalar
-from optax import GradientTransformation, OptState, Schedule
-
-from bayinx.core import Model
-
-
-class Variational(eqx.Module):
-    """
-    A superclass used to define variational methods.
-
-    # Attributes
-    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
-    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
-    """
-
-    _unflatten: Callable[[Float[Array, "..."]], Model]
-    _constraints: Model
-
-    @abstractmethod
-    def sample(self, n: int, key: Key) -> Array:
-        """
-        Sample from the variational distribution.
-        """
-        pass
-
-    @abstractmethod
-    def eval(self, draws: Array) -> Array:
-        """
-        Evaluate the variational distribution at `draws`.
-        """
-        pass
-
-    @abstractmethod
-    def elbo(self, n: int, key: Key, data: Any = None) -> Array:
-        """
-        Evaluate the ELBO.
-        """
-        pass
-
-    @abstractmethod
-    def elbo_grad(self, n: int, key: Key, data: Any = None) -> PyTree:
-        """
-        Evaluate the gradient of the ELBO.
-        """
-        pass
-
-    @abstractmethod
-    def filter_spec(self):
-        """
-        Filter specification for dynamic and static components of the `Variational`.
-        """
-        pass
-
-    def __init_subclass__(cls):
-        """
-        Construct methods that are shared across all VI methods.
-        """
-
-        def eval_model(self, draws: Array, data: Any = None) -> Array:
-            """
-            Reconstruct models from variational draws and evaluate their posterior density.
-
-            # Parameters
-            - `draws`: A set of variational draws.
-            - `data`: Data used to evaluate the posterior(if needed).
-            """
-            # Unflatten variational draw
-            model: Model = self._unflatten(draws)
-
-            # Combine with constraints
-            model: Model = eqx.combine(model, self._constraints)
-
-            # Evaluate posterior density
-            return model.eval(data)
-
-        cls.eval_model = jax.vmap(eqx.filter_jit(eval_model), (None, 0, None))
-
-        def fit(
-            self,
-            max_iters: int,
-            data: Any = None,
-            learning_rate: float = 1,
-            weight_decay: float = 1e-4,
-            tolerance: float = 1e-4,
-            var_draws: int = 1,
-            key: Key = jr.PRNGKey(0),
-        ) -> Self:
-            """
-            Optimize the variational distribution.
-
-            # Parameters
-            - `max_iters`: Maximum number of iterations for the optimization loop.
-            - `data`: Data to evaluate the posterior density with(if available).
-            - `learning_rate`: Initial learning rate for optimization.
-            - `tolerance`: Relative tolerance of ELBO decrease for stopping early.
-            - `var_draws`: Number of variational draws to draw each iteration.
-            - `key`: A PRNG key.
-            """
-            # Partition variational
-            dyn, static = eqx.partition(self, self.filter_spec())
-
-            # Construct scheduler
-            schedule: Schedule = opx.cosine_decay_schedule(
-                init_value=learning_rate, decay_steps=max_iters
-            )
-
-            # Initialize optimizer
-            optim: GradientTransformation = opx.chain(
-                opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
-            )
-            opt_state: OptState = optim.init(dyn)
-
-            # Optimization loop helper functions
-            @eqx.filter_jit
-            def condition(state: Tuple[Self, OptState, Scalar, Key]):
-                # Unpack iteration state
-                dyn, opt_state, i, key = state
-
-                return i < max_iters
-
-            @eqx.filter_jit
-            def body(state: Tuple[Self, OptState, Scalar, Key]):
-                # Unpack iteration state
-                dyn, opt_state, i, key = state
-
-                # Update iteration
-                i = i + 1
-
-                # Update PRNG key
-                key, _ = jr.split(key)
-
-                # Combine variational
-                vari = eqx.combine(dyn, static)
-
-                # Compute gradient of the ELBO
-                updates: PyTree = vari.elbo_grad(var_draws, key, data)
-
-                # Compute updates
-                updates, opt_state = optim.update(
-                    updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
-                )
-
-                # Update variational distribution
-                dyn = eqx.apply_updates(dyn, updates)
-
-                return dyn, opt_state, i, key
-
-            # Run optimization loop
-            dyn = lax.while_loop(
-                cond_fun=condition,
-                body_fun=body,
-                init_val=(dyn, opt_state, jnp.array(0, jnp.uint32), key),
-            )[0]
-
-            # Return optimized variational
-            return eqx.combine(dyn, static)
-
-        cls.fit = eqx.filter_jit(fit)