bayinx 0.3.1__tar.gz → 0.3.3__tar.gz

{bayinx-0.3.1 → bayinx-0.3.3}/PKG-INFO

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

{bayinx-0.3.1 → bayinx-0.3.3}/pyproject.toml

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

{bayinx-0.3.1 → bayinx-0.3.3}/src/bayinx/__init__.py

@@ -1,2 +1,3 @@
 from bayinx.core import Model as Model
 from bayinx.core import Parameter as Parameter
+from bayinx.core.model import constrain as constrain

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

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

{bayinx-0.3.1 → bayinx-0.3.3}/src/bayinx/core/flow.py

@@ -11,7 +11,7 @@ class Flow(eqx.Module):
     An abstract base class for a flow(of a normalizing flow).
 
     # Attributes
-    - `pars`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
+    - `params`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
     - `constraints`: A dictionary of simple functions that constrain their corresponding parameter.
     """
 

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

@@ -0,0 +1,102 @@
+from abc import abstractmethod
+from dataclasses import field, fields
+from typing import Any, Self, Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Scalar
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.parameter import Parameter
+
+
+def constrain(constraint: Constraint):
+    """Defines constraint metadata."""
+    return field(metadata={'constraint': constraint})
+
+
+class Model(eqx.Module):
+    """
+    An abstract base class used to define probabilistic models.
+
+    Annotate parameter attributes with `Parameter`.
+
+    Include constraints by setting them equal to `constrain(Constraint)`.
+    """
+
+    @abstractmethod
+    def eval(self, data: Any) -> Scalar:
+        pass
+
+    # Default filter specification
+    @property
+    @eqx.filter_jit
+    def filter_spec(self) -> Self:
+        """
+        Generates a filter specification to subset relevant parameters for the model.
+        """
+        # Generate empty specification
+        filter_spec: Self = jt.map(lambda _: False, self)
+
+        for f in fields(self):
+            # Extract attribute from field
+            attr = getattr(self, f.name)
+
+            # Check if attribute is a parameter
+            if isinstance(attr, Parameter):
+                # Update filter specification for parameter
+                filter_spec = eqx.tree_at(
+                    lambda model: getattr(model, f.name),
+                    filter_spec,
+                    replace=attr.filter_spec
+                )
+
+        return filter_spec
+
+
+    @eqx.filter_jit
+    def constrain_params(self) -> Tuple[Self, Scalar]:
+        """
+        Constrain parameters to the appropriate domain.
+
+        # Returns
+        A constrained `Model` object and the adjustment to the posterior.
+        """
+        constrained: Self = self
+        target: Scalar = jnp.array(0.0)
+
+        for f in fields(self):
+            # Extract attribute
+            attr = getattr(self, f.name)
+
+            # Check if constrained parameter
+            if isinstance(attr, Parameter) and 'constraint' in f.metadata:
+                param = attr
+                constraint = f.metadata['constraint']
+
+                # Apply constraint
+                param, laj = constraint.constrain(param)
+
+                # Update parameters for constrained model
+                constrained = eqx.tree_at(
+                    lambda model: getattr(model, f.name),
+                    constrained,
+                    replace=param
+                )
+
+                # Adjust posterior density
+                target += laj
+
+        return constrained, target
+
+
+    @eqx.filter_jit
+    def transform_params(self) -> Tuple[Self, Scalar]:
+        """
+        Apply a custom transformation to parameters if needed(defaults to constrained parameters).
+
+        # Returns
+        A transformed `Model` object and the adjustment to the posterior.
+        """
+        return self.constrain_params()

{bayinx-0.3.1 → bayinx-0.3.3}/src/bayinx/core/parameter.py

@@ -21,6 +21,9 @@ class Parameter(eqx.Module, Generic[T]):
         # Insert parameter values
         self.vals = values
 
+    def __call__(self) -> T:
+        return self.vals
+
     # Default filter specification
     @property
     @eqx.filter_jit

{bayinx-0.3.1 → bayinx-0.3.3}/tests/test_variational.py

@@ -1,7 +1,5 @@
-
 from typing import Dict
 
-import equinox as eqx
 import jax.numpy as jnp
 import pytest
 from jaxtyping import Array
@@ -12,30 +10,25 @@ 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[Array]):
-        params: Dict[str, Parameter[Array]]
+class NormalDist(Model):
+    x: Parameter[Array]
 
-        def __init__(self):
-            self.params = {"mu": Parameter(jnp.array([0.0, 0.0]))}
-            self.constraints = {}
+    def __init__(self):
+        self.x = Parameter(jnp.array([0.0, 0.0]))
 
-        @eqx.filter_jit
-        def eval(self, data = None):
-            # Get constrained parameters
-            params, target = self.constrain_params()
+    def eval(self, data: Dict[str, Array]):
+        # Constrain parameters
+        self, target = self.constrain_params()
 
-            # Evaluate mu ~ N(10,1)
-            target += normal.logprob(
-                x=params["mu"].vals, mu=jnp.array(10.0), sigma=jnp.array(1.0)
-            ).sum()
+        # Evaluate x ~ Normal(10.0, 1.0)
+        target += jnp.sum(normal.logprob(self.x(), jnp.array(10.0), jnp.array(1.0)))
 
-            # Evaluate mu ~ N(10,1)
-            return target
+        return target
 
+
+# Tests ----
+@pytest.mark.parametrize("var_draws", [1, 100])
+def test_meanfield(benchmark, var_draws):
     # Construct model
     model = NormalDist()
 
@@ -46,8 +39,8 @@ def test_meanfield(benchmark, var_draws):
     def benchmark_fit():
         vari.fit(10000, var_draws=var_draws)
 
-    benchmark(benchmark_fit)
     vari = vari.fit(20000, var_draws=var_draws)
+    benchmark(benchmark_fit)
 
     # Assert parameters are roughly correct
     assert all(abs(10.0 - vari.var_params["mean"]) < 0.1) and all(
@@ -55,28 +48,8 @@ def test_meanfield(benchmark, var_draws):
     )
 
 
-@pytest.mark.parametrize("var_draws", [1, 10, 100])
+@pytest.mark.parametrize("var_draws", [1, 100])
 def test_affine(benchmark, var_draws):
-    # Construct model definition
-    class NormalDist(Model):
-
-        def __init__(self):
-            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_params()
-
-            # 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)
-            return target
-
     # Construct model
     model = NormalDist()
 
@@ -87,8 +60,8 @@ def test_affine(benchmark, var_draws):
     def benchmark_fit():
         vari.fit(10000, var_draws=var_draws)
 
-    benchmark(benchmark_fit)
     vari = vari.fit(20000, var_draws=var_draws)
+    benchmark(benchmark_fit)
 
     params = vari.flows[0].transform_params()
     assert (abs(10.0 - vari.flows[0].params["shift"]) < 0.1).all() and (
@@ -96,26 +69,8 @@ def test_affine(benchmark, var_draws):
     ).all()
 
 
-@pytest.mark.parametrize("var_draws", [1, 10, 100])
+@pytest.mark.parametrize("var_draws", [1, 100])
 def test_flows(benchmark, var_draws):
-    # Construct model definition
-    class NormalDist(Model):
-        def __init__(self):
-            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_params()
-
-            # Evaluate mu ~ N(10,1)
-            target += normal.logprob(
-                x=params["mu"].vals, mu=jnp.array(10.0), sigma=jnp.array(1.0)
-            ).sum()
-
-            return target
-
     # Construct model
     model = NormalDist()
 
@@ -128,8 +83,9 @@ def test_flows(benchmark, var_draws):
     def benchmark_fit():
         vari.fit(10000, var_draws=var_draws)
 
-    benchmark(benchmark_fit)
     vari = vari.fit(20000, var_draws=var_draws)
+    benchmark(benchmark_fit)
+
 
     mean = vari.sample(1000).mean(0)
     var = vari.sample(1000).var(0)

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

@@ -1,78 +0,0 @@
-from abc import abstractmethod
-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 PyTree, Scalar
-
-from bayinx.core.constraint import Constraint
-from bayinx.core.parameter import Parameter
-
-T = TypeVar('T', bound=PyTree)
-class Model(eqx.Module, Generic[T]):
-    """
-    An abstract base class used to define probabilistic models.
-
-    # Attributes
-    - `params`: A dictionary of parameters.
-    - `constraints`: A dictionary of constraints.
-    """
-
-    params: Dict[str, Parameter[T]]
-    constraints: Dict[str, Constraint]
-
-    @abstractmethod
-    def eval(self, data: Any) -> Scalar:
-        pass
-
-    # Default filter specification
-    @property
-    @eqx.filter_jit
-    def filter_spec(self):
-        """
-        Generates a filter specification to subset relevant parameters for the model.
-        """
-        # Generate empty specification
-        filter_spec = jt.map(lambda _: False, self)
-
-        # Specify relevant parameters
-        filter_spec = eqx.tree_at(
-            lambda model: model.params,
-            filter_spec,
-            replace={key: param.filter_spec for key, param in self.params.items()}
-        )
-
-        return filter_spec
-
-    # Add constrain method
-    @eqx.filter_jit
-    def constrain_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
-        """
-        Constrain `params` to the appropriate domain.
-
-        # Returns
-        A dictionary of PyTrees representing the constrained parameters and the adjustment to the posterior density.
-        """
-        t_params: Dict[str, Parameter[T]] = 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
-
-    # Add default transform method
-    @eqx.filter_jit
-    def transform_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
-        """
-        Apply a custom transformation to `params` if needed.
-
-        # Returns
-        A dictionary of transformed JAX Arrays representing the transformed parameters.
-        """
-        return self.constrain_params()