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

bayinx/core/node.py

@@ -0,0 +1,201 @@
+from typing import Any, Generic, Iterator, Self
+
+import equinox as eqx
+import jax.tree as jt
+from jaxtyping import PyTree
+
+from bayinx.core.types import T
+from bayinx.core.utils import _extract_obj, _merge_filter_specs
+
+
+class Node(eqx.Module, Generic[T]):
+    """
+    A thin wrapper for nodes of a probabilistic model.
+
+
+    # Attributes
+    - `obj`: An internal realization of the node.
+    - `_filter_spec`: An internal filter specification for `obj`.
+    """
+
+    obj: T
+    _filter_spec: PyTree
+
+    def __init__(self, obj, filter_spec):
+        self.obj = obj
+        self._filter_spec = filter_spec
+
+    @property
+    def filter_spec(self) -> Self:
+        """
+        An outer filter specification for the full node.
+        """
+        # Generate empty specification
+        node_filter_spec: Self = jt.map(lambda _: False, self)
+
+        # Filter based on inner filter specification for 'obj'
+        node_filter_spec = eqx.tree_at(
+            lambda node: node.obj,
+            node_filter_spec,
+            replace=self._filter_spec,
+        )
+
+        return node_filter_spec
+
+    # Wrappers around internal dunder methods ----
+    def __getitem__(self, key: Any) -> "Node":
+        if isinstance(key, Node):
+            raise TypeError("Subsetting nodes with nodes is not yet supported.")
+
+        # Subset internally
+        new_obj = self.obj[key]
+        if type(self.obj) is type(self._filter_spec):
+            new_filter_spec = self._filter_spec[key]
+        else:
+            new_filter_spec = self._filter_spec
+
+        # Create new subsetted node
+        return type(self)(new_obj, new_filter_spec)
+
+    def __iter__(self) -> Iterator["Node"]:
+        for obj_i, spec_i in zip(self.obj, self._filter_spec):
+            # Create a new Node for the current element
+            yield Node(obj_i, spec_i)
+
+
+    ## Arithmetic ----
+    def __add__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform addition
+        new_obj = lhs_obj + rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __sub__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform subtraction
+        new_obj = lhs_obj - rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __mul__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform multiplication
+        new_obj = lhs_obj * rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __matmul__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform matrix multiplication
+        new_obj = lhs_obj @ rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __truediv__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform true division
+        new_obj = lhs_obj / rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __floordiv__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform floor division
+        new_obj = lhs_obj // rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __pow__(self, other: Any) -> "Node":
+         # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform floor division
+        new_obj = lhs_obj ** rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)
+
+    def __mod__(self, other: Any) -> "Node":
+        # Extract internal objects and their filter specifications
+        lhs_obj, lhs_filter_spec = _extract_obj(self)
+        rhs_obj, rhs_filter_spec = _extract_obj(other)
+
+        # Perform modulus
+        new_obj = lhs_obj % rhs_obj
+
+        # Merge filter specifications
+        new_filter_spec = _merge_filter_specs(
+            [lhs_filter_spec, rhs_filter_spec],
+            [lhs_obj, rhs_obj],
+            new_obj
+        )
+
+        return Node(new_obj, new_filter_spec)

bayinx/core/types.py

@@ -0,0 +1,17 @@
+from typing import Generic, Protocol, TypeVar, runtime_checkable
+
+from jaxtyping import PyTree
+
+from bayinx.core.constraint import Constraint
+
+T = TypeVar("T", bound=PyTree)
+
+@runtime_checkable
+class HasConstraint(Protocol, Generic[T]):
+    """
+    Protocol for probabilistic nodes that have constraints.
+    """
+
+    obj: T
+    _filter_spec: PyTree
+    _constraint: Constraint

bayinx/core/utils.py

