bayinx 0.2.25__py3-none-any.whl → 0.3.2__py3-none-any.whl

bayinx/__init__.py

@@ -1 +1,2 @@
-from bayinx.core.model import Model as Model
+from bayinx.core import Model as Model
+from bayinx.core import Parameter as Parameter

bayinx/constraints/__init__.py

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

bayinx/constraints/lower.py

@@ -0,0 +1,51 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import PyTree, Scalar, ScalarLike
+
+from bayinx.core.constraint import Constraint
+from bayinx.core.parameter import Parameter
+
+
+class Lower(Constraint):
+    """
+    Enforces a lower bound on the parameter.
+    """
+
+    lb: Scalar
+
+    def __init__(self, lb: ScalarLike):
+        self.lb = jnp.array(lb)
+
+    @eqx.filter_jit
+    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+        """
+        Enforces a lower bound on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `x`: The unconstrained `Parameter`.
+
+        # Parameters
+        A tuple containing:
+            - A modified `Parameter` with relevant leaves satisfying the constraint.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        # Extract relevant filter specification
+        filter_spec = x.filter_spec
+
+        # Extract relevant parameters(all Array)
+        dyn_params, static_params = eqx.partition(x, filter_spec)
+
+        # Compute density adjustment
+        laj: PyTree = jt.map(jnp.sum, dyn_params) # pyright: ignore
+        laj: Scalar = jt.reduce(lambda a,b: a + b, laj)
+
+        # Compute transformation
+        dyn_params = jt.map(lambda v: jnp.exp(v) + self.lb, dyn_params)
+
+        # Combine into full parameter object
+        x = eqx.combine(dyn_params, static_params)
+
+        return x, laj

bayinx/core/__init__.py

@@ -1,3 +1,4 @@
 from bayinx.core.flow import Flow as Flow
 from bayinx.core.model import Model as Model
+from bayinx.core.parameter import Parameter as Parameter
 from bayinx.core.variational import Variational as Variational

bayinx/core/constraint.py

@@ -0,0 +1,28 @@
+from abc import abstractmethod
+from typing import Tuple
+
+import equinox as eqx
+from jaxtyping import Scalar
+
+from bayinx.core.parameter import Parameter
+
+
+class Constraint(eqx.Module):
+    """
+    Abstract base class for defining parameter constraints.
+    """
+
+    @abstractmethod
+    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+        """
+        Applies the constraining transformation to a parameter and computes the log-absolute-Jacobian of the transformation.
+
+        # Parameters
+        - `x`: The unconstrained `Parameter`.
+
+        # Returns
+        A tuple containing:
+            - The constrained `Parameter`.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        pass

bayinx/core/flow.py

@@ -8,11 +8,11 @@ from jaxtyping import Array, Float
 
 class Flow(eqx.Module):
     """
-    A superclass used to define continuously parameterized diffeomorphisms for normalizing flows.
+    An abstract base class for a flow(of a normalizing flow).
 
     # Attributes
     - `pars`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
-    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    - `constraints`: A dictionary of simple functions that constrain their corresponding parameter.
     """
 
     params: Dict[str, Float[Array, "..."]]
@@ -28,14 +28,16 @@ class Flow(eqx.Module):
     @abstractmethod
     def adjust_density(self, draws: Array) -> Tuple[Array, Array]:
         """
-        Computes the log-absolute-determinant of the Jacobian at `draws` and applies the forward transformation.
+        Computes the log-absolute-Jacobian at `draws` and applies the forward transformation.
 
         # Returns
-        A tuple of JAX Arrays containing the log-absolute-determinant of the Jacobians and transformed draws.
+            A tuple of JAX Arrays containing the transformed draws and log-absolute-Jacobians.
         """
         pass
 
     # Default filter specification
+    @property
+    @eqx.filter_jit
     def filter_spec(self):
         """
         Generates a filter specification to subset relevant parameters for the flow.
@@ -53,7 +55,7 @@ class Flow(eqx.Module):
         return filter_spec
 
     @eqx.filter_jit
-    def constrain_pars(self: Self):
+    def constrain_params(self: Self):
         """
         Constrain `params` to the appropriate domain.
 
