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

bayinx/core/_parameter.py

@@ -1,44 +0,0 @@
-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
-
-    def __call__(self) -> T:
-        return self.vals
-
-    # 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

@@ -1,181 +0,0 @@
-from abc import abstractmethod
-from functools import partial
-from typing import Any, 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, Key, PyTree, Scalar
-from optax import GradientTransformation, OptState, Schedule
-
-from ._model import Model
-
-M = TypeVar('M', bound=Model)
-class Variational(eqx.Module, Generic[M]):
-    """
-    An abstract base class used to define variational methods.
-
-    # Attributes
-    - `_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[[Array], M]
-    _constraints: M
-
-    @abstractmethod
-    def filter_spec(self):
-        """
-        Filter specification for dynamic and static components of the `Variational`.
-        """
-        pass
-
-    @abstractmethod
-    def sample(self, n: int, key: Key = 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, key: Key, data: Any = None) -> Array:
-        """
-        Evaluate the ELBO.
-        """
-        pass
-
-    @abstractmethod
-    def elbo_grad(self, n: int, key: Key, data: Any = None) -> PyTree:
-        """
-        Evaluate the gradient of the ELBO.
-        """
-        pass
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0, None))
-    def eval_model(self, draws: Array, data: Any = None) -> Array:
-        """
-        Reconstruct models from variational draws and evaluate their posterior density.
-
-        # Parameters
-        - `draws`: A set of variational draws.
-        - `data`: Data used to evaluate the posterior(if needed).
-        """
-        # Unflatten variational draw
-        model: M = self._unflatten(draws)
-
-        # Combine with constraints
-        model: M = eqx.combine(model, self._constraints)
-
-        # Evaluate posterior density
-        return model.eval(data)
-
-    @eqx.filter_jit
-    def fit(
-        self,
-        max_iters: int,
-        data: Any = None,
-        learning_rate: float = 1,
-        weight_decay: float = 1e-4,
-        tolerance: float = 1e-4,
-        var_draws: int = 1,
-        key: Key = jr.PRNGKey(0),
-    ) -> Self:
-        """
-        Optimize the variational distribution.
-
-        # Parameters
-        - `max_iters`: Maximum number of iterations for the optimization loop.
-        - `data`: Data to evaluate the posterior density with(if available).
-        - `learning_rate`: Initial learning rate for optimization.
-        - `tolerance`: Relative tolerance of ELBO decrease for stopping early.
-        - `var_draws`: Number of variational draws to draw each iteration.
-        - `key`: A PRNG key.
-        """
-        # Partition variational
-        dyn, static = eqx.partition(self, self.filter_spec)
-
-        # Construct scheduler
-        schedule: Schedule = opx.warmup_cosine_decay_schedule(
-            init_value=1e-16, peak_value=learning_rate, warmup_steps=int(max_iters/10), decay_steps=max_iters-int(max_iters/10)
-        )
-
-        # Initialize optimizer
-        optim: GradientTransformation = opx.chain(
-            opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
-        )
-        opt_state: OptState = optim.init(dyn)
-
-        # Optimization loop helper functions
-        @eqx.filter_jit
-        def condition(state: Tuple[Self, OptState, Scalar, Key]):
-            # Unpack iteration state
-            dyn, opt_state, i, key = state
-
-            return i < max_iters
-
-        @eqx.filter_jit
-        def body(state: Tuple[Self, OptState, Scalar, Key]):
-            # Unpack iteration state
-            dyn, opt_state, i, key = state
-
-            # Update iteration
-            i = i + 1
-
-            # Update PRNG key
-            key, _ = jr.split(key)
-
-            # Reconstruct variational
-            vari = eqx.combine(dyn, static)
-
-            # Compute gradient of the ELBO
-            updates: PyTree = vari.elbo_grad(var_draws, key, data)
-
-            # Compute updates
-            updates, opt_state = optim.update(
-                updates, opt_state, eqx.filter(dyn, dyn.filter_spec)
-            )
-
-            # Update variational distribution
-            dyn = eqx.apply_updates(dyn, updates)
-
-            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)
-
-    @eqx.filter_jit
-    def posterior_predictive(
-        self, func: Callable[[M, Any], Array], n: int, data: Any = None, key: Key = jr.PRNGKey(0)
-    ) -> Array:
-        # Sample draws from the variational approximation
-        draws: Array = self.sample(n, key)
-
-        # Evaluate posterior predictive
-        @jax.jit
-        @jax.vmap
-        def evaluate(draw: Array, data: Any = None):
-            # Reconstruct model
-            model: M = self._unflatten(draw)
-
-            # Evaluate
-            return func(model, data)
-
-        return evaluate(draws, data)

bayinx/dists/censored/__init__.py

@@ -1,3 +0,0 @@
-from . import posnormal
-
-__all__ = ["posnormal"]

bayinx/dists/censored/gamma2/__init__.py

@@ -1,3 +0,0 @@
-from . import r
-
-__all__ = ["r"]

bayinx/dists/censored/gamma2/r.py

@@ -1,68 +0,0 @@
-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.
-    - `censor`: The positive censor value.
-
-    # Returns
-    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `nu`, and `censor`.
-    """
-    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 probability mass/density 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.
-    - `censor`: The positive censor value.
-
-    # Returns
-    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `nu`, and `censor`.
-    """
-    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
-
-    # Evaluate log probability mass/density function
-    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/censored/posnormal/__init__.py

@@ -1,3 +0,0 @@
-from . import r
-
-__all__ = ["r"]

bayinx/dists/censored/posnormal/r.py

@@ -1,116 +0,0 @@
-import jax.numpy as jnp
-import jax.random as jr
-from jaxtyping import Array, ArrayLike, Float, Key
-
-from bayinx.dists import posnormal
-
-
-def prob(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-    censor: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
-    """
-    The mixed probability mass/density function (PMF/PDF) for a right-censored positive Normal distribution.
-
-    # Parameters
-    - `x`:  Value(s) at which to evaluate the PMF/PDF.
-    - `mu`: The mean.
-    - `sigma`: The positive standard deviation.
-    - `censor`: The positive censor value.
-
-    # Returns
-    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `sigma`, and `censor`.
-    """
-    # Cast to Array
-    x, mu, sigma, censor = (
-        jnp.asarray(x),
-        jnp.asarray(mu),
-        jnp.asarray(sigma),
-        jnp.asarray(censor),
-    )
-
-    # Construct boolean masks
-    uncensored: Array = jnp.logical_and(0.0 < x, x < censor)
-    censored: Array = x == censor
-
-    # Evaluate probability mass/density function
-    evals = jnp.where(uncensored, posnormal.prob(x, mu, sigma), 0.0)
-    evals = jnp.where(censored, posnormal.ccdf(x, mu, sigma), evals)
-
-    return evals
-
-
-def logprob(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-    censor: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
-    """
-    The log-transformed mixed probability mass/density function (log PMF/PDF) for a right-censored positive Normal distribution.
-
-    # Parameters
-    - `x`: Where to evaluate the log PMF/PDF.
-    - `mu`: The mean.
-    - `sigma`: The standard deviation.
-    - `censor`: The censor.
-
-    # Returns
-    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `sigma`, and `censor`.
-    """
-    # Cast to Array
-    x, mu, sigma, censor = (
-        jnp.asarray(x),
-        jnp.asarray(mu),
-        jnp.asarray(sigma),
-        jnp.asarray(censor),
-    )
-
-    # Construct boolean masks for censoring
-    uncensored: Array = jnp.logical_and(jnp.asarray(0.0) < x, x < censor)
-    censored: Array = x == censor
-
-    # Evaluate log probability mass/density function
-    evals = jnp.where(uncensored, posnormal.logprob(x, mu, sigma), -jnp.inf)
-    evals = jnp.where(censored, posnormal.logccdf(x, mu, sigma), evals)
-
-    return evals
-
-def sample(
-    n: int,
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-    censor: Float[ArrayLike, "..."] = jnp.inf,
-    key: Key = jr.PRNGKey(0)
-) -> Float[Array, "..."]:
-    """
-    Sample from a right-censored positive Normal distribution.
-
-    # Parameters
-    - `n`: Number of draws to sample per-parameter.
-    - `mu`: The mean.
-    - `sigma`: The standard deviation.
-    - `censor`: The censor.
-
-    # Returns
-    Draws from a right-censored positive Normal distribution. The output will have the shape of (n,) + the broadcasted shapes of `mu`, `sigma`, and `censor`.
-    """
-    # Cast to Array
-    mu, sigma, censor = (
-        jnp.asarray(mu),
-        jnp.asarray(sigma),
-        jnp.asarray(censor),
-    )
-
-    # Derive shape
-    shape = (n,) + jnp.broadcast_shapes(mu.shape, sigma.shape, censor.shape)
-
-    # Draw from positive normal
-    draws = jr.truncated_normal(key, 0.0, jnp.inf, shape) * sigma + mu
-
-    # Censor values
-    draws = jnp.where(censor <= draws, censor, draws)
-
-    return draws

bayinx/dists/gamma2.py

@@ -1,49 +0,0 @@
-import jax.lax as lax
-import jax.numpy as jnp
-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`.
-    """
-    # Cast to Array
-    x, mu, nu = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(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`.
-    """
-    # Cast to Array
-    x, mu, nu = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(nu)
-
-    return (
-        -gammaln(nu)
-        + nu * (lax.log(nu) - lax.log(mu))
-        + (nu - 1.0) * lax.log(x)
-        - (x * nu / mu)
-    )  # pyright: ignore