@@ -0,0 +1,109 @@
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import jax.tree as jt
+from jaxtyping import PyTree
+
+
+def _extract_shape_params(shape_spec: int | str | Tuple[int | str, ...]) -> set[str]:
+    """
+    Extract parameter names from shape specification.
+    """
+    params = set()
+    if isinstance(shape_spec, str):
+        params.add(shape_spec)
+    elif isinstance(shape_spec, tuple):
+        for item in shape_spec:
+            if isinstance(item, str):
+                params.add(item)
+    #
+    return params
+
+def _resolve_shape_spec(
+    shape_spec: None | int | str | Tuple[int | str, ...],
+    shape_values: Dict[str, int]
+) -> None | Tuple[int, ...]:
+    """
+    Replaces named dimensions in a shape specification with their integer or tuple values.
+
+    # Example
+    For `shape_values = {'k': 5, 's': (3,2,1)}`:
+        `('k', 4, 's')` --> `(5, 4, 3, 2, 1)`
+    """
+    if shape_spec is None:
+        return None
+
+    # Coerce to tuple for uniform processing
+    if isinstance(shape_spec, (int, str)):
+        shape_spec = (shape_spec,)
+
+    resolved_spec: List[int] = []
+    for dim in shape_spec:
+        if isinstance(dim, str):
+            if dim in shape_values:
+                resolved_value = shape_values[dim] # Grab initialized value
+
+                if isinstance(resolved_value, int):
+                    # Scalar shape dimension (e.g., 'k' -> 3)
+                    resolved_spec.append(resolved_value)
+                elif isinstance(resolved_value, tuple):
+                    # Packed shape dimension (e.g., 'shape' -> (3, 2, 1))
+                    resolved_spec.extend(resolved_value)
+                else:
+                    raise TypeError(f"Shape parameter '{dim}' resolved to an unsupported type: {type(resolved_value).__name__}")
+            else: # dim not in shape_values
+                raise TypeError(f"Shape parameter '{dim}' was not initialized with a value.")
+        elif isinstance(dim, int):
+            # Literal integer (e.g., 3 -> 3)
+            resolved_spec.append(dim)
+        else:
+            raise TypeError(f"Shape parameter {dim} was incorrectly specified (must be 'int' or 'str', got '{type(dim).__name__}').")
+
+    return tuple(resolved_spec)
+
+def _extract_obj(x: Any) -> Tuple[Any, Any]:
+    """
+    Extract the object and its (potentially implicit) filter specification.
+    """
+    from bayinx.core.node import Node
+
+    if isinstance(x, Node):
+        obj = x.obj
+        filter_spec = x._filter_spec
+    else:
+        obj: Any = x # type: ignore
+        filter_spec = True # implicit filter specification
+
+    return (obj, filter_spec)
+
+
+def _merge_filter_specs(
+    filter_specs: List[PyTree],
+    objs: Optional[List[PyTree]] = None,
+    obj: Optional[PyTree] = None
+) -> PyTree:
+    """
+    Merge filter specifications that share the type of `obj`.
+
+    If `obj` and `objs` are not provided then `filter_specs` will be merged as is.
+    """
+
+    def _merge(*args):
+        return all(args)
+
+    if objs is None and obj is None:
+        filter_spec: Any = jt.map(_merge, *filter_specs) # type: ignore
+    else:
+        selected_specs: List[PyTree] = []
+
+        # Include filter specs whose objects share the correct type
+        for cur_obj, cur_spec in zip(objs, filter_specs): # type: ignore
+            if type(obj) is type(cur_obj):
+                selected_specs.append(cur_spec)
+
+        if len(selected_specs) != 0:
+            filter_spec = jt.map(_merge, *selected_specs)
+        else:
+            filter_spec = True
+
+    return filter_spec

bayinx/core/variational.py