@@ -68,11 +70,11 @@ class Flow(eqx.Module):
         return t_params
 
     @eqx.filter_jit
-    def transform_pars(self: Self) -> Dict[str, Array]:
+    def transform_params(self: Self) -> Dict[str, Array]:
         """
         Apply a custom transformation to `params` if needed.
 
         # Returns
         A dictionary of transformed JAX Arrays representing the transformed parameters.
         """
-        return self.constrain_pars()
+        return self.constrain_params()

bayinx/core/model.py

@@ -1,24 +1,25 @@
 from abc import abstractmethod
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Generic, Tuple, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
-import jax.tree_util as jtu
-from jaxtyping import Array, Scalar
+import jax.tree as jt
+from jaxtyping import PyTree, Scalar
 
-from bayinx.core.constraints import Constraint
+from bayinx.core.constraint import Constraint
+from bayinx.core.parameter import Parameter
 
-
-class Model(eqx.Module):
+P = TypeVar('P', bound=Dict[str, Parameter[PyTree]])
+class Model(eqx.Module, Generic[P]):
     """
-    A superclass used to define probabilistic models.
+    An abstract base class used to define probabilistic models.
 
     # Attributes
-    - `params`: A dictionary of JAX Arrays representing parameters of the model.
-    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    - `params`: A dictionary of parameters.
+    - `constraints`: A dictionary of constraints.
     """
 
-    params: Dict[str, Array]
+    params: P
     constraints: Dict[str, Constraint]
 
     @abstractmethod
@@ -26,32 +27,34 @@ class Model(eqx.Module):
         pass
 
     # Default filter specification
+    @property
+    @eqx.filter_jit
     def filter_spec(self):
         """
         Generates a filter specification to subset relevant parameters for the model.
         """
         # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
+        filter_spec = jt.map(lambda _: False, self)
 
-        # Specify JAX Array parameters
+        # Specify relevant parameters
         filter_spec = eqx.tree_at(
             lambda model: model.params,
             filter_spec,
-            replace=jtu.tree_map(eqx.is_array, self.params),
+            replace={key: param.filter_spec for key, param in self.params.items()}
         )
 
         return filter_spec
 
     # Add constrain method
     @eqx.filter_jit
-    def constrain_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+    def constrain_params(self) -> Tuple[P, Scalar]:
         """
         Constrain `params` to the appropriate domain.
 
         # Returns
-        A dictionary of transformed JAX Arrays representing the constrained parameters and the adjustment to the posterior density.
+        A dictionary of PyTrees representing the constrained parameters and the adjustment to the posterior density.
         """
-        t_params: Dict[str, Array] = self.params
+        t_params: P = self.params
         target: Scalar = jnp.array(0.0)
 
         for par, map in self.constraints.items():
@@ -63,12 +66,13 @@ class Model(eqx.Module):
 
         return t_params, target
 
-
-    def transform_pars(self) -> Tuple[Dict[str, Array], Scalar]:
+    # Add default transform method
+    @eqx.filter_jit
+    def transform_params(self) -> Tuple[P, Scalar]:
         """
         Apply a custom transformation to `params` if needed.
 
         # Returns
         A dictionary of transformed JAX Arrays representing the transformed parameters.
         """
-        return self.constrain_pars()
+        return self.constrain_params()

bayinx/core/parameter.py

@@ -0,0 +1,41 @@
+from typing import Generic, Self, TypeVar
+
+import equinox as eqx
+import jax.tree as jt
+from jaxtyping import PyTree
+
+T = TypeVar('T', bound=PyTree)
+class Parameter(eqx.Module, Generic[T]):
+    """
+    A container for a parameter of a `Model`.
+
+    Subclasses can be constructed for custom filter specifications(`filter_spec`).
+
+    # Attributes
+    - `vals`: The parameter's value(s).
+    """
+    vals: T
+
+
+    def __init__(self, values: T):
+        # Insert parameter values
+        self.vals = values
+
+    # Default filter specification
+    @property
+    @eqx.filter_jit
+    def filter_spec(self) -> Self:
+        """
+        Generates a filter specification to filter out static parameters.
+        """
+        # Generate empty specification
+        filter_spec = jt.map(lambda _: False, self)
+
+        # Specify Array leaves
+        filter_spec = eqx.tree_at(
+            lambda params: params.vals,
+            filter_spec,
+            replace=jt.map(eqx.is_array_like, self.vals),
+        )
+
+        return filter_spec

