bayinx 0.3.10__py3-none-any.whl → 0.5.3__py3-none-any.whl

bayinx/__init__.py

@@ -1,3 +1,3 @@
-from bayinx.core import Model, Parameter, constrain
-
-__all__ = ["Model", "Parameter", "constrain"]
+from .core.model import Model as Model
+from .core.model import define as define
+from .posterior import Posterior as Posterior

bayinx/constraints/__init__.py

@@ -1,3 +1,4 @@
-from bayinx.constraints.lower import Lower
-
-__all__ = ['Lower']
+from .identity import Identity as Identity
+from .interval import Interval as Interval
+from .lower import Lower as Lower
+from .upper import Upper as Upper

bayinx/constraints/identity.py

@@ -0,0 +1,26 @@
+from typing import Tuple
+
+import jax.numpy as jnp
+from jaxtyping import PyTree, Scalar
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.types import T
+
+
+class Identity(Constraint):
+    """
+    Does nothing.
+    """
+
+    def constrain(self, obj: T, filter_spec: PyTree) -> Tuple[T, Scalar]:
+        """
+        Applies the identity transformation (does nothing) and computes its log-Jacobian adjustment (0).
+
+        # Returns
+        A tuple containing:
+            - The same `PyTree`.
+            - A scalar `Array` containing 0.
+        """
+        log_jac: Scalar = jnp.array(0.0)
+
+        return obj, log_jac

bayinx/constraints/interval.py

@@ -0,0 +1,62 @@
+from typing import Any, Tuple
+
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Scalar, ScalarLike
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.types import T
+
+
+class Interval(Constraint):
+    """
+    Enforces that the parameter lies in the (lb, ub) interval using a scaled
+    and shifted sigmoid transformation.
+
+    # Attributes
+    - `lb`: The lower bound.
+    - `ub`: The upper bound.
+    """
+
+    lb: Scalar
+    ub: Scalar
+
+    def __init__(self, lb: ScalarLike, ub: ScalarLike):
+        self.lb = jnp.asarray(lb)
+        self.ub = jnp.asarray(ub)
+
+    def constrain(self, obj: T, filter_spec: T) -> Tuple[T, Scalar]:
+        """
+        Applies the scaled Sigmoid transformation to the leaves of a `PyTree` and
+        computes the log-Jacobian adjustment.
+
+        # Parameters
+        - `obj`: The unconstrained `PyTree` (values are in R).
+
+        # Returns
+        A tuple containing:
+            - A `PyTree` with its values `y` now satisfying lb < y < ub.
+            - A scalar `Array` containing the log-absolute-Jacobian of the
+              transformation.
+        """
+        log_jac: Scalar = jnp.array(0.0)
+
+        def constrain_leaf(leaf: Any, filter: bool):
+            nonlocal log_jac  # Reference outer variable
+
+            if filter:
+                # Apply transformation
+                constrained = self.lb + (self.ub - self.lb) * jnp.exp(leaf) / (1.0 + jnp.exp(leaf))
+
+                log_jac = log_jac + (jnp.log(constrained - self.lb) +
+                    jnp.log(self.ub - constrained) -
+                    jnp.log(self.ub - self.lb)).sum()
+
+                return constrained
+            else:
+                return leaf
+
+        # Constrain leaves
+        obj = jt.map(constrain_leaf, obj, filter_spec)
+
+        return obj, log_jac

bayinx/constraints/lower.py

@@ -1,50 +1,57 @@
-from typing import Tuple
+from typing import Any, Tuple
 
-import equinox as eqx
 import jax.numpy as jnp
 import jax.tree as jt
-from jaxtyping import PyTree, Scalar, ScalarLike
+from jaxtyping import Scalar, ScalarLike
 
-from bayinx.core import Constraint, Parameter
+from bayinx.core.constraint import Constraint
+from bayinx.core.types import T
 
 
 class Lower(Constraint):
     """
     Enforces a lower bound on the parameter.
+
+    # Attributes
+    - `lb`: The lower bound.
     """
 
     lb: Scalar
 
     def __init__(self, lb: ScalarLike):
-        self.lb = jnp.array(lb)
+        self.lb = jnp.asarray(lb)
 
