bayinx 0.3.1__tar.gz → 0.3.2__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.3.1"
+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.3.1"
+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/constraints/__init__.py

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

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

@@ -9,8 +9,8 @@ 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]):
+P = TypeVar('P', bound=Dict[str, Parameter[PyTree]])
+class Model(eqx.Module, Generic[P]):
     """
     An abstract base class used to define probabilistic models.
 
@@ -19,7 +19,7 @@ class Model(eqx.Module, Generic[T]):
     - `constraints`: A dictionary of constraints.
     """
 
-    params: Dict[str, Parameter[T]]
+    params: P
     constraints: Dict[str, Constraint]
 
     @abstractmethod
@@ -47,14 +47,14 @@ class Model(eqx.Module, Generic[T]):
 
     # Add constrain method
     @eqx.filter_jit
-    def constrain_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
+    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: Dict[str, Parameter[T]] = self.params
+        t_params: P = self.params
         target: Scalar = jnp.array(0.0)
 
         for par, map in self.constraints.items():
@@ -68,7 +68,7 @@ class Model(eqx.Module, Generic[T]):
 
     # Add default transform method
     @eqx.filter_jit
-    def transform_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
+    def transform_params(self) -> Tuple[P, Scalar]:
         """
         Apply a custom transformation to `params` if needed.
 

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

@@ -1,4 +1,3 @@
-
 from typing import Dict
 
 import equinox as eqx
@@ -12,30 +11,26 @@ 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[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": Parameter(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 = None):
-            # 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)
-            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
+        return target
 
+# Tests ----
+@pytest.mark.parametrize("var_draws", [1, 10, 100])
+def test_meanfield(benchmark, var_draws):
     # Construct model
     model = NormalDist()
 
@@ -57,26 +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):
-
-        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()