bayinx/core/variational.py

@@ -8,7 +8,7 @@ import jax.lax as lax
 import jax.numpy as jnp
 import jax.random as jr
 import optax as opx
-from jaxtyping import Array, Float, Key, PyTree, Scalar
+from jaxtyping import Array, Key, PyTree, Scalar
 from optax import GradientTransformation, OptState, Schedule
 
 from bayinx.core import Model
@@ -16,16 +16,23 @@ from bayinx.core import Model
 
 class Variational(eqx.Module):
     """
-    A superclass used to define variational methods.
+    An abstract base class used to define variational methods.
 
     # Attributes
-    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
-    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
+    - `_unflatten`: A function to transform draws from the variational distribution back to a `Model`.
+    - `_constraints`: The static component of a partitioned `Model` used to initialize the `Variational` object.
     """
 
-    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _unflatten: Callable[[Array], Model]
     _constraints: Model
 
+    @abstractmethod
+    def filter_spec(self):
+        """
+        Filter specification for dynamic and static components of the `Variational`.
+        """
+        pass
+
     @abstractmethod
     def sample(self, n: int, key: Key) -> Array:
         """
@@ -54,13 +61,6 @@ class Variational(eqx.Module):
         """
         pass
 
-    @abstractmethod
-    def filter_spec(self):
-        """
-        Filter specification for dynamic and static components of the `Variational`.
-        """
-        pass
-
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0, None))
     def eval_model(self, draws: Array, data: Any = None) -> Array:
@@ -103,7 +103,7 @@ class Variational(eqx.Module):
         - `key`: A PRNG key.
         """
         # Partition variational
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         # Construct scheduler
         schedule: Schedule = opx.cosine_decay_schedule(
@@ -135,7 +135,7 @@ class Variational(eqx.Module):
             # Update PRNG key
             key, _ = jr.split(key)
 
-            # Combine variational
+            # Reconstruct variational
             vari = eqx.combine(dyn, static)
 
             # Compute gradient of the ELBO
@@ -143,7 +143,7 @@ class Variational(eqx.Module):
 
             # Compute updates
             updates, opt_state = optim.update(
-                updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
+                updates, opt_state, eqx.filter(dyn, dyn.filter_spec)
             )
 
             # Update variational distribution

bayinx/dists/censored/gamma2/r.py

@@ -0,0 +1,65 @@
+import jax.lax as lax
+import jax.numpy as jnp
+from jax.scipy.special import gammaincc
+from jaxtyping import Array, ArrayLike, Float
+
+from bayinx.dists import gamma2
+
+
+def prob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    nu: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The mixed probability mass/density function (PMF/PDF) for a (mean-inverse dispersion parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:  Value(s) at which to evaluate the PMF/PDF.
+    - `mu`: The positive mean.
+    - `nu`: The positive inverse dispersion.
+
+    # Returns
+    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+    evals: Array = jnp.zeros_like(x * 1.0) # ensure float dtype
+
+    # Construct boolean masks
+    uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
+    censored: Array = jnp.array(x == censor) # pyright: ignore
+
+    # Evaluate mixed probability (?) function
+    evals = jnp.where(uncensored, gamma2.prob(x, mu, nu), evals)
+    evals = jnp.where(censored, gammaincc(nu, x * nu / mu), evals) # pyright: ignore
+
+    return evals
+
+
+def logprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    nu: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed mixed probability mass/density function (log PMF/PDF) for a (mean-inverse dispersion parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PMF/PDF.
+    - `mu`:     The positive mean/location.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+    evals: Array = jnp.full_like(x * 1.0, -jnp.inf) # ensure float dtype
+
+    # Construct boolean masks
+    uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
+    censored: Array = jnp.array(x == censor) # pyright: ignore
+
+    evals = jnp.where(uncensored, gamma2.logprob(x, mu, nu), evals)
+    evals = jnp.where(censored, lax.log(gammaincc(nu, x * nu / mu)), evals) # pyright: ignore
+
+    return evals

bayinx/dists/gamma2.py

@@ -0,0 +1,39 @@
+import jax.lax as lax
+from jax.scipy.special import gammaln
+from jaxtyping import Array, ArrayLike, Float
+
+
+def prob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], nu: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The probability density function (PDF) for a (mean-precision parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `mu`:     The positive mean.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+
+    return lax.exp(logprob(x, mu, nu))
+
+
+def logprob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], nu: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed probability density function (log PDF) for a (mean-precision parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PDF.
+    - `mu`:     The positive mean/location.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+
+    return - gammaln(nu) + nu * (lax.log(nu) - lax.log(mu)) + (nu - 1.0) * lax.log(x) - (x * nu / mu) # pyright: ignore

bayinx/dists/normal.py

@@ -1,31 +1,31 @@
-import jax.lax as _lax
-from jaxtyping import Array, ArrayLike, Float, Real
+import jax.lax as lax
+from jaxtyping import Array, ArrayLike, Float
 
 __PI = 3.141592653589793
 
 
 def prob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Normal distribution.
 
     # Parameters
     - `x`:      Value(s) at which to evaluate the PDF.
-    - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `mu`:     The mean/location.
+    - `sigma`:  The positive standard deviation.
 
     # Returns
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / ( # pyright: ignore
-        sigma * _lax.sqrt(2.0 * __PI)
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / (  # pyright: ignore
+        sigma * lax.sqrt(2.0 * __PI)
     )
 
 
 def logprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Normal distribution.
@@ -36,14 +36,16 @@ def logprob(
     - `sigma`:  The non-negative standard deviation parameter(s).
 
     # Returns
-    The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return -_lax.log(sigma * _lax.sqrt(2.0 * __PI)) - 0.5 * _lax.square((x - mu) / sigma) # pyright: ignore
+    return -lax.log(sigma * lax.sqrt(2.0 * __PI)) - 0.5 * lax.square(
+        (x - mu) / sigma  # pyright: ignore
+    )
 
 
 def uprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Normal distribution.
@@ -51,17 +53,17 @@ def uprob(
     # Parameters
     - `x`:      Value(s) at which to evaluate the uPDF.
     - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `sigma`:  The positive standard deviation parameter(s).
 
     # Returns
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / sigma # pyright: ignore
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / sigma  # pyright: ignore
 
 
 def ulogprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
@@ -75,4 +77,4 @@ def ulogprob(
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return -_lax.log(sigma) - 0.5 * _lax.square((x - mu) / sigma) # pyright: ignore
+    return -lax.log(sigma) - 0.5 * lax.square((x - mu) / sigma)  # pyright: ignore

bayinx/dists/uniform.py

@@ -1,10 +1,10 @@
 import jax.lax as _lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float, Real
+from jaxtyping import Array, ArrayLike, Float
 
 
 def prob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Uniform distribution.
@@ -18,11 +18,11 @@ def prob(
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
 
-    return 1.0 / (ub - lb) # pyright: ignore
+    return 1.0 / (ub - lb)  # pyright: ignore
 
 
 def logprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Uniform distribution.
@@ -36,11 +36,11 @@ def logprob(
     The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
 
-    return _lax.log(1.0) - _lax.log(ub - lb) # pyright: ignore
+    return _lax.log(1.0) - _lax.log(ub - lb)  # pyright: ignore
 
 
 def uprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Uniform distribution.
@@ -54,11 +54,11 @@ def uprob(
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
 
-    return jnp.ones(jnp.broadcast_arrays(x,lb,ub))
+    return jnp.ones(jnp.broadcast_arrays(x, lb, ub))
 
 
 def ulogprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Uniform distribution.
@@ -72,4 +72,4 @@ def ulogprob(
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
 
-    return jnp.zeros(jnp.broadcast_arrays(x,lb,ub))
+    return jnp.zeros(jnp.broadcast_arrays(x, lb, ub))

bayinx/mhx/vi/flows/fullaffine.py

@@ -1,29 +1,26 @@
 from functools import partial
-from typing import Callable, Dict, Tuple
+from typing import Tuple
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, Float, Scalar
+from jaxtyping import Array, Scalar
 
 from bayinx.core import Flow
 
 
 class FullAffine(Flow):
     """