-    @eqx.filter_jit
-    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+    def constrain(self, obj: T, filter_spec: T) -> Tuple[T, Scalar]:
         """
-        Enforces a lower bound on the parameter and adjusts the posterior density.
+        Applies the exponential transformation to the leaves of a `PyTree` and
+        computes the log-Jacobian adjustment of the transformation.
 
         # Parameters
-        - `x`: The unconstrained `Parameter`.
+        - `x`: The unconstrained `PyTree`.
 
-        # Parameters
+        # Returns
         A tuple containing:
-            - A modified `Parameter` with relevant leaves satisfying the constraint.
-            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+            - A `PyTree` with its values `x` now satisfying `lb <= x`.
+            - A scalar `Array` containing the log-absolute-Jacobian of the
+                transformation.
         """
-        # Extract relevant filter specification
-        filter_spec = x.filter_spec
+        log_jac: Scalar = jnp.array(0.0)
+
+        def constrain_leaf(leaf: Any, filter: bool):
+            nonlocal log_jac  # Reference outer variable
 
-        # Extract relevant parameters(all Array)
-        dyn_params, static_params = eqx.partition(x, filter_spec)
+            if filter:
+                # Apply transformation
+                constrained = jnp.exp(leaf) + self.lb
 
-        # Compute density adjustment
-        laj: PyTree = jt.map(jnp.sum, dyn_params)  # pyright: ignore
-        laj: Scalar = jt.reduce(lambda a, b: a + b, laj)
+                # Accumulate Jacobian adjustment
+                log_jac = log_jac + jnp.sum(leaf)
 
-        # Compute transformation
-        dyn_params = jt.map(lambda v: jnp.exp(v) + self.lb, dyn_params)
+                return constrained
+            else:
+                return leaf
 
-        # Combine into full parameter object
-        x = eqx.combine(dyn_params, static_params)
+        # Constrain leaves
+        obj = jt.map(constrain_leaf, obj, filter_spec)
 
-        return x, laj
+        return obj, log_jac

bayinx/constraints/upper.py

@@ -0,0 +1,57 @@
+from typing import Any, Tuple
+
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Scalar, ScalarLike
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.types import T
+
+
+class Upper(Constraint):
+    """
+    Enforces an upper bound on the parameter.
+
+    # Attributes
+    - `ub`: The upper bound.
+    """
+
+    ub: Scalar
+
+    def __init__(self, ub: ScalarLike):
+        self.ub = jnp.asarray(ub)
+
+    def constrain(self, obj: T, filter_spec: T) -> Tuple[T, Scalar]:
+        """
+        Applies the negated exponential transformation to the leaves of a `PyTree` and
+        computes the log-Jacobian adjustment of the transformation.
+
+        # Parameters
+        - `x`: The unconstrained `PyTree`.
+
+        # Returns
+        A tuple containing:
+            - A `PyTree` with its values `x` now satisfying `x <= ub`.
+            - A scalar `Array` containing the log-absolute-Jacobian of the
+                transformation.
+        """
+        log_jac: Scalar = jnp.array(0.0)
+
+        def constrain_leaf(leaf: Any, include: bool):
+            nonlocal log_jac  # Reference outer variable
+
+            if include:
+                # Apply transformation
+                constrained = -jnp.exp(leaf) + self.ub
+
+                # Accumulate Jacobian adjustment
+                log_jac = log_jac + jnp.sum(leaf)
+
+                return constrained
+            else:
+                return leaf
+
+        # Constrain leaves
+        obj = jt.map(constrain_leaf, obj, filter_spec)
+
+        return obj, log_jac

bayinx/core/__init__.py

@@ -1,7 +0,0 @@
-from ._constraint import Constraint
-from ._flow import Flow
-from ._model import Model, constrain
-from ._parameter import Parameter
-from ._variational import Variational
-
-__all__ = ["Constraint", "Flow", "Model", "constrain", "Parameter", "Variational"]

bayinx/core/constraint.py

