bayinx 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

bayinx/__init__.py

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

bayinx/constraints/__init__.py

@@ -1 +1,3 @@
-from bayinx.constraints.lower import Lower as Lower
+from bayinx.constraints.lower import Lower
+
+__all__ = ['Lower']

bayinx/constraints/lower.py

@@ -5,8 +5,7 @@ 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
+from bayinx.core import Constraint, Parameter
 
 
 class Lower(Constraint):
@@ -39,8 +38,8 @@ class Lower(Constraint):
         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)
+        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)

bayinx/core/__init__.py

@@ -1,4 +1,7 @@
-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
+from ._constraint import Constraint
+from ._flow import Flow
+from ._model import Model, constrain
+from ._parameter import Parameter
+from ._variational import Variational
+
+__all__ = ["Constraint", "Flow", "Model", "constrain", "Parameter", "Variational"]

bayinx/core/{constraint.py → _constraint.py}

@@ -4,7 +4,7 @@ from typing import Tuple
 import equinox as eqx
 from jaxtyping import Scalar
 
-from bayinx.core.parameter import Parameter
+from bayinx.core._parameter import Parameter
 
 
 class Constraint(eqx.Module):

bayinx/core/{model.py → _model.py}

@@ -7,13 +7,13 @@ import jax.numpy as jnp
 import jax.tree as jt
 from jaxtyping import Scalar
 
-from bayinx.core.constraint import Constraint
-from bayinx.core.parameter import Parameter
+from ._constraint import Constraint
+from ._parameter import Parameter
 
 
 def constrain(constraint: Constraint):
     """Defines constraint metadata."""
-    return field(metadata={'constraint': constraint})
+    return field(metadata={"constraint": constraint})
 
 
 class Model(eqx.Module):
@@ -49,12 +49,11 @@ class Model(eqx.Module):
                 filter_spec = eqx.tree_at(
                     lambda model: getattr(model, f.name),
                     filter_spec,
-                    replace=attr.filter_spec
+                    replace=attr.filter_spec,
                 )
 
         return filter_spec
 
-
     @eqx.filter_jit
     def constrain_params(self) -> Tuple[Self, Scalar]:
         """
@@ -71,18 +70,16 @@ class Model(eqx.Module):
             attr = getattr(self, f.name)
 
             # Check if constrained parameter
-            if isinstance(attr, Parameter) and 'constraint' in f.metadata:
+            if isinstance(attr, Parameter) and "constraint" in f.metadata:
                 param = attr
-                constraint = f.metadata['constraint']
+                constraint = f.metadata["constraint"]
 
                 # Apply constraint
                 param, laj = constraint.constrain(param)
 
                 # Update parameters for constrained model
                 constrained = eqx.tree_at(
-                    lambda model: getattr(model, f.name),
-                    constrained,
-                    replace=param
+                    lambda model: getattr(model, f.name), constrained, replace=param
                 )
 
                 # Adjust posterior density
@@ -90,7 +87,6 @@ class Model(eqx.Module):
 
         return constrained, target
 
-
     @eqx.filter_jit
     def transform_params(self) -> Tuple[Self, Scalar]:
         """

bayinx/core/{parameter.py → _parameter.py}

@@ -4,7 +4,7 @@ import equinox as eqx
 import jax.tree as jt
 from jaxtyping import PyTree
 
-T = TypeVar('T', bound=PyTree)
+T = TypeVar("T", bound=PyTree)
 class Parameter(eqx.Module, Generic[T]):
     """
     A container for a parameter of a `Model`.
@@ -14,8 +14,8 @@ class Parameter(eqx.Module, Generic[T]):
     # Attributes
     - `vals`: The parameter's value(s).
     """
-    vals: T
 
+    vals: T
 
     def __init__(self, values: T):
         # Insert parameter values

bayinx/core/{variational.py → _variational.py}

@@ -1,6 +1,6 @@
 from abc import abstractmethod
 from functools import partial
-from typing import Any, Callable, Self, Tuple
+from typing import Any, Callable, Generic, Self, Tuple, TypeVar
 
 import equinox as eqx
 import jax
