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

bayinx/nodes/continuous.py

@@ -0,0 +1,64 @@
+from typing import Any, Optional
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree as jt
+import numpy as np
+from jaxtyping import Array, PyTree
+
+from bayinx.constraints import Identity
+from bayinx.core.constraint import Constraint
+from bayinx.core.types import T
+from bayinx.nodes.stochastic import Stochastic
+
+
+def is_float_like(element: Any) -> bool:
+    """
+    Check if `element` is float-like.
+
+    The structure of this function is borrowed from the `Equinox` library.
+    """
+    if hasattr(element, "__jax_array__"):
+        element = element.__jax_array__()
+    if isinstance(element, (np.ndarray, np.generic)):
+        return bool(np.issubdtype(element.dtype, np.floating))
+    elif isinstance(element, Array):
+        return jnp.issubdtype(element.dtype, jnp.floating)
+    else:
+        return isinstance(element, float)
+
+
+class Continuous(Stochastic[T]):
+    """
+    A container for continuous stochastic nodes of a probabilistic model.
+
+
+    # Attributes
+    - `obj`: An internal realization of the node.
+    - `_filter_spec`: An internal filter specification for `obj`.
+    - `_constraint`: A constraining transformation.
+    """
+
+    _constraint: Constraint
+
+
+    def __init__(
+        self,
+        obj: T,
+        constraint: Constraint = Identity(),
+        filter_spec: Optional[PyTree] = None
+    ):
+        if filter_spec is None: # Default filter specification
+            # Generate empty specification
+            filter_spec = jt.map(lambda _: False, obj)
+
+            # Specify float-like leaves
+            filter_spec = eqx.tree_at(
+                where=lambda obj: obj,
+                pytree=filter_spec,
+                replace=jt.map(is_float_like, obj),
+            )
+
+        self.obj = obj
+        self._filter_spec = filter_spec
+        self._constraint = constraint

bayinx/nodes/observed.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+import equinox as eqx
+import jax.tree as jt
+from jaxtyping import PyTree
+
+from bayinx.core.node import Node
+from bayinx.core.types import T
+
+
+class Observed(Node[T]):
+    """
+    A container for observed nodes of a probabilistic model.
+
+
+    # Attributes
+    - `obj`: An internal realization of the node.
+    - `_filter_spec`: An internal filter specification for `obj`.
+    """
+
+    def __init__(
+        self, obj: T, filter_spec: Optional[PyTree] = None
+    ):
+        if filter_spec is None: # Default filter specification
+            # Generate empty specification
+            filter_spec = jt.map(lambda _: False, obj)
+
+            # Specify array-like leaves
+            filter_spec = eqx.tree_at(
+                where=lambda obj: obj,
+                pytree=filter_spec,
+                replace=jt.map(eqx.is_array_like, obj),
+            )
+
+        self.obj = obj
+        self._filter_spec = filter_spec

bayinx/nodes/stochastic.py

