bayinx 0.2.33__tar.gz → 0.3.2__tar.gz

{bayinx-0.2.33 → bayinx-0.3.2}/PKG-INFO

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

{bayinx-0.2.33 → bayinx-0.3.2}/pyproject.toml

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

bayinx-0.3.2/src/bayinx/__init__.py

@@ -0,0 +1,2 @@
+from bayinx.core import Model as Model
+from bayinx.core import Parameter as Parameter

bayinx-0.3.2/src/bayinx/constraints/__init__.py

@@ -0,0 +1 @@
+from bayinx.constraints.lower import Lower as Lower

bayinx-0.3.2/src/bayinx/constraints/lower.py

@@ -0,0 +1,51 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import PyTree, Scalar, ScalarLike
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.parameter import Parameter
+
+
+class Lower(Constraint):
+    """
+    Enforces a lower bound on the parameter.
+    """
+
+    lb: Scalar
+
+    def __init__(self, lb: ScalarLike):
+        self.lb = jnp.array(lb)
+
+    @eqx.filter_jit
+    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+        """
+        Enforces a lower bound on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `x`: The unconstrained `Parameter`.
+
+        # Parameters
+        A tuple containing:
+            - A modified `Parameter` with relevant leaves satisfying the constraint.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        # Extract relevant filter specification
+        filter_spec = x.filter_spec
+
+        # Extract relevant parameters(all Array)
+        dyn_params, static_params = eqx.partition(x, filter_spec)
+
+        # Compute density adjustment
+        laj: PyTree = jt.map(jnp.sum, dyn_params) # pyright: ignore
+        laj: Scalar = jt.reduce(lambda a,b: a + b, laj)
+
+        # Compute transformation
+        dyn_params = jt.map(lambda v: jnp.exp(v) + self.lb, dyn_params)
+
+        # Combine into full parameter object
+        x = eqx.combine(dyn_params, static_params)
+
+        return x, laj

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/core/__init__.py

@@ -1,3 +1,4 @@
 from bayinx.core.flow import Flow as Flow
 from bayinx.core.model import Model as Model
+from bayinx.core.parameter import Parameter as Parameter
 from bayinx.core.variational import Variational as Variational

bayinx-0.3.2/src/bayinx/core/constraint.py

@@ -0,0 +1,28 @@
+from abc import abstractmethod
+from typing import Tuple
+
+import equinox as eqx
+from jaxtyping import Scalar
+
+from bayinx.core.parameter import Parameter
+
+
+class Constraint(eqx.Module):
+    """
+    Abstract base class for defining parameter constraints.
+    """
+
+    @abstractmethod
+    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+        """
+        Applies the constraining transformation to a parameter and computes the log-absolute-Jacobian of the transformation.
+
+        # Parameters
+        - `x`: The unconstrained `Parameter`.
+
+        # Returns
+        A tuple containing:
+            - The constrained `Parameter`.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        pass

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/core/flow.py

@@ -31,11 +31,13 @@ class Flow(eqx.Module):
         Computes the log-absolute-Jacobian at `draws` and applies the forward transformation.
 
         # Returns
-        A tuple of JAX Arrays containing the transformed draws and log-absolute-Jacobians.
+            A tuple of JAX Arrays containing the transformed draws and log-absolute-Jacobians.
         """
         pass
 
     # Default filter specification
+    @property
+    @eqx.filter_jit
     def filter_spec(self):
         """
         Generates a filter specification to subset relevant parameters for the flow.
@@ -53,7 +55,7 @@ class Flow(eqx.Module):
         return filter_spec
 
     @eqx.filter_jit
-    def constrain_pars(self: Self):
+    def constrain_params(self: Self):
         """
         Constrain `params` to the appropriate domain.
 
@@ -68,11 +70,11 @@ class Flow(eqx.Module):
         return t_params
 
     @eqx.filter_jit
-    def transform_pars(self: Self) -> Dict[str, Array]:
+    def transform_params(self: Self) -> 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()
+        return self.constrain_params()

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/core/model.py

@@ -1,24 +1,25 @@
 from abc import abstractmethod
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Generic, Tuple, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
 import jax.tree as jt
-from jaxtyping import Array, PyTree, Scalar
+from jaxtyping import PyTree, Scalar
 
 from bayinx.core.constraint import Constraint
+from bayinx.core.parameter import Parameter
 
-
-class Model(eqx.Module):
+P = TypeVar('P', bound=Dict[str, Parameter[PyTree]])
+class Model(eqx.Module, Generic[P]):
     """
     An abstract base class used to define probabilistic models.
 
     # Attributes