@@ -11,10 +11,10 @@ import optax as opx
 from jaxtyping import Array, Key, PyTree, Scalar
 from optax import GradientTransformation, OptState, Schedule
 
-from bayinx.core import Model
+from ._model import Model
 
-
-class Variational(eqx.Module):
+M = TypeVar('M', bound=Model)
+class Variational(eqx.Module, Generic[M]):
     """
     An abstract base class used to define variational methods.
 
@@ -23,8 +23,8 @@ class Variational(eqx.Module):
     - `_constraints`: The static component of a partitioned `Model` used to initialize the `Variational` object.
     """
 
-    _unflatten: Callable[[Array], Model]
-    _constraints: Model
+    _unflatten: Callable[[Array], M]
+    _constraints: M
 
     @abstractmethod
     def filter_spec(self):
@@ -34,7 +34,7 @@ class Variational(eqx.Module):
         pass
 
     @abstractmethod
-    def sample(self, n: int, key: Key) -> Array:
+    def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         """
         Sample from the variational distribution.
         """
@@ -72,10 +72,10 @@ class Variational(eqx.Module):
         - `data`: Data used to evaluate the posterior(if needed).
         """
         # Unflatten variational draw
-        model: Model = self._unflatten(draws)
+        model: M = self._unflatten(draws)
 
         # Combine with constraints
-        model: Model = eqx.combine(model, self._constraints)
+        model: M = eqx.combine(model, self._constraints)
 
         # Evaluate posterior density
         return model.eval(data)
@@ -160,3 +160,22 @@ class Variational(eqx.Module):
 
         # Return optimized variational
         return eqx.combine(dyn, static)
+
+    @eqx.filter_jit
+    def posterior_predictive(
+        self, func: Callable[[M], Array], n: int, 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):
+            # Reconstruct model
+            model: M = self._unflatten(draw)
+
+            # Evaluate
+            return func(model)
+
+        return evaluate(draws)

bayinx/dists/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.dists import censored, gamma2, normal, posnormal
+
+__all__ = ['censored', "gamma2", "normal", "posnormal"]

bayinx/dists/censored/__init__.py

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

bayinx/dists/censored/gamma2/__init__.py

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

bayinx/dists/censored/gamma2/r.py

@@ -10,7 +10,7 @@ def prob(
     x: Float[ArrayLike, "..."],
     mu: Float[ArrayLike, "..."],
     nu: Float[ArrayLike, "..."],
-    censor: Float[ArrayLike, "..."]
+    censor: Float[ArrayLike, "..."],
 ) -> Float[Array, "..."]:
     """
     The mixed probability mass/density function (PMF/PDF) for a (mean-inverse dispersion parameterized) Gamma distribution.