@@ -0,0 +1,25 @@
+from abc import abstractmethod
+from typing import Optional
+
+from bayinx.core.node import Node
+from bayinx.core.types import T
+
+
+class Stochastic(Node[T]):
+    """
+    A container for stochastic (unobserved) nodes of a probabilistic model.
+
+    Subclasses can be constructed with defined filter specifications (implement the `filter_spec` property).
+
+    # Attributes
+    - `obj`: An internal realization of the node.
+    - `_filter_spec`: An internal filter specification for `obj`.
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        obj: T,
+        filter_spec: Optional[T],
+    ):
+        pass

bayinx/ops.py

@@ -0,0 +1,104 @@
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Array, ArrayLike, Real
+
+from bayinx.core.node import Node
+from bayinx.core.utils import _extract_obj
+
+
+def exp(node: Node) -> Node:
+    """
+    Apply the exponential transformation (jnp.exp) to a node.
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf exponential transform
+    def leaf_exp(x: Real[ArrayLike, "..."]) -> Array:
+        return jnp.exp(x)
+
+    # Apply exponential
+    new_obj = jt.map(leaf_exp, obj)
+
+    return type(node)(new_obj, filter_spec)
+
+
+def log(node: Node) -> Node:
+    """
+    Apply the natural logarithm transformation (jnp.log) to a node.
+    Handles input value restrictions (must be positive).
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf log transform
+    def leaf_log(x: Real[ArrayLike, "..."]) -> Array:
+        return jnp.log(x)
+
+    # Apply logarithm
+    new_obj = jt.map(leaf_log, obj)
+
+    return type(node)(new_obj, filter_spec)
+
+
+def sin(node: Node) -> Node:
+    """
+    Apply the sine transformation (jnp.sin) to a node.
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf sine transform
+    def leaf_sin(x: Real[ArrayLike, "..."]) -> Array:
+        return jnp.sin(x)
+
+    # Apply sine
+    new_obj = jt.map(leaf_sin, obj)
+
+    return type(node)(new_obj, filter_spec)
+
+
+def cos(node: Node) -> Node:
+    """
+    Apply the cosine transformation (jnp.cos) to a node.
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf cosine transform
+    def leaf_cos(x: Real[ArrayLike, "..."]) -> Array:
+        return jnp.cos(x)
+
+    # Apply cosine
+    new_obj = jt.map(leaf_cos, obj)
+
+    return type(node)(new_obj, filter_spec)
+
+
+def tanh(node: Node) -> Node:
+    """
+    Apply the hyperbolic tangent transformation (jnp.tanh) to a node.
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf tanh transform
+    def leaf_tanh(x: Real[ArrayLike, "..."]) -> Array:
+        return jnp.tanh(x)
+
+    # Apply tanh
+    new_obj = jt.map(leaf_tanh, obj)
+
+    return type(node)(new_obj, filter_spec)
+
+
+def sigmoid(node: Node) -> Node:
+    """
+    Apply the sigmoid (logistic) transformation to a node.
+    Sigmoid formula: 1 / (1 + exp(-x))
+    """
+    obj, filter_spec = _extract_obj(node)
+
+    # Helper function for the single-leaf sigmoid transform
+    def leaf_sigmoid(x: Real[ArrayLike, "..."]) -> Array:
+        return 1.0 / (1.0 + jnp.exp(-x)) # type: ignore
+
+    # Apply sigmoid
+    new_obj = jt.map(leaf_sigmoid, obj)
+
+    return type(node)(new_obj, filter_spec)

bayinx/posterior.py