-    - `params`: A dictionary of JAX Arrays representing parameters of the model.
+    - `params`: A dictionary of parameters.
     - `constraints`: A dictionary of constraints.
     """
 
-    params: Dict[str, PyTree]
+    params: P
     constraints: Dict[str, Constraint]
 
     @abstractmethod
@@ -26,6 +27,8 @@ class Model(eqx.Module):
         pass
 
     # Default filter specification
+    @property
+    @eqx.filter_jit
     def filter_spec(self):
         """
         Generates a filter specification to subset relevant parameters for the model.
@@ -33,25 +36,25 @@ class Model(eqx.Module):
         # Generate empty specification
         filter_spec = jt.map(lambda _: False, self)
 
-        # Specify JAX Array parameters
+        # Specify relevant parameters
         filter_spec = eqx.tree_at(
             lambda model: model.params,
             filter_spec,
-            replace=jt.map(eqx.is_array, self.params),
+            replace={key: param.filter_spec for key, param in self.params.items()}
         )
 
         return filter_spec
 
     # Add constrain method
     @eqx.filter_jit
-    def constrain_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+    def constrain_params(self) -> Tuple[P, 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.
+        A dictionary of PyTrees representing the constrained parameters and the adjustment to the posterior density.
         """
-        t_params: Dict[str, Array] = self.params
+        t_params: P = self.params
         target: Scalar = jnp.array(0.0)
 
         for par, map in self.constraints.items():
@@ -64,11 +67,12 @@ class Model(eqx.Module):
         return t_params, target
 
     # Add default transform method
-    def transform_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+    @eqx.filter_jit
+    def transform_params(self) -> Tuple[P, Scalar]:
         """
         Apply a custom transformation to `params` if needed.
 
         # Returns
         A dictionary of transformed JAX Arrays representing the transformed parameters.
         """
-        return self.constrain_pars()
+        return self.constrain_params()

bayinx-0.3.2/src/bayinx/core/parameter.py

@@ -0,0 +1,41 @@
+from typing import Generic, Self, TypeVar
+
+import equinox as eqx
+import jax.tree as jt
+from jaxtyping import PyTree
+
+T = TypeVar('T', bound=PyTree)
+class Parameter(eqx.Module, Generic[T]):
+    """
+    A container for a parameter of a `Model`.
+
+    Subclasses can be constructed for custom filter specifications(`filter_spec`).
+
+    # Attributes
+    - `vals`: The parameter's value(s).
+    """
+    vals: T
+
+
+    def __init__(self, values: T):
+        # Insert parameter values
+        self.vals = values
+
+    # Default filter specification
+    @property
+    @eqx.filter_jit
+    def filter_spec(self) -> Self:
+        """
+        Generates a filter specification to filter out static parameters.
+        """
+        # Generate empty specification
+        filter_spec = jt.map(lambda _: False, self)
+
+        # Specify Array leaves
+        filter_spec = eqx.tree_at(
+            lambda params: params.vals,
+            filter_spec,
+            replace=jt.map(eqx.is_array_like, self.vals),
+        )
+
+        return filter_spec

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/core/variational.py

@@ -103,7 +103,7 @@ class Variational(eqx.Module):
         - `key`: A PRNG key.
         """
         # Partition variational
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         # Construct scheduler
         schedule: Schedule = opx.cosine_decay_schedule(
@@ -143,7 +143,7 @@ class Variational(eqx.Module):
 
             # Compute updates
             updates, opt_state = optim.update(
-                updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
+                updates, opt_state, eqx.filter(dyn, dyn.filter_spec)
             )
 
             # Update variational distribution

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/flows/fullaffine.py

@@ -46,7 +46,7 @@ class FullAffine(Flow):
 
     @eqx.filter_jit
     def forward(self, draws: Array) -> Array:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         shift: Array = params["shift"]
@@ -60,7 +60,7 @@ class FullAffine(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         shift: Array = params["shift"]

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/flows/planar.py

@@ -39,7 +39,7 @@ class Planar(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         w: Array = params["w"]
@@ -54,7 +54,7 @@ class Planar(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         w: Array = params["w"]

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/flows/radial.py

@@ -49,7 +49,7 @@ class Radial(Flow):
         # Returns
         The transformed samples.
         """
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         alpha = params["alpha"]
@@ -67,7 +67,7 @@ class Radial(Flow):
     @partial(jax.vmap, in_axes=(None, 0))
     @eqx.filter_jit
     def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         alpha = params["alpha"]

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/meanfield.py

@@ -29,7 +29,7 @@ class MeanField(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec())
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
@@ -40,6 +40,22 @@ class MeanField(Variational):
             "log_std": jnp.zeros(params.size, dtype=params.dtype),
         }
 
+    @property
+    @eqx.filter_jit
+    def filter_spec(self):
+        # Generate empty specification
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        # Specify variational parameters
+        filter_spec = eqx.tree_at(
+            lambda mf: mf.var_params,
+            filter_spec,
+            replace=True,
+        )
+
+        return filter_spec
+
+
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws
@@ -59,23 +75,9 @@ class MeanField(Variational):
             sigma=jnp.exp(self.var_params["log_std"]),
         ).sum(axis=1)
 
-    @eqx.filter_jit
-    def filter_spec(self):
-        # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        # Specify variational parameters
-        filter_spec = eqx.tree_at(
-            lambda mf: mf.var_params,
-            filter_spec,
-            replace=True,
-        )
-
-        return filter_spec
-
     @eqx.filter_jit
     def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_jit
         def elbo(dyn: Self, n: int, key: Key, data: Any = None) -> Scalar:
@@ -97,7 +99,7 @@ class MeanField(Variational):
 
     @eqx.filter_jit
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_grad
         @eqx.filter_jit

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -33,7 +33,7 @@ class NormalizingFlow(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, eqx.is_array)
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         _, self._unflatten = jfu.ravel_pytree(params)
@@ -41,6 +41,8 @@ class NormalizingFlow(Variational):
         self.base = base
         self.flows = flows
 
+    @property
+    @eqx.filter_jit
     def filter_spec(self):
         # Generate empty specification
         filter_spec = jtu.tree_map(lambda _: False, self)
@@ -49,7 +51,7 @@ class NormalizingFlow(Variational):
         filter_spec = eqx.tree_at(
             lambda vari: vari.flows,
             filter_spec,
-            replace=[flow.filter_spec() for flow in self.flows],
+            replace=[flow.filter_spec for flow in self.flows],
         )
 
         return filter_spec
@@ -112,7 +114,7 @@ class NormalizingFlow(Variational):
 
     @eqx.filter_jit
     def elbo(self, n: int, key: Key = jr.PRNGKey(0), data: Any = None) -> Scalar:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_jit
         def elbo(dyn: Self, n: int, key: Key, data: Any = None):
@@ -129,7 +131,7 @@ class NormalizingFlow(Variational):
 
     @eqx.filter_jit
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_grad
         @eqx.filter_jit

{bayinx-0.2.33 → bayinx-0.3.2}/src/bayinx/mhx/vi/standard.py

@@ -19,7 +19,7 @@ class Standard(Variational):
     - `dim`: Dimension of the parameter space.
     """
 
-    dim: int = eqx.field(static=True)
+    dim: int
     _unflatten: Callable[[Float[Array, "..."]], Model]
     _constraints: Model
 
@@ -31,7 +31,7 @@ class Standard(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec())
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
@@ -54,7 +54,7 @@ class Standard(Variational):
             sigma=jnp.array(1.0),
         ).sum(axis=1, keepdims=True)
 
-    @eqx.filter_jit
+    @property
     def filter_spec(self):
         filter_spec = jtu.tree_map(lambda _: False, self)
 

{bayinx-0.2.33 → bayinx-0.3.2}/tests/test_variational.py

@@ -1,41 +1,36 @@
-from typing import Callable, Dict
+from typing import Dict
 
 import equinox as eqx
 import jax.numpy as jnp
 import pytest
 from jaxtyping import Array
 
-from bayinx import Model
+from bayinx import Model, Parameter
 from bayinx.dists import normal
 from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
 from bayinx.mhx.vi.flows import FullAffine, Planar, Radial
 
 
-# Tests ----
-@pytest.mark.parametrize("var_draws", [1, 10, 100])
-def test_meanfield(benchmark, var_draws):
-    # Construct model definition
-    class NormalDist(Model):
-        params: Dict[str, Array]
-        constraints: Dict[str, Callable[[Array], Array]]
+class NormalDist(Model[Dict[str, Parameter[Array]]]):
+    def __init__(self):
+        self.params = {"mu": Parameter(jnp.array([0.0, 0.0]))}
+        self.constraints = {}
 
-        def __init__(self):
-            self.params = {"mu": jnp.array([0.0, 0.0])}
-            self.constraints = {}
+    @eqx.filter_jit
+    def eval(self, data = None):
+        # Get constrained parameters
+        params, target = self.constrain_params()
 
-        @eqx.filter_jit
-        def eval(self, data: dict):
-            # Get constrained parameters
-            params, target = self.constrain_pars()
+        # Evaluate mu ~ N(10,1)
+        target += normal.logprob(
+            x=params["mu"].vals, mu=jnp.array(10.0), sigma=jnp.array(1.0)
+        ).sum()
 
-            # Evaluate mu ~ N(10,1)
-            target += normal.logprob(
-                x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0)
-            ).sum()
-
-            # Evaluate mu ~ N(10,1)
-            return target
+        return target
 
+# Tests ----
+@pytest.mark.parametrize("var_draws", [1, 10, 100])
+def test_meanfield(benchmark, var_draws):
     # Construct model
     model = NormalDist()
 
@@ -57,28 +52,6 @@ def test_meanfield(benchmark, var_draws):
 
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_affine(benchmark, var_draws):
-    # Construct model definition
-    class NormalDist(Model):
-        params: Dict[str, Array]
-        constraints: Dict[str, Callable[[Array], Array]]
-
-        def __init__(self):
-            self.params = {"mu": jnp.array([0.0, 0.0])}
-            self.constraints = {}
-
-        @eqx.filter_jit
-        def eval(self, data: dict):
-            # Get constrained parameters
-            params, target = self.constrain_pars()
-
-            # Evaluate mu ~ N(10,1)
-            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()
 
@@ -92,7 +65,7 @@ def test_affine(benchmark, var_draws):
     benchmark(benchmark_fit)
     vari = vari.fit(20000, var_draws=var_draws)
 
-    params = vari.flows[0].transform_pars()
+    params = vari.flows[0].transform_params()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
         abs(jnp.eye(2) - params["scale"]) < 0.1
     ).all()
@@ -102,21 +75,18 @@ def test_affine(benchmark, var_draws):
 def test_flows(benchmark, var_draws):
     # Construct model definition
     class NormalDist(Model):
-        params: Dict[str, Array]
-        constraints: Dict[str, Callable[[Array], Array]]
-
         def __init__(self):
-            self.params = {"mu": jnp.array([0.0, 0.0])}
+            self.params = {"mu": Parameter(jnp.array([0.0, 0.0]))}
             self.constraints = {}
 
         @eqx.filter_jit
         def eval(self, data: dict):
             # Get constrained parameters
-            params, target = self.constrain_pars()
+            params, target = self.constrain_params()
 
             # Evaluate mu ~ N(10,1)
             target += normal.logprob(
-                x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0)
+                x=params["mu"].vals, mu=jnp.array(10.0), sigma=jnp.array(1.0)
             ).sum()
 
             return target

bayinx-0.2.33/src/bayinx/__init__.py

@@ -1 +0,0 @@
-from bayinx.core.model import Model as Model

bayinx-0.2.33/src/bayinx/constraints/lower.py

@@ -1,37 +0,0 @@
-from typing import Tuple
-
-import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Scalar, ScalarLike
-
-from bayinx.core.constraint import Constraint
-
-
-class Lower(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 adjusts the posterior density.
-
-        # 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 log-absolute-Jacobian of the transformation.
-        """
-        # Compute transformation adjustment
-        laj: Scalar = jnp.sum(x)
-
-        # Compute transformation
-        x = jnp.exp(x) + self.lb
-
-        return x, laj

bayinx-0.2.33/src/bayinx/core/constraint.py

@@ -1,26 +0,0 @@
-from abc import abstractmethod
-from typing import Tuple
-
-import equinox as eqx
-from jaxtyping import Array, ArrayLike, Scalar
-
-
-class Constraint(eqx.Module):
-    """
-    Abstract base class for defining parameter constraints.
-    """
-
-    @abstractmethod
-    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
-        """
-        Applies the constraining transformation to an unconstrained input and computes the log-absolute-Jacobian of the transformation.
-
-        # Parameters
-        - `x`: The unconstrained JAX Array-like input.
-
-        # Returns
-        A tuple containing:
-            - The constrained JAX Array.
-            - A scalar JAX Array representing the log-absolute-Jacobian of the transformation.
-        """
-        pass