@@ -0,0 +1,32 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Tuple
+
+import equinox as eqx
+from jaxtyping import PyTree, Scalar
+
+if TYPE_CHECKING:
+    from bayinx.core.types import T
+
+
+class Constraint(eqx.Module):
+    """
+    Abstract base class for defining constraints (for stochastic nodes).
+    """
+
+    @abstractmethod
+    def constrain(self, obj: "T", filter_spec: PyTree) -> Tuple["T", Scalar]:
+        """
+        Applies the constraining transformation to the leaves of a `PyTree` and
+        computes the log-Jacobian adjustment of the transformation.
+
+        # Parameters
+        - `x`: The unconstrained values.
+        - `filter_spec`: The filter specification for `values`.
+
+        # Returns
+        A tuple containing:
+            - A `PyTree` with its leaves now constrained.
+            - A scalar `Array` containing the log-Jacobian adjustment of the
+                transformation.
+        """
+        pass

bayinx/core/context.py

@@ -0,0 +1,42 @@
+import threading
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from jaxtyping import Scalar
+
+# Local storage for model context
+_model_context = threading.local()
+
+
+@dataclass
+class Target:
+    value: Scalar
+
+    # Define relevant methods to get around lack of explicit mutability in JAX
+    def __iadd__(self, other):
+        self.value = self.value + other
+        return self
+
+    def __add__(self, other):
+        return self.value + other
+
+    def __radd__(self, other):
+        return self.value + other
+
+
+@contextmanager
+def model_context():
+    """
+    Context manager that sets up the implicit `target` accumulator.
+
+    This context allows the `<<` operator for Node classes to automatically
+    accumulate log probabilities without requiring explicit handling of target.
+    """
+    _model_context.target = Target(jnp.array(0.0))
+
+    try:
+        yield _model_context.target
+    finally: # Remove old context if present
+        if hasattr(_model_context, "target"):
+            delattr(_model_context, "target")

bayinx/core/distribution.py

@@ -0,0 +1,34 @@
+from typing import Protocol, Tuple
+
+import jax.random as jr
+from jaxtyping import PRNGKeyArray, Scalar
+
+from bayinx.core.node import Node
+
+
+class Distribution(Protocol):
+    """
+    A protocol used for defining the structure of distributions.
+    """
+
+    def logprob(self, node: Node) -> Scalar: ...
+
+    def sample(self, shape: int | Tuple[int, ...], key: PRNGKeyArray = jr.key(0)): ...
+
+    def __rlshift__(self, node: Node):
+        """
+        Implicitly accumulate the log probability into the current model context.
+        """
+        from bayinx.core.context import _model_context
+
+        # Evaluate log posterior
+        log_prob = self.logprob(node)
+
+        # Accumulate log probability into context
+        if hasattr(_model_context, "target"):
+            _model_context.target += log_prob
+        else:
+            raise RuntimeError(
+                "Model context doesn't exist. Make sure you're calling "
+                + "this within the 'model' method."
+            )

bayinx/core/flow.py