-    An affine flow.
+    A full affine flow.
 
     # Attributes
     - `params`: A dictionary containing the JAX Arrays representing the scale and shift parameters.
     - `constraints`: A dictionary of constraining transformations.
     """
 
-    params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]
-
     def __init__(self, dim: int):
         """
-        Initializes an affine flow.
+        Initializes a full affine flow.
 
         # Parameters
         - `dim`: The dimension of the parameter space.
@@ -33,11 +30,23 @@ class FullAffine(Flow):
             "scale": jnp.zeros((dim, dim)),
         }
 
-        self.constraints = {"scale": lambda m: jnp.tril(jnp.exp(m))}
+        if dim == 1:
+            self.constraints = {}
+        else:
+
+            @eqx.filter_jit
+            def constrain_scale(scale: Array):
+                # Extract diagonal and apply exponential
+                diag: Array = jnp.exp(jnp.diag(scale))
+
+                # Return matrix with modified diagonal
+                return jnp.fill_diagonal(jnp.tril(scale), diag, inplace=False)
+
+            self.constraints = {"scale": constrain_scale}
 
     @eqx.filter_jit
     def forward(self, draws: Array) -> Array:
-        params = self.constrain_pars()
+        params = self.transform_params()
 
         # Extract parameters
         shift: Array = params["shift"]
@@ -50,8 +59,8 @@ class FullAffine(Flow):
 
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
-    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.constrain_pars()
+    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
+        params = self.transform_params()
 
         # Extract parameters
         shift: Array = params["shift"]
@@ -60,7 +69,7 @@ class FullAffine(Flow):
         # Compute forward transformation
         draws = draws @ scale + shift
 
-        # Compute ladj
-        ladj: Scalar = jnp.log(jnp.diag(scale)).sum()
+        # Compute laj
+        laj: Scalar = jnp.log(jnp.diag(scale)).sum()
 
-        return ladj, draws
+        return draws, laj

bayinx/mhx/vi/flows/planar.py

@@ -39,7 +39,7 @@ class Planar(Flow):
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         w: Array = params["w"]
@@ -53,8 +53,8 @@ class Planar(Flow):
 
     @eqx.filter_jit
     @partial(jax.vmap, in_axes=(None, 0))
-    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.transform_pars()
+    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
+        params = self.transform_params()
 
         # Extract parameters
         w: Array = params["w"]
@@ -67,8 +67,8 @@ class Planar(Flow):
         # Compute forward transformation
         draws = draws + u * jnp.tanh(x)
 
-        # Compute ladj
+        # Compute laj
         h_prime: Scalar = 1.0 - jnp.square(jnp.tanh(x))
-        ladj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
+        laj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
 
-        return ladj, draws
+        return draws, laj

bayinx/mhx/vi/flows/radial.py

@@ -49,7 +49,7 @@ class Radial(Flow):
         # Returns
         The transformed samples.
         """
-        params = self.transform_pars()
+        params = self.transform_params()
 
         # Extract parameters
         alpha = params["alpha"]
@@ -66,8 +66,8 @@ class Radial(Flow):
 
     @partial(jax.vmap, in_axes=(None, 0))
     @eqx.filter_jit
-    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.transform_pars()
+    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
+        params = self.transform_params()
 
         # Extract parameters
         alpha = params["alpha"]
@@ -84,11 +84,11 @@ class Radial(Flow):
         draws = draws + (x) * (draws - center)
 
         # Compute density adjustment
-        ladj = jnp.log(
+        laj = jnp.log(
             jnp.abs(
                 (1.0 + alpha * beta / (alpha + r) ** 2.0)
                 * (1.0 + x) ** (center.size - 1.0)
             )
         )
 
-        return ladj, draws
+        return draws, laj

bayinx/mhx/vi/meanfield.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Self
+from typing import Any, Dict, Self
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -20,8 +20,6 @@ class MeanField(Variational):
     """
 
     var_params: Dict[str, Float[Array, "..."]]