@@ -0,0 +1,220 @@
+
+from functools import partial
+from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, Type
+
+import equinox as eqx
+import jax
+import jax.random as jr
+from jax.lax import scan
+from jaxtyping import Array, PRNGKeyArray
+
+from bayinx.core.flow import FlowSpec
+from bayinx.core.node import Node
+from bayinx.core.variational import M
+from bayinx.vi.normalizing_flow import NormalizingFlow
+from bayinx.vi.standard import Standard
+
+
+class Posterior(Generic[M]):
+    """
+    The posterior distribution for a model.
+
+    # Attributes
+    - `vari`: The variational approximation of the posterior.
+    - `config` The configuration for the posterior.
+    """
+    vari: NormalizingFlow[M]
+    config: Dict[str, Any]
+
+    def __init__(self, model_def: Type[M], **kwargs: Any):
+        # (hopefully) omit intermediate model construction through jit
+        @eqx.filter_jit
+        def construct_base(model_def):
+            # Construct model
+            model = model_def(**kwargs)
+
+            return Standard(model)
+
+        # Construct standard normal base distribution
+        self.vari = construct_base(model_def)
+
+        # Include default attributes
+        self.config = {
+            "learning_rate": 0.1 / self.vari.dim**0.5,
+            "tolerance": 1e-4,
+            "grad_draws": 4,
+            "batch_size": 1
+        }
+
+
+    def configure(
+        self,
+        flowspecs: Optional[List[FlowSpec]] = None,
+        learning_rate: Optional[float] = None,
+        tolerance: Optional[float] = None,
+        grad_draws: Optional[int] = None,
+        batch_size: Optional[int] = None
+    ):
+        """
+        Configure the variational approximation.
+
+        # Parameters
+        - `flowspecs`: The specification for a sequence of flows.
+        """
+        # Append new NF architecture
+        if flowspecs is not None:
+            # Initialize NF architecture
+            flows = [
+                flowspec.construct(self.vari.dim) for flowspec in flowspecs
+            ]
+
+            if isinstance(self.vari, Standard):
+                # Construct new normalizing flow
+                self.vari = NormalizingFlow(
+                    base = self.vari,
+                    flows = flows,
+                    _static = self.vari._static,
+                    _unflatten = self.vari._unflatten
+                )
+            elif isinstance(self.vari, NormalizingFlow):
+                # Freeze current flows
+                for flow in self.vari.flows:
+                    object.__setattr__(flow, 'static', True) # kind of illegal but I need to avoid copies
+
+                # Append new flows
+                self.vari.flows.extend(flows)
+
+        # Include other settings
+        if learning_rate is not None:
+            self.config["learning_rate"] = learning_rate
+        if tolerance is not None:
+            self.config["tolerance"] = tolerance
+        if grad_draws is not None:
+            self.config["grad_draws"] = grad_draws
+        if batch_size is not None:
+            self.config["batch_size"] = batch_size
+
+
+    def fit(
+        self,
+        max_iters: int = 50_000,
+        learning_rate: Optional[float] = None,
+        tolerance: Optional[float] = None,
+        grad_draws: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        key: PRNGKeyArray = jr.key(0),
+    ):
+        # Include settings
+        if learning_rate is not None:
+            self.config["learning_rate"] = learning_rate
+        if tolerance is not None:
+            self.config["tolerance"] = tolerance
+        if grad_draws is not None:
+            self.config["grad_draws"] = grad_draws
+        if batch_size is not None:
+            self.config["batch_size"] = batch_size
+
+        if isinstance(self.vari, Standard):
+            # Construct default sequence of optimization
+            print("TODO")
+        else:
+            # Optimize variational approximation with user-specified flows
+            self.vari = self.vari.fit(
+                max_iters,
+                self.config["learning_rate"],
+                self.config["tolerance"],
+                self.config["grad_draws"],
+                self.config["batch_size"],
+                key
+            )
+
+    def sample(
+        self,
+        node: str,
+        n_draws: int,
+        batch_size: Optional[int] = None,
+        key: PRNGKeyArray = jr.key(0)
+    ) -> Array:
+        """
+        Sample a node from the posterior distribution.
+
+        # Parameters
+        - `node`: The name of the node.
+        - `n_draws`: The number of draws from the posterior.
+        - `batch_size`: The number of draws for the full model ever initialized in memory at once.
+        - `key`: The PRNG key.
+        """
+        if batch_size is None:
+            batch_size = n_draws
+
+        # Split keys
+        keys = jr.split(key, n_draws // batch_size)
+
+        @partial(jax.vmap, in_axes = 0)
+        def reconstruct_and_subset(draw: Array):
+            model = self.vari.reconstruct_model(draw).constrain()[0]
+
+            return getattr(model, node).obj
+
+        def batched_sample(carry: None, key: PRNGKeyArray):
+            # Sample draws
+            draws = self.vari.sample(batch_size, key)
+
+            return None, reconstruct_and_subset(draws)
+
+        posterior_draws = scan(
+            batched_sample,
+            init=None,
+            xs=keys,
+            length=n_draws // batch_size
+        )[1].squeeze()
+
+        return posterior_draws
+
+
+    def predictive(
+        self,
+        func: Callable[[M, PRNGKeyArray], Node[Array] | Array],
+        n_draws: int,
+        batch_size: Optional[int] = None,
+        key: PRNGKeyArray = jr.key(0)
+    ) -> Array:
+        """
+
+        """
+        if batch_size is None:
+            batch_size = n_draws
+
+        # Split keys
+        keys = jr.split(key, n_draws // batch_size)
+
+        @partial(jax.vmap, in_axes = (0, 0))
+        def reconstruct_and_predict(draw: Array, key: PRNGKeyArray) -> Array:
+            model = self.vari.reconstruct_model(draw).constrain()[0]
+
+            # Compute predictive
+            obj = func(model, key)
+
+            # Coerce from Node if needed
+            if isinstance(obj, Node):
+                obj: Array = obj.obj # type: ignore
+
+            return obj
+
+        def batched_sample(carry: None, key: PRNGKeyArray) -> Tuple[None, Array]:
+            # Sample draws
+            draws = self.vari.sample(batch_size, key)
+
+            # Generate additional keys for each draw
+            keys = jr.split(key, batch_size)
+
+            return None, reconstruct_and_predict(draws, keys)
+
+        posterior_draws: Array = scan(
+            batched_sample,
+            init=None,
+            xs=keys,
+            length=n_draws // batch_size
+        )[1].squeeze()
+
+        return posterior_draws

bayinx/{mhx/vi → vi}/meanfield.py

@@ -1,25 +1,28 @@
-from typing import Any, Dict, Generic, Self, TypeVar
+from typing import Generic, Self
 
 import equinox as eqx
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
 from jax.flatten_util import ravel_pytree
-from jaxtyping import Array, Float, Key, Scalar
+from jaxtyping import Array, PRNGKeyArray, Scalar
 
-from bayinx.core import Model, Variational
+from bayinx.core.variational import M, Variational
 from bayinx.dists import normal
 
-M = TypeVar('M', bound=Model)
+
 class MeanField(Variational, Generic[M]):
     """
     A fully factorized Gaussian approximation to a posterior distribution.
 
     # Attributes
-    - `var_params`: The variational parameters for the approximation.
+    - `dim`: The dimension of the support.
+    - `mean`: The mean of the unconstrained approximation.
+    - `log_std` The log-transformed standard deviation of the unconstrained approximation.
     """
 
-    var_params: Dict[str, Float[Array, "..."]] #todo: just expand to attributes
+    mean: Array
+    log_std: Array
 
     def __init__(self, model: M):
         """
@@ -27,28 +30,31 @@ class MeanField(Variational, Generic[M]):
 
         # Parameters
         - `model`: A probabilistic `Model` object.
+        - `init_log_std`: The initial log-transformed standard deviation of the Gaussian approximation.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec)
+        params, self._static = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
 
         # Initialize variational parameters
-        self.var_params = {
-            "mean": params,
-            "log_std": jnp.zeros(params.size, dtype=params.dtype),
-        }
+        self.mean = params
+        self.log_std = jnp.full(params.size, 0.0)
 
     @property
-    @eqx.filter_jit
     def filter_spec(self):
         # Generate empty specification
         filter_spec = jtu.tree_map(lambda _: False, self)
 
         # Specify variational parameters
         filter_spec = eqx.tree_at(
-            lambda mf: mf.var_params,
+            lambda mf: mf.mean,
+            filter_spec,
+            replace=True,
+        )
+        filter_spec = eqx.tree_at(
+            lambda mf: mf.log_std,
             filter_spec,
             replace=True,
         )
@@ -56,12 +62,11 @@ class MeanField(Variational, Generic[M]):
         return filter_spec
 
     @eqx.filter_jit
-    def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
+    def sample(self, n: int, key: PRNGKeyArray = jr.PRNGKey(0)) -> Array:
         # Sample variational draws
         draws: Array = (
-            jr.normal(key=key, shape=(n, self.var_params["mean"].size))
-            * jnp.exp(self.var_params["log_std"])
-            + self.var_params["mean"]
+            jr.normal(key=key, shape=(n, self.mean.size)) * jnp.exp(self.log_std)
+            + self.mean
         )
 
         return draws
@@ -70,23 +75,23 @@ class MeanField(Variational, Generic[M]):
     def eval(self, draws: Array) -> Array:
         return normal.logprob(
             x=draws,
-            mu=self.var_params["mean"],
-            sigma=jnp.exp(self.var_params["log_std"]),
+            mu=self.mean,
+            sigma=jnp.exp(self.log_std),
         ).sum(axis=1)
 
     @eqx.filter_jit
-    def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
+    def elbo(self, n: int, batch_size: int, key: PRNGKeyArray) -> Scalar:
         dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_jit
-        def elbo(dyn: Self, n: int, key: Key, data: Any = None) -> Scalar:
+        def elbo(dyn: Self, n: int, key: PRNGKeyArray) -> Scalar:
             vari = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution
             draws: Array = vari.sample(n, key)
 
             # Evaluate posterior density for each draw
-            posterior_evals: Array = vari.eval_model(draws, data)
+            posterior_evals: Array = vari.eval_model(draws)
 
             # Evaluate variational density for each draw
             variational_evals: Array = vari.eval(draws)
@@ -94,23 +99,22 @@ class MeanField(Variational, Generic[M]):
             # Evaluate ELBO
             return jnp.mean(posterior_evals - variational_evals)
 
-        return elbo(dyn, n, key, data)
+        return elbo(dyn, n, key)
 
     @eqx.filter_jit
-    def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
+    def elbo_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Self:
         dyn, static = eqx.partition(self, self.filter_spec)
 
-        @eqx.filter_grad
         @eqx.filter_jit
-        def elbo_grad(dyn: Self, n: int, key: Key, data: Any = None):
-            # Combine
+        @eqx.filter_grad
+        def elbo_grad(dyn: Self, n: int, key: PRNGKeyArray):
             vari = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution
             draws: Array = vari.sample(n, key)
 
             # Evaluate posterior density for each draw
-            posterior_evals: Array = vari.eval_model(draws, data)
+            posterior_evals: Array = vari.eval_model(draws)
 
             # Evaluate variational density for each draw
             variational_evals: Array = vari.eval(draws)
@@ -118,4 +122,4 @@ class MeanField(Variational, Generic[M]):
             # Evaluate ELBO
             return jnp.mean(posterior_evals - variational_evals)
 
-        return elbo_grad(dyn, n, key, data)
+        return elbo_grad(dyn, n, key)