bayinx 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

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/core/model.py

@@ -1,10 +1,12 @@
 from abc import abstractmethod
-from typing import Any, Callable, Dict
+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
 
 
@@ -18,7 +20,7 @@ class Model(eqx.Module, metaclass=__MyMeta):
     """
 
     params: Dict[str, Array]
-    constraints: Dict[str, Callable[[Array], Array]]
+    constraints: Dict[str, Constraint]
 
     @abstractmethod
     def eval(self, data: Any) -> Scalar:
@@ -41,34 +43,33 @@ class Model(eqx.Module, metaclass=__MyMeta):
 
         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
+    # Add constrain method
+    @eqx.filter_jit
+    def constrain_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+        """
+        Constrain `params` to the appropriate domain.
 
-            for par, map in self.constraints.items():
-                t_params[par] = map(t_params[par])
+        # 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)
 
-            return t_params
+        for par, map in self.constraints.items():
+            # Apply transformation
+            t_params[par], ladj = map.constrain(t_params[par])
 
-        cls.constrain_pars = eqx.filter_jit(constrain_pars)
+            # Adjust posterior density
+            target -= ladj
 
-        # Add transform_pars method if not present
-        if not callable(getattr(cls, "transform_pars", None)):
+        return t_params, target
 
-            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()
+    def transform_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+        """
+        Apply a custom transformation to `params` if needed.
 
-            cls.transform_pars = eqx.filter_jit(transform_pars)
+        # Returns
+        A dictionary of transformed JAX Arrays representing the transformed parameters.
+        """
+        return self.constrain_pars()

bayinx/core/variational.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from functools import partial
 from typing import Any, Callable, Self, Tuple
 
 import equinox as eqx
@@ -60,108 +61,102 @@ class Variational(eqx.Module):
         """
         pass
 
-    def __init_subclass__(cls):
-        """
-        Construct methods that are shared across all VI methods.
+    @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.
 
-        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
-            )
+        # 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())
 
-            # Initialize optimizer
-            optim: GradientTransformation = opx.chain(
-                opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
-            )
-            opt_state: OptState = optim.init(dyn)
+        # Construct scheduler
+        schedule: Schedule = opx.cosine_decay_schedule(
+            init_value=learning_rate, decay_steps=max_iters
+        )
 
-            # Optimization loop helper functions
-            @eqx.filter_jit
-            def condition(state: Tuple[Self, OptState, Scalar, Key]):
-                # Unpack iteration state
-                dyn, opt_state, i, key = state
+        # Initialize optimizer
+        optim: GradientTransformation = opx.chain(
+            opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
+        )
+        opt_state: OptState = optim.init(dyn)
 
-                return i < max_iters
+        # Optimization loop helper functions
+        @eqx.filter_jit
+        def condition(state: Tuple[Self, OptState, Scalar, Key]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
 
-            @eqx.filter_jit
-            def body(state: Tuple[Self, OptState, Scalar, Key]):
-                # Unpack iteration state
-                dyn, opt_state, i, key = state
+            return i < max_iters
 
-                # Update iteration
-                i = i + 1
+        @eqx.filter_jit
+        def body(state: Tuple[Self, OptState, Scalar, Key]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
 
-                # Update PRNG key
-                key, _ = jr.split(key)
+            # Update iteration
+            i = i + 1
 
-                # Combine variational
-                vari = eqx.combine(dyn, static)
+            # Update PRNG key
+            key, _ = jr.split(key)
 
-                # Compute gradient of the ELBO
-                updates: PyTree = vari.elbo_grad(var_draws, key, data)
+            # Combine variational
+            vari = eqx.combine(dyn, static)
 
-                # Compute updates
-                updates, opt_state = optim.update(
-                    updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
-                )
+            # Compute gradient of the ELBO
+            updates: PyTree = vari.elbo_grad(var_draws, key, data)
 
-                # Update variational distribution
-                dyn = eqx.apply_updates(dyn, updates)
+            # Compute updates
+            updates, opt_state = optim.update(
+                updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
+            )
 
-                return dyn, opt_state, i, key
+            # Update variational distribution
+            dyn = eqx.apply_updates(dyn, updates)
 
-            # 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 dyn, opt_state, i, key
 
-            # Return optimized variational
-            return eqx.combine(dyn, static)
+        # 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]
 
-        cls.fit = eqx.filter_jit(fit)
+        # Return optimized variational
+        return eqx.combine(dyn, static)

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/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.dist-info → bayinx-0.2.24.dist-info}/METADATA

@@ -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.dist-info → bayinx-0.2.24.dist-info}/RECORD