-    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
-    _constraints: Model = eqx.field(static=True)
 
     def __init__(self, model: Model):
         """
@@ -31,7 +29,7 @@ class MeanField(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec())
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
@@ -42,6 +40,22 @@ class MeanField(Variational):
             "log_std": jnp.zeros(params.size, dtype=params.dtype),
         }
 
+    @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,
+            filter_spec,
+            replace=True,
+        )
+
+        return filter_spec
+
+
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws
@@ -61,27 +75,12 @@ class MeanField(Variational):
             sigma=jnp.exp(self.var_params["log_std"]),
         ).sum(axis=1)
 
-    @eqx.filter_jit
-    def filter_spec(self):
-        filter_spec = jtu.tree_map(lambda _: False, self)
-        filter_spec = eqx.tree_at(
-            lambda mf: mf.var_params,
-            filter_spec,
-            replace=True,
-        )
-        return filter_spec
-
     @eqx.filter_jit
     def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
-        """
-        Estimate the ELBO and its gradient(w.r.t the variational parameters).
-        """
-        # Partition variational
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_jit
         def elbo(dyn: Self, n: int, key: Key, data: Any = None) -> Scalar:
-            # Combine
             vari = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution
@@ -100,8 +99,7 @@ class MeanField(Variational):
 
     @eqx.filter_jit
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        # Partition
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_grad
         @eqx.filter_jit

bayinx/mhx/vi/normalizing_flow.py

@@ -1,11 +1,11 @@
-from typing import Any, Callable, Self, Tuple
+from typing import Any, Self, Tuple
 
 import equinox as eqx
 import jax.flatten_util as jfu
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
-from jaxtyping import Array, Float, Key, Scalar
+from jaxtyping import Array, Key, Scalar
 
 from bayinx.core import Flow, Model, Variational
 
@@ -17,14 +17,11 @@ class NormalizingFlow(Variational):
 
     # Attributes
     - `base`: A base variational distribution.
-    - `flows`: An ordered collection of continuously parameterized
-        diffeomorphisms.
+    - `flows`: An ordered collection of continuously parameterized diffeomorphisms.
     """
 
     flows: list[Flow]
     base: Variational
-    _unflatten: Callable[[Float[Array, "..."]], Model]
-    _constraints: Model
 
     def __init__(self, base: Variational, flows: list[Flow], model: Model):
         """
@@ -36,7 +33,7 @@ class NormalizingFlow(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, eqx.is_array)
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         _, self._unflatten = jfu.ravel_pytree(params)
@@ -44,6 +41,21 @@ class NormalizingFlow(Variational):
         self.base = base
         self.flows = flows
 
+    @property
+    @eqx.filter_jit
+    def filter_spec(self):
+        # Generate empty specification
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        # Specify variational parameters based on each flow's filter spec.
+        filter_spec = eqx.tree_at(
+            lambda vari: vari.flows,
+            filter_spec,
+            replace=[flow.filter_spec for flow in self.flows],
+        )
+
+        return filter_spec
+
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)):
         """
@@ -65,19 +77,18 @@ class NormalizingFlow(Variational):
 
         for map in self.flows:
             # Compute adjustment
-            ladj, draws = map.adjust_density(draws)
+            draws, laj = map.adjust_density(draws)
 
             # Adjust variational density