@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Dict, Self, Tuple
+
+import equinox as eqx
+import jax.tree_util as jtu
+from jaxtyping import Array, PyTree
+
+
+class FlowLayer(eqx.Module):
+    """
+    An abstract base class for a flow layer.
+
+    # Attributes
+    - `params`: The parameters of the diffeomorphism. # TODO FOR ALL FLOWS
+    - `constraints`: The constraining transformations for parameters.
+    - `static`: Whether the flow layer is frozen (parameters are not subject to further optimization).
+    """
+
+    params: Dict[str, PyTree]
+    constraints: Dict[str, Callable[[PyTree], Array]]
+    static: bool
+
+    @abstractmethod
+    def forward(self, draws: Array) -> Array:
+        """
+        Computes the forward transformation of `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def adjust(self, draws: Array) -> Array:
+        """
+        Computes the log-Jacobian adjustment for each draw in `draws`.
+
+        # Returns
+            An array of the log-Jacobian adjustments.
+        """
+        pass
+
+    @abstractmethod
+    def forward_and_adjust(self, draws: Array) -> Tuple[Array, Array]:
+        """
+        Computes the log-Jacobian adjustment at `draws` and applies the forward transformation.
+
+        # Returns
+            A tuple of JAX Arrays containing the transformed draws and log-absolute-Jacobians.
+        """
+        pass
+
+    # Default filter specification
+    @property
+    def filter_spec(self) -> "FlowLayer":
+        """
+        Generates a filter specification to subset relevant parameters for the flow.
+        """
+        # Generate empty specification
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        if self.static is False:
+            # Specify parameters
+            filter_spec = eqx.tree_at(
+                lambda flow: flow.params,
+                filter_spec,
+                replace=True,
+            )
+
+        return filter_spec
+
+    def constrain_params(self: Self) -> Dict[str, PyTree]:
+        """
+        Constrain flow parameters to the appropriate domain.
+
+        # Returns
+        The constrained parameters of the diffeomorphism.
+        """
+        params = self.params
+
+        for par, map in self.constraints.items():
+            params[par] = map(params[par])
+
+        return params
+
+    def transform_params(self: Self) -> Dict[str, PyTree]:
+        """
+        Apply a custom transformation to `params` if needed. Defaults to `constrain_params()`.
+
+        # Returns
+        The transformed parameters of the diffeomorphism.
+        """
+        return self.constrain_params()
+
+
+class FlowSpec(ABC):
+    """
+    A specification for a flow layer.
+    """
+
+    @abstractmethod
+    def construct(self, dim: int) -> FlowLayer: ...

bayinx/core/model.py

@@ -0,0 +1,228 @@
+from abc import abstractmethod
+from dataclasses import field, fields
+from typing import (
+    Dict,
+    Optional,
+    Self,
+    Tuple,
+    Type,
+    get_origin,
+    get_type_hints,
+)
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import PyTree, Scalar
+
+from bayinx.constraints import Identity, Interval, Lower, Upper
+from bayinx.core.context import _model_context, model_context
+from bayinx.core.node import Node
+from bayinx.core.types import HasConstraint
+from bayinx.core.utils import _extract_shape_params, _resolve_shape_spec
+from bayinx.nodes import Continuous, Observed, Stochastic
+
+
+def define(
+    shape: Optional[int | str | Tuple[int | str, ...]] = None,
+    init: Optional[PyTree] = None,
+    lower: Optional[float] = None,
+    upper: Optional[float] = None
+):
+    """
+    Define a stochastic node.
+    """
+    metadata: Dict = {}
+
+    if shape is not None:
+        metadata["shape"] = shape
+    if init is not None:
+        metadata["init"] = init
+
+    match (lower, upper):
+        case (float() | int(), None):
+            metadata["constraint"] = Lower(lower) # type: ignore
+        case (None, float() | int()):
+            metadata["constraint"] = Upper(upper) # type: ignore
+        case (float() | int(), float() | int()):
+            metadata["constraint"] = Interval(lower, upper) # type: ignore
+        case (None, None):
+            metadata["constraint"] = Identity()
+        case (_):
+            raise TypeError("TODO.")
+
+    return field(metadata=metadata)
+
+
+class Model(eqx.Module):
+    """
+    A base class used to define probabilistic models.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        # Consume 'init' argument before passing it up to Equinox
+        kwargs.pop('init', None)
+        super().__init_subclass__(**kwargs)
+
+    def __init__(self, **kwargs):
+        cls = self.__class__
+
+        # Grab initialized parameters
+        init_params: set[str] = {f.name for f in fields(cls) if f.name in kwargs.keys()}
+
+        # Grab shape parameters from model definition
+        shape_params: set[str] = set()
+        for node_defn in fields(cls):
+            if (shape_spec := node_defn.metadata.get("shape")) is not None:
+                shape_params = shape_params | _extract_shape_params(shape_spec)
+
+        # Check all shape parameters are passed as arguments
+        if not shape_params.issubset(kwargs.keys()):
+            missing_params = shape_params - kwargs.keys()
+            raise TypeError(
+                f"Following shape parameters were not specified during model initialization: {", ".join(missing_params)}"
+            )
+
+        # Define all initialized dimensions
+        shape_values: dict = {
+            shape_param: kwargs[shape_param]
+            for shape_param in shape_params
+        }
+
+        # Grab node types
+        node_types: dict[str, Type[Node]] = {k: get_origin(v) for k, v in get_type_hints(cls).items()}
+
+
+        # Auto-initialize parameters based on field metadata and type annotations
+        for node_defn in fields(cls):
+            # Grab node type
+            node_type = node_types[node_defn.name]
+
+            # Grab shape information if available
+            shape_spec: str | None = node_defn.metadata.get("shape")
+            shape = _resolve_shape_spec(shape_spec, shape_values)
+
+            # Construct object
+            if node_defn.name in init_params: # Initialized in model construction
+                obj = kwargs[node_defn.name]
+            elif "init" in node_defn.metadata: # Initialized in model definition
+                obj = node_defn.metadata["init"]
+            elif shape is not None: # Shape defined in model definition
+                obj = jnp.zeros(shape) # TODO: will change later for discrete objects
+            else:
+                raise ValueError(f"Node '{node_defn.name}' not initialized or defined.")
+
+            # Check shape
+            if shape is not None and jnp.shape(obj) != shape:
+                raise ValueError(f"Expected shape {shape} for {node_defn.name} but got {jnp.shape(obj)}.")
+
+            if issubclass(node_type, Stochastic):
+                if node_type == Continuous:
+                    setattr(
+                        self,
+                        node_defn.name,
+                        Continuous(obj, node_defn.metadata["constraint"]),
+                    )
+                else:
+                    TypeError(f"{node_type.__name__} is not implemented yet")
+            elif issubclass(node_type, Observed):
+                setattr(self, node_defn.name, Observed(obj))
+            else:
+                raise TypeError(f"{node_defn.name} node is neither Stochastic nor Observed but {node_type.__name__}.")
+
+
+    @property
+    def filter_spec(self) -> Self:
+        """
+        Generates a filter specification to subset stochastic elements of the model.
+        """
+        # Generate empty specification
+        filter_spec: Self = jt.map(lambda _: False, self)
+
+        for f in fields(self): # type: ignore
+            # Extract attribute
+            node: Node = getattr(self, f.name)
+
+            # Check if attribute is stochastic
+            if isinstance(node, Stochastic):
+                # Update model's filter specification at node
+                filter_spec: Self = eqx.tree_at(
+                    lambda model: getattr(model, f.name),
+                    filter_spec,
+                    replace=node.filter_spec
+                )
+
+        return filter_spec
+
+    def filter_for(self, node_type: Type[Stochastic]) -> Self:
+        """
+        Generates a filter specification to subset stochastic elements of a certain type of the model.
+        """
+        # Generate empty specification
+        filter_spec: Self = jt.map(lambda _: False, self)
+
+        for f in fields(self): # type: ignore
+            # Extract node
+            node: Node = getattr(self, f.name)
+
+            if isinstance(node, node_type):
+                # Update model's filter specification for node
+                filter_spec: Self = eqx.tree_at(
+                    lambda model: getattr(model, f.name),
+                    filter_spec,
+                    replace=node.filter_spec,
+                )
+
+        return filter_spec
+
+
+    def constrain(self, jacobian: bool = True) -> Tuple[Self, Scalar]:
+        """
+        Constrain nodes to the appropriate domain.
+
+        # Returns
+        A tuple containing the constrained `Model` object and the log-Jacobian adjustment.
+        """
+        model: Self = self
+        target: Scalar = jnp.array(0.0)
+
+        for f in fields(self): # type: ignore
+            # Extract attribute
+            node = getattr(self, f.name)
+
+            # Check if node has a constraint
+            if isinstance(node, HasConstraint):
+                # Apply constraint
+                obj, laj = node._constraint.constrain(node.obj, node._filter_spec)
+
+                # Update values with constrained counterpart
+                model = eqx.tree_at(
+                    where=lambda model: getattr(model, f.name).obj,
+                    pytree=model,
+                    replace=obj,
+                )
+
+                # Adjust posterior density
+                if jacobian:
+                    target += laj
+
+        return model, target
+
+    @abstractmethod
+    def model(self, target: Scalar) -> Scalar:
+        pass
+
+    @eqx.filter_jit
+    def __call__(self) -> Scalar:
+        with model_context():
+            # Constrain the model and accumulate Jacobian adjustments
+            self, target = self.constrain()
+
+            # Accumulate manual increments
+            target += self.model(jnp.array(0.0))
+
+            # Accumulate implicit increments
+            target += _model_context.target.value
+
+            # Return the accumulated target
+            return target