bayinx 0.3.2__tar.gz → 0.3.4__tar.gz

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

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

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

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

{bayinx-0.3.2 → bayinx-0.3.4}/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.2 → bayinx-0.3.4}/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.4/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.2 → bayinx-0.3.4}/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.4/src/bayinx/dists/censored/gamma2/__init__.py

@@ -0,0 +1 @@
+from . import r as r

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

@@ -1,6 +1,5 @@
 from typing import Dict
 
-import equinox as eqx
 import jax.numpy as jnp
 import pytest
 from jaxtyping import Array
@@ -11,25 +10,24 @@ from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
 from bayinx.mhx.vi.flows import FullAffine, Planar, Radial
 
 
-class NormalDist(Model[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 = {}
+        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)))
 
         return target
 
+
 # Tests ----
-@pytest.mark.parametrize("var_draws", [1, 10, 100])
+@pytest.mark.parametrize("var_draws", [1, 100])
 def test_meanfield(benchmark, var_draws):
     # Construct model
     model = NormalDist()
@@ -41,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(
@@ -50,7 +48,7 @@ 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
     model = NormalDist()
@@ -62,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 (
@@ -71,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()
 
@@ -103,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.2 → bayinx-0.3.4}/uv.lock

@@ -17,7 +17,7 @@ wheels = [
 
 [[package]]
 name = "bayinx"
-version = "0.2.11"
+version = "0.3.3"
 source = { editable = "." }
 dependencies = [
     { name = "equinox" },

bayinx-0.3.2/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
-
-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 parameters.
-    - `constraints`: A dictionary of constraints.
-    """
-
-    params: P
-    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[P, 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: P = 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[P, Scalar]:
-        """
-        Apply a custom transformation to `params` if needed.
-
-        # Returns
-        A dictionary of transformed JAX Arrays representing the transformed parameters.
-        """
-        return self.constrain_params()