@@ -19,19 +19,20 @@ def prob(
     - `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`, and `nu`.
+    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
+    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
+    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
+    # 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
+    evals = jnp.where(censored, gammaincc(nu, x * nu / mu), evals)  # pyright: ignore
 
     return evals
 
@@ -40,7 +41,7 @@ def logprob(
     x: Float[ArrayLike, "..."],
     mu: Float[ArrayLike, "..."],
     nu: Float[ArrayLike, "..."],
-    censor: 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.
@@ -49,17 +50,19 @@ def logprob(
     - `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`, and `nu`.
+    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
+    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
+    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
+    evals = jnp.where(censored, lax.log(gammaincc(nu, x * nu / mu)), evals)  # pyright: ignore
 
     return evals

bayinx/dists/censored/posnormal/__init__.py

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

bayinx/dists/censored/posnormal/r.py

@@ -0,0 +1,78 @@
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Float
+
+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

bayinx/dists/gamma2.py

@@ -1,4 +1,5 @@
 import jax.lax as lax
+import jax.numpy as jnp
 from jax.scipy.special import gammaln
 from jaxtyping import Array, ArrayLike, Float
 
@@ -17,6 +18,8 @@ def prob(
     # 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))
 
@@ -35,5 +38,12 @@ def logprob(
     # 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
+    # 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

bayinx/dists/normal.py

@@ -1,80 +1,138 @@
 import jax.lax as lax
+import jax.numpy as jnp
+import jax.scipy.special as jss
 from jaxtyping import Array, ArrayLike, Float
 
 __PI = 3.141592653589793
 
 
 def prob(
-    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[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.
-    - `sigma`:  The positive standard deviation.
+    - `x`: Where to evaluate the PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
 
     # Returns
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(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)) / (sigma * lax.sqrt(2.0 * __PI))
 
 
 def logprob(
-    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[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.
 
     # Parameters
-    - `x`:      Value(s) at which to evaluate the log PDF.
-    - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `x`: Where to evaluate the log PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
 
     # Returns
     The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(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)
 
 
 def uprob(
-    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
 ) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Normal distribution.
 
     # Parameters
-    - `x`:      Value(s) at which to evaluate the uPDF.
-    - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The positive standard deviation parameter(s).
+    - `x`: Where to evaluate the PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
 
     # Returns
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(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
 
 
 def ulogprob(
-    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[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.
 
     # Parameters
-    - `x`:      Value(s) at which to evaluate the log uPDF.
-    - `mu`:     The mean/location parameter(s).
-    - `sigma`:  The non-negative standard deviation parameter(s).
+    - `x`: Where to evaluate the PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
 
     # Returns
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(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)
+
+
+def cdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jss.ndtr((x - mu) / sigma)
+
+
+def logcdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jss.log_ndtr((x - mu) / sigma)
+
+
+def ccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jss.ndtr((mu - x) / sigma)
+
+
+def logccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jss.log_ndtr((mu - x) / sigma)

bayinx/dists/posnormal.py

@@ -0,0 +1,260 @@
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Float
+
+from bayinx.dists import normal
+
+
+def prob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The probability density function (PDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    # Evaluate PDF
+    evals = jnp.where(
+        non_negative,
+        normal.prob(x, mu, sigma) / normal.cdf(mu / sigma, 0.0, 1.0),
+        jnp.asarray(0.0),
+    )
+
+    return evals
+
+
+def logprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    # Evaluate log PDF
+    evals = jnp.where(
+        non_negative,
+        normal.logprob(x, mu, sigma) - normal.logcdf(mu / sigma, 0.0, 1.0),
+        -jnp.inf,
+    )
+
+    return evals
+
+
+def uprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The unnormalized probability density function (uPDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the uPDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    # Evaluate PDF
+    evals = jnp.where(non_negative, normal.prob(x, mu, sigma), jnp.asarray(0.0))
+
+    return evals
+
+
+def ulogprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The log of the unnormalized probability density function (log uPDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log uPDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    # Evaluate log PDF
+    evals = jnp.where(non_negative, normal.logprob(x, mu, sigma), -jnp.inf)
+
+    return evals
+
+
+def cdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The cumulative density function (CDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the CDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    # Compute intermediates
+    A: Array = normal.cdf(x, mu, sigma)
+    B: Array = normal.cdf(-mu / sigma, 0.0, 1.0)
+    C: Array = normal.cdf(mu / sigma, 0.0, 1.0)
+
+    # Evaluate CDF
+    evals = jnp.where(non_negative, (A - B) / C, jnp.asarray(0.0))
+
+    return evals
+
+
+# TODO: make numerically stable
+def logcdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The log-transformed cumulative density function (log CDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log CDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.asarray(0.0) <= x
+
+    A: Array = normal.logcdf(x, mu, sigma)
+    B: Array = normal.logcdf(-mu / sigma, 0.0, 1.0)
+    C: Array = normal.logcdf(mu / sigma, 0.0, 1.0)
+
+    # Evaluate log CDF
+    evals = jnp.where(non_negative, A + jnp.log1p(-jnp.exp(B - A)) - C, -jnp.inf)
+
+    return evals
+
+
+def ccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the cCDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The cCDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to arrays
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = 0.0 <= x
+
+    # Compute intermediates
+    A: Array = normal.cdf(-x, -mu, sigma)
+    B: Array = normal.cdf(mu / sigma, 0.0, 1.0)
+
+    # Evaluate cCDF
+    evals = jnp.where(non_negative, A / B, jnp.asarray(1.0))
+
+    return evals
+
+
+def logccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The log-transformed complementary cumulative density function (log cCDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log cCDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The log cCDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to arrays
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = 0.0 <= x
+
+    # Compute intermediates
+    A: Array = normal.logcdf(-x, -mu, sigma)
+    B: Array = normal.logcdf(mu / sigma, 0.0, 1.0)
+
+    # Evaluate log cCDF
+    evals = jnp.where(non_negative, A - B, jnp.asarray(0.0))
+
+    return evals

bayinx/mhx/vi/__init__.py

@@ -1,3 +1,5 @@
-from bayinx.mhx.vi.meanfield import MeanField as MeanField
-from bayinx.mhx.vi.normalizing_flow import NormalizingFlow as NormalizingFlow
-from bayinx.mhx.vi.standard import Standard as Standard
+from bayinx.mhx.vi.meanfield import MeanField
+from bayinx.mhx.vi.normalizing_flow import NormalizingFlow
+from bayinx.mhx.vi.standard import Standard
+
+__all__ = ['MeanField', 'NormalizingFlow', 'Standard']

bayinx/mhx/vi/meanfield.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Self
+from typing import Any, Dict, Generic, Self, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -10,8 +10,8 @@ from jaxtyping import Array, Float, Key, Scalar
 from bayinx.core import Model, Variational
 from bayinx.dists import normal
 
-
-class MeanField(Variational):
+M = TypeVar('M', bound=Model)
+class MeanField(Variational, Generic[M]):
     """
     A fully factorized Gaussian approximation to a posterior distribution.
 
@@ -19,9 +19,9 @@ class MeanField(Variational):
     - `var_params`: The variational parameters for the approximation.
     """
 
-    var_params: Dict[str, Float[Array, "..."]]
+    var_params: Dict[str, Float[Array, "..."]] #todo: just expand to attributes
 
-    def __init__(self, model: Model):
+    def __init__(self, model: M):
         """
         Constructs an unoptimized meanfield posterior approximation.
 
@@ -55,7 +55,6 @@ class MeanField(Variational):
 
         return filter_spec
 
-
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws

bayinx/mhx/vi/normalizing_flow.py

@@ -1,4 +1,4 @@
-from typing import Any, Self, Tuple
+from typing import Any, Generic, Self, Tuple, TypeVar
 
 import equinox as eqx
 import jax.flatten_util as jfu
@@ -9,8 +9,8 @@ from jaxtyping import Array, Key, Scalar
 
 from bayinx.core import Flow, Model, Variational
 
-
-class NormalizingFlow(Variational):
+M = TypeVar('M', bound=Model)
+class NormalizingFlow(Variational, Generic[M]):
     """
     An ordered collection of diffeomorphisms that map a base distribution to a
     normalized approximation of a posterior distribution.
@@ -23,7 +23,7 @@ class NormalizingFlow(Variational):
     flows: list[Flow]
     base: Variational
 
-    def __init__(self, base: Variational, flows: list[Flow], model: Model):
+    def __init__(self, base: Variational, flows: list[Flow], model: M):
         """
         Constructs an unoptimized normalizing flow posterior approximation.
 

bayinx/mhx/vi/standard.py

@@ -1,29 +1,25 @@
-from typing import Callable
 
 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
+from jaxtyping import Array, Key
 
-from bayinx.core import Model, Variational
+from bayinx.core._variational import M, Variational
 from bayinx.dists import normal
 
 
-class Standard(Variational):
+class Standard(Variational[M]):
     """
     A standard normal approximation to a posterior distribution.
 
     # Attributes
     - `dim`: Dimension of the parameter space.
     """
-
     dim: int
-    _unflatten: Callable[[Float[Array, "..."]], Model]
-    _constraints: Model
 
-    def __init__(self, model: Model):
+    def __init__(self, model: M):
         """
         Constructs a standard normal approximation to a posterior distribution.
 

{bayinx-0.3.4.dist-info → bayinx-0.3.6.dist-info}/METADATA

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.4
+Version: 0.3.6
 Summary: Bayesian Inference with JAX
+License-File: LICENSE
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12
 Requires-Dist: jax>=0.4.38

bayinx-0.3.6.dist-info/RECORD

@@ -0,0 +1,35 @@
+bayinx/__init__.py,sha256=TM-aoRaPX6jSYtCM7Jv59TPV-H6bcDk1-VMttYP1KME,99
+bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/constraints/__init__.py,sha256=PiWXZKi7YdbTMKvw-OE5f-t87jJT893uAFrwWWBfOdg,64
+bayinx/constraints/lower.py,sha256=30y0l6PF-tbS9LR_tto9AvwmsvXq1ExU-v8DLrJD4g4,1446
+bayinx/core/__init__.py,sha256=bZvQITgW0DWuPKl3wCLKt6WHKogYKx8Zz36g8z9Aung,253
+bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
+bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
+bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
+bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
+bayinx/core/_variational.py,sha256=szm1WuUh_3pxzFfQy92TR4p2Sk-fR6rO-4-LrJMeVGI,5356
+bayinx/dists/__init__.py,sha256=9DdPea7HAnBOzaV_4gM5noPX8YCb_p06d8PJvGfFy3Y,118
+bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
+bayinx/dists/gamma2.py,sha256=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
+bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
+bayinx/dists/posnormal.py,sha256=w9plA1EctXwXOiY0doc4ZndjnwptbEZBHHCGdc4gviY,7292
+bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
+bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
+bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
+bayinx/dists/censored/gamma2/r.py,sha256=dKAOYstufwgDwibQZHrJxA1d2gawj-7K3IkaCRCzNTg,2446
+bayinx/dists/censored/posnormal/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
+bayinx/dists/censored/posnormal/r.py,sha256=hyuNR3HZY-Tgtso-WwjcZT6Ejxfyax_VKwIvVix44Jc,2362
+bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+bayinx/mhx/vi/__init__.py,sha256=2woNB5oZxfs8pZCkOfzriGahRFLzkLdkTj8_keTN0I0,205
+bayinx/mhx/vi/meanfield.py,sha256=Z7kGQAyp5iB8rEdjbwAbVTFH4GwxlTKDZFbdJ-FN5Vs,3739
+bayinx/mhx/vi/normalizing_flow.py,sha256=8pLMDdZPIt5wlgbhHWSFY1ChSWM9pvSD2bQx3zgz1F8,4710
+bayinx/mhx/vi/standard.py,sha256=W-ZvigJkUpqVlREgiFm9io8ansT1XpZwq5AqSmdv--E,1578
+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.6.dist-info/METADATA,sha256=WEdMVyISWGgK0KJvuSlkpbObsxiVfGvIxky7OsuYdXg,3079
+bayinx-0.3.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.6.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.6.dist-info/RECORD,,

bayinx-0.3.6.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Todd McCready
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bayinx-0.3.4.dist-info/RECORD

@@ -1,31 +0,0 @@
-bayinx/__init__.py,sha256=5fb_tGeEVnrNt6IQqu7gZaJskBJHqjcg08JRPrY2ANo,139
-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=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
-bayinx/core/model.py,sha256=1vQPVjE0ebCdW7mLuabgQcCTi95o8n8CC6GuzJdNL1s,2956
-bayinx/core/parameter.py,sha256=eECqvfMNWSU8_CkGYaAfOCneMMQGZI21kF0mErsh2Rc,1080
-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/__init__.py,sha256=2EaQcgCXEwaRoHChVlD02ZMfgiwQAqey6uLPov1lcwE,21
-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.4.dist-info/METADATA,sha256=EpVIXPifXNloZfCCWNuNaVhWO_dMEujN3V_kVZz2Q6Y,3057
-bayinx-0.3.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.4.dist-info/RECORD,,