@@ -1,10 +1,11 @@
 bayinx/__init__.py,sha256=l20JdkSsE_XGZlZFNEtySXf4NIlbjrao14vXPB-H6aQ,45
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
+bayinx/core/constraints.py,sha256=Y8FJX3CkgnLQ3HXuTPGuzvLtXlKs0B7z0-YymoHgdfg,1682
 bayinx/core/flow.py,sha256=oZE0OHCninIHjp-WVLFWd1DaN0-qXxNWFAUAdgIDmRU,2423
-bayinx/core/model.py,sha256=-rT3NHjxqGB0lDBMi0Mr9XNOz1_TUnJWtd4ITj0rsus,2257
+bayinx/core/model.py,sha256=t7s5Yt4E3iC_MPvynJnk6lb4OLal7Gnk59tIZ6e5M4I,2203
 bayinx/core/utils.py,sha256=-YewhqzMFL3GJEjVdm3LgaZyHwDs9IVYllU9wAXZrtw,1859
-bayinx/core/variational.py,sha256=k9wWn7Tnj3eET-qK1pZtzDyPZVvQTRUexJUBVSdGXOA,5251
+bayinx/core/variational.py,sha256=vUZ6u5CXCHfs6ZrA8PF4PHfmUXHTK2RJGHyZ3afFfsg,4820
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/binomial.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -15,13 +16,13 @@ bayinx/dists/uniform.py,sha256=PSZIIc2QfNF5XYi-TLGltnr_vnAIA-MZsn1rKV8QXAo,2353
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
 bayinx/mhx/vi/meanfield.py,sha256=LNLwfjKO9os7YBmRBpGEiFxlxonuN7DHVFEmXV3hvfI,3876
-bayinx/mhx/vi/normalizing_flow.py,sha256=XBUWYZpm_Ipi6X9oTnGhqIs3ARY-5QFiuxM7uAWFRps,4790
+bayinx/mhx/vi/normalizing_flow.py,sha256=nj7bpIoMJl6GTOXPxQCAsPArchbHd5vwwqMm3cLbMII,4791
 bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
 bayinx/mhx/vi/flows/fullaffine.py,sha256=2QbOtA1Jmu-yRcJeFmCKc8N1atm8G7JXYMLEZaEXKV0,1711
-bayinx/mhx/vi/flows/planar.py,sha256=qmtWpIBXRct2seI78pkmtF0X7cASUBELqmZmf2QS5Gs,1918
+bayinx/mhx/vi/flows/planar.py,sha256=u9ZVwEeOv4fMfwiORlseCz463atsWMuid_9crRg05Z8,1919
 bayinx/mhx/vi/flows/radial.py,sha256=c-SWybGn_jmgBQk9ZMQ5uHKPzcdhowNp8MD8t1-8VZM,2501
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.22.dist-info/METADATA,sha256=TwM3DxTXPTttN0TGJOPpM1pRMzDvoUjFZWkWWCR9vNI,3058
-bayinx-0.2.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.22.dist-info/RECORD,,
+bayinx-0.2.24.dist-info/METADATA,sha256=sR0C0Pk5vrAmdvAtB3faXZO-hIDpKzqLjnXcfMsikjw,3058
+bayinx-0.2.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.24.dist-info/RECORD,,