@@ -0,0 +1,170 @@
+from abc import abstractmethod
+from functools import partial
+from typing import Callable, Generic, Self, Tuple, TypeVar
+
+import equinox as eqx
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import optax as opx
+from jaxtyping import Array, PRNGKeyArray, PyTree, Scalar
+from optax import GradientTransformation, OptState
+
+from bayinx.core.model import Model
+
+M = TypeVar("M", bound=Model)
+
+class Variational(eqx.Module, Generic[M]):
+    """
+    An abstract base class used to define variational methods.
+
+    # Attributes
+    - `dim`: The dimension of the support.
+    - `_unflatten`: A function to transform draws from the variational distribution back to a `Model`.
+    - `_static`: The static component of a partitioned `Model` used to initialize the `Variational` object.
+    """
+    dim: int
+    _unflatten: Callable[[Array], M]
+    _static: M
+
+    @property
+    @abstractmethod
+    def filter_spec(self) -> Self:
+        """
+        Filter specification for dynamic and static components of the
+        `Variational` object.
+        """
+        pass
+
+    @abstractmethod
+    def sample(self, n: int, key: PRNGKeyArray = jr.PRNGKey(0)) -> Array:
+        """
+        Sample from the variational distribution.
+        """
+        pass
+
+    @abstractmethod
+    def eval(self, draws: Array) -> Array:
+        """
+        Evaluate the variational distribution at `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def elbo(self, n: int, batch_size: int, key: PRNGKeyArray) -> Array:
+        """
+        Evaluate the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def elbo_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> M:
+        """
+        Evaluate the gradient of the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def elbo_and_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Tuple[Scalar, PyTree]:
+        """
+        Evaluate the ELBO and its gradient.
+        """
+        pass
+
+    @eqx.filter_jit
+    def reconstruct_model(self, draw: Array) -> M:
+        # Unflatten variational draw
+        model: M = self._unflatten(draw)
+
+        # Combine with static components
+        model: M = eqx.combine(model, self._static)
+
+        return model
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def eval_model(self, draws: Array) -> Array:
+        """
+        Reconstruct models from variational draws and evaluate their posterior.
+
+        # Parameters
+        - `draws`: A set of variational draws.
+        """
+        # Unflatten variational draw
+        model: M = self.reconstruct_model(draws)
+
+        # Evaluate posterior
+        return model()
+
+    @eqx.filter_jit
+    def fit(
+        self,
+        max_iters: int,
+        learning_rate: float,
+        tolerance: float,
+        grad_draws: int,
+        batch_size: int,
+        key: PRNGKeyArray = jr.key(0),
+    ) -> Self:
+        """
+        Optimize the variational distribution.
+        """
+        # Partition variational
+        dyn, static = eqx.partition(self, self.filter_spec)
+
+        # Construct scheduler
+        schedule = opx.warmup_cosine_decay_schedule(
+            1e-8,
+            learning_rate,
+            int(max_iters / 10),
+            max_iters - int(max_iters / 10)
+        )
+
+        # Initialize optimizer
+        optim: GradientTransformation = opx.chain(
+            opx.scale(-1.0), opx.adam(schedule) # replace learning_rate with scheduler
+        )
+        opt_state: OptState = optim.init(dyn)
+
+        def condition(state: Tuple[Self, OptState, Scalar, PRNGKeyArray]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
+
+            return i < max_iters
+
+        def body(state: Tuple[Self, OptState, Scalar, PRNGKeyArray]):
+            # Unpack iteration state
+            dyn, opt_state, i, key = state
+
+            # Update iteration
+            i = i + 1
+
+            # Update PRNG key
+            key, _ = jr.split(key)
+
+            # Reconstruct variational
+            vari: Self = eqx.combine(dyn, static)
+
+            # Compute gradient of the ELBO for update
+            update: M = vari.elbo_grad(grad_draws, batch_size, key)
+
+            # Transform update through optimizer
+            update, opt_state = optim.update( # type: ignore
+                update, opt_state, eqx.filter(dyn, dyn.filter_spec) # type: ignore
+            )
+
+            # Update variational distribution
+            dyn: Self = eqx.apply_updates(dyn, update)
+
+            return dyn, opt_state, i, key
+
+        # Run optimization loop
+        dyn = lax.while_loop(
+            cond_fun=condition,
+            body_fun=body,
+            init_val=(dyn, opt_state, jnp.array(0, jnp.uint32), key),
+        )[0]
+
+        # Return optimized variational
+        return eqx.combine(dyn, static)

bayinx/dists/__init__.py

@@ -1,3 +1,5 @@
-from bayinx.dists import censored, gamma2, normal, posnormal
-
-__all__ = ['censored', "gamma2", "normal", "posnormal"]
+from .bernoulli import Bernoulli as Bernoulli
+from .binomial import Binomial as Binomial
+from .exponential import Exponential as Exponential
+from .normal import Normal as Normal
+from .poisson import Poisson as Poisson