bayinx 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

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/constraints/__init__.py

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

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

@@ -1,26 +1,29 @@
 from abc import abstractmethod
-from typing import Any, Dict, Generic, Tuple, TypeVar
+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 PyTree, Scalar
+from jaxtyping import Scalar
 
 from bayinx.core.constraint import Constraint
 from bayinx.core.parameter import Parameter
 
-T = TypeVar('T', bound=PyTree)
-class Model(eqx.Module, Generic[T]):
+
+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.
 
-    # Attributes
-    - `params`: A dictionary of parameters.
-    - `constraints`: A dictionary of constraints.
-    """
+    Annotate parameter attributes with `Parameter`.
 
-    params: Dict[str, Parameter[T]]
-    constraints: Dict[str, Constraint]
+    Include constraints by setting them equal to `constrain(Constraint)`.
+    """
 
     @abstractmethod
     def eval(self, data: Any) -> Scalar:
@@ -29,50 +32,71 @@ class Model(eqx.Module, Generic[T]):
     # Default filter specification
     @property
     @eqx.filter_jit
-    def filter_spec(self):
+    def filter_spec(self) -> Self:
         """
         Generates a filter specification to subset relevant parameters for the model.
         """
         # Generate empty specification
-        filter_spec = jt.map(lambda _: False, self)
+        filter_spec: Self = jt.map(lambda _: False, self)
+
+        for f in fields(self):
+            # Extract attribute from field
+            attr = getattr(self, f.name)
 
-        # 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()}
-        )
+            # 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
 
-    # Add constrain method
+
     @eqx.filter_jit
-    def constrain_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
+    def constrain_params(self) -> Tuple[Self, Scalar]:
         """
-        Constrain `params` to the appropriate domain.
+        Constrain parameters to the appropriate domain.
 
         # Returns
-        A dictionary of PyTrees representing the constrained parameters and the adjustment to the posterior density.
+        A constrained `Model` object and the adjustment to the posterior.
         """
-        t_params: Dict[str, Parameter[T]] = self.params
+        constrained: Self = self
         target: Scalar = jnp.array(0.0)
 
-        for par, map in self.constraints.items():
-            # Apply transformation
-            t_params[par], ladj = map.constrain(t_params[par])
+        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
 
-            # Adjust posterior density
-            target -= ladj
+        return constrained, target
 
-        return t_params, target
 
-    # Add default transform method
     @eqx.filter_jit
-    def transform_params(self) -> Tuple[Dict[str, Parameter[T]], Scalar]:
+    def transform_params(self) -> Tuple[Self, Scalar]:
         """
-        Apply a custom transformation to `params` if needed.
+        Apply a custom transformation to parameters if needed(defaults to constrained parameters).
 
         # Returns
-        A dictionary of transformed JAX Arrays representing the transformed parameters.
+        A transformed `Model` object and the adjustment to the posterior.
         """
         return self.constrain_params()

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

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

@@ -1,12 +1,12 @@
-bayinx/__init__.py,sha256=htihTsJ54k-ljBLzt4ye8DR7ORwZhxv-bLMcEaDQeqY,86
+bayinx/__init__.py,sha256=5fb_tGeEVnrNt6IQqu7gZaJskBJHqjcg08JRPrY2ANo,139
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/constraints/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/constraints/__init__.py,sha256=PSxvcuSox2JL61AG1iag2PTNKPcid_DbOQzHpYdj5RE,52
 bayinx/constraints/lower.py,sha256=wkYnWjaAEGQeXKfBo_gY0pcK9ElJUMkzGdAmWI8ykCk,1488
 bayinx/core/__init__.py,sha256=jSwEFdXqi-Bj_X8_H-YuaXp5ebEQpZTG2T18zpquzPo,207
 bayinx/core/constraint.py,sha256=F6-TXQjzt-tcNm8bHkRcGEtyE9bZQf2RbAh_MKDuM20,760
-bayinx/core/flow.py,sha256=lAPJdQnrIxC3JoowTp77Gvm0p0v_xQn8FMjFJWMnWbc,2340
-bayinx/core/model.py,sha256=RKAtsLXc6xDnXWz5upmx6Vz6JOoorw4WTfxTA7B7Lmg,2294
-bayinx/core/parameter.py,sha256=oxCCZcZ-DDBvfWzexfhQkSJPxNQnE1vYXtBhiEZG2bM,1025
+bayinx/core/flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
+bayinx/core/model.py,sha256=1vQPVjE0ebCdW7mLuabgQcCTi95o8n8CC6GuzJdNL1s,2956
+bayinx/core/parameter.py,sha256=eECqvfMNWSU8_CkGYaAfOCneMMQGZI21kF0mErsh2Rc,1080
 bayinx/core/variational.py,sha256=lqENISRrKY8ODLtl0D-D7TAA2gD7HGh37BnROM7p5hI,4783
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
@@ -25,6 +25,6 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvAD
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.1.dist-info/METADATA,sha256=Slp2nxR8HISCwCqIXY2El3GgqsO1v9_UbVeJq726w7k,3057
-bayinx-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.1.dist-info/RECORD,,
+bayinx-0.3.3.dist-info/METADATA,sha256=5tZVPxDvYVajnoDrRIF-KskSMWBix7Zq4h_glyJa-_M,3057
+bayinx-0.3.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.3.dist-info/RECORD,,