-            variational_evals = variational_evals - ladj
+            variational_evals = variational_evals - laj
 
         return variational_evals
 
     @eqx.filter_jit
     def __eval(self, draws: Array, data=None) -> Tuple[Array, Array]:
         """
-        Evaluate the posterior and variational densities at the transformed
-        `draws` to avoid extra compute when requiring variational draws for
-        the posterior evaluation.
+        Evaluate the posterior and variational densities together at the
+        transformed `draws` to avoid extra compute.
 
         # Parameters
         - `draws`: Draws from the base variational distribution.
@@ -91,32 +102,19 @@ class NormalizingFlow(Variational):
 
         for map in self.flows:
             # Compute adjustment
-            ladj, draws = map.adjust_density(draws)
+            draws, laj = map.adjust_density(draws)
 
             # Adjust variational density
-            variational_evals = variational_evals - ladj
+            variational_evals = variational_evals - laj
 
         # Evaluate posterior at final variational draws
         posterior_evals = self.eval_model(draws, data)
 
         return posterior_evals, variational_evals
 
-    def filter_spec(self):
-        # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        # Specify variational parameters based on each flow's filter spec.
-        filter_spec = eqx.tree_at(
-            lambda vari: vari.flows,
-            filter_spec,
-            replace=[flow.filter_spec() for flow in self.flows],
-        )
-
-        return filter_spec
-
     @eqx.filter_jit
     def elbo(self, n: int, key: Key = jr.PRNGKey(0), data: Any = None) -> Scalar:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_jit
         def elbo(dyn: Self, n: int, key: Key, data: Any = None):
@@ -133,7 +131,7 @@ class NormalizingFlow(Variational):
 
     @eqx.filter_jit
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        dyn, static = eqx.partition(self, self.filter_spec())
+        dyn, static = eqx.partition(self, self.filter_spec)
 
         @eqx.filter_grad
         @eqx.filter_jit

bayinx/mhx/vi/standard.py

@@ -19,7 +19,7 @@ class Standard(Variational):
     - `dim`: Dimension of the parameter space.
     """
 
-    dim: int = eqx.field(static=True)
+    dim: int
     _unflatten: Callable[[Float[Array, "..."]], Model]
     _constraints: Model
 
@@ -31,7 +31,7 @@ class Standard(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec())
+        params, self._constraints = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
@@ -54,7 +54,7 @@ class Standard(Variational):
             sigma=jnp.array(1.0),
         ).sum(axis=1, keepdims=True)
 
-    @eqx.filter_jit
+    @property
     def filter_spec(self):
         filter_spec = jtu.tree_map(lambda _: False, self)
 

{bayinx-0.2.25.dist-info → bayinx-0.3.2.dist-info}/METADATA

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

bayinx-0.3.2.dist-info/RECORD

@@ -0,0 +1,30 @@
+bayinx/__init__.py,sha256=htihTsJ54k-ljBLzt4ye8DR7ORwZhxv-bLMcEaDQeqY,86
+bayinx/py.typed,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=ADSMapUJGyvKf_TpeC7Foaa3BJ03_Kc7FZxIEKNQkZE,2228
+bayinx/core/parameter.py,sha256=oxCCZcZ-DDBvfWzexfhQkSJPxNQnE1vYXtBhiEZG2bM,1025
+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
+bayinx/dists/gamma2.py,sha256=8XYaOtcYJCrr5q1yHWfZaMJmASpLOrfyhrH_J06ksj8,1333
+bayinx/dists/normal.py,sha256=mvm6EoAlORy-yivuhMcExYCZUo0vJzMKMOWH-9iQBZU,2634
+bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
+bayinx/dists/censored/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/censored/gamma2/r.py,sha256=3brRCKhE-74mRXyIyPcnyaWY2OJv8CZyUWPP9T1t09Y,2274
+bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
+bayinx/mhx/vi/meanfield.py,sha256=M4QrOuHaIMLTuQSD6JNF9vELnTm370tXV68JPB7B67M,3652
+bayinx/mhx/vi/normalizing_flow.py,sha256=9c5ayMJ_Wgq6pUb1GYHIFIzw3Bf1AsIdOjcerLoYMrA,4655
+bayinx/mhx/vi/standard.py,sha256=DfSV0r9oXzp9UM8OsZBpoJPRUhiDoAq_X2_2z_M83lA,1685
+bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
+bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvADwVYQ,1954
+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.2.dist-info/METADATA,sha256=9cltWLDiwqg6VpnufQfKYEw_5ZCywJRp7gAPZAogLlA,3057
+bayinx-0.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.2.dist-info/RECORD,,

bayinx/core/constraints.py

@@ -1,61 +0,0 @@
-from abc import abstractmethod
-from typing import Tuple
-
-import equinox as eqx
-import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Scalar, ScalarLike
-
-
-class Constraint(eqx.Module):
-    """
-    Abstract base class for defining parameter constraints.
-
-    Subclasses should implement the `constrain` method to apply the
-    transformation and compute the ladj adjustment.
-    """
-    @abstractmethod
-    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
-        """
-        Applies the constraining transformation to an unconstrained input
-        and computes the log absolute determinant of the Jacobian (ladj)
-        of this transformation.
-
-        # Parameters
-        - `x`: The unconstrained JAX Array-like input.
-
-        # Returns
-        A tuple containing:
-            - The constrained JAX Array.
-            - A scalar JAX Array representing the ladj of the transformation.
-        """
-        pass
-
-
-class LowerBound(Constraint):
-    """
-    Enforces a lower bound on the parameter.
-    """
-    lb: ScalarLike
-
-    def __init__(self, lb: ScalarLike):
-        self.lb = lb
-
-    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
-        """
-        Applies the lower bound constraint and computes the ladj.
-
-        # Parameters
-        - `x`: The unconstrained JAX Array-like input.
-
-        # Parameters
-        A tuple containing:
-            - The constrained JAX Array (x > self.lb).
-            - A scalar JAX Array representing the ladj of the transformation.
-        """
-        # Compute transformation adjustment
-        ladj = jnp.sum(x)
-
-        # Compute transformation
-        x = jnp.exp(x) + self.lb
-
-        return x, ladj

bayinx/core/utils.py

@@ -1 +0,0 @@
-

bayinx-0.2.25.dist-info/RECORD

@@ -1,28 +0,0 @@
-bayinx/__init__.py,sha256=l20JdkSsE_XGZlZFNEtySXf4NIlbjrao14vXPB-H6aQ,45
-bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
-bayinx/core/constraints.py,sha256=Y8FJX3CkgnLQ3HXuTPGuzvLtXlKs0B7z0-YymoHgdfg,1682
-bayinx/core/flow.py,sha256=9swS5wh7AsIZWgG_IhQS-upcPlr7G-juaP_5rsbX6G0,2363
-bayinx/core/model.py,sha256=U1xBnAXnIvFJjWF-XIT8BYjP9PtoRZY_PwyhRwJf-HA,2144
-bayinx/core/utils.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-bayinx/core/variational.py,sha256=vUZ6u5CXCHfs6ZrA8PF4PHfmUXHTK2RJGHyZ3afFfsg,4820
-bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
-bayinx/dists/binomial.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/gamma.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/gamma2.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/normal.py,sha256=rtSDi0NAObH1LGRWiPZk_6cbSVv2dOPHkgxtWn6gFgM,2662
-bayinx/dists/uniform.py,sha256=PSZIIc2QfNF5XYi-TLGltnr_vnAIA-MZsn1rKV8QXAo,2353
-bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
-bayinx/mhx/vi/meanfield.py,sha256=LNLwfjKO9os7YBmRBpGEiFxlxonuN7DHVFEmXV3hvfI,3876
-bayinx/mhx/vi/normalizing_flow.py,sha256=nj7bpIoMJl6GTOXPxQCAsPArchbHd5vwwqMm3cLbMII,4791
-bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
-bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
-bayinx/mhx/vi/flows/fullaffine.py,sha256=2QbOtA1Jmu-yRcJeFmCKc8N1atm8G7JXYMLEZaEXKV0,1711
-bayinx/mhx/vi/flows/planar.py,sha256=u9ZVwEeOv4fMfwiORlseCz463atsWMuid_9crRg05Z8,1919
-bayinx/mhx/vi/flows/radial.py,sha256=c-SWybGn_jmgBQk9ZMQ5uHKPzcdhowNp8MD8t1-8VZM,2501
-bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.25.dist-info/METADATA,sha256=9TI5NY4M1EBtYwA8E-EDds-QkOECwrfVaaEWq3pTdu4,3058
-bayinx-0.2.25.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.25.dist-info/RECORD,,