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

bayinx/dists/bernoulli.py

@@ -1,33 +1,202 @@
+from typing import Tuple
+
+import equinox as eqx
 import jax.lax as lax
-from jaxtyping import Array, ArrayLike, Real, UInt
+import jax.numpy as jnp
+import jax.random as jr
+import jax.tree as jt
+from jaxtyping import Array, ArrayLike, Float, Integer, PRNGKeyArray, Real, Scalar
+
+from bayinx.core.distribution import Distribution
+from bayinx.core.node import Node
+from bayinx.nodes import Observed
 
 
-# MARK: Functions ----
-def prob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+def prob(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
     """
     The probability mass function (PMF) for a Bernoulli distribution.
 
     # Parameters
-    - `x`:      Value(s) at which to evaluate the PDF.
-    - `p`:     The probability parameter(s).
+    - `x`: Where to evaluate the PMF (pst be 0 or 1).
+    - `p`: The probability of success (p).
 
     # Returns
     The PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
     """
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
 
-    return lax.pow(p, x) * lax.pow(1 - p, 1 - x)
+    # Bernoulli PMF: p^x * (1-p)^(1-x)
+    return lax.exp(x * lax.log(p) + (1.0 - x) * lax.log1p(-p))
 
 
-def logprob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+def logprob(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
     """
-    The log probability mass function (log PMF) for a Bernoulli distribution.
+    The log of the probability mass function (log PMF) for a Bernoulli distribution.
 
     # Parameters
-    - `x`:  Value(s) at which to evaluate the log PMF.
-    - `p`:  The probability parameter(s).
+    - `x`: Where to evaluate the log PMF (pst be 0 or 1).
+    - `p`: The probability of success (p).
 
     # Returns
     The log PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
     """
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
+
+    return x * lax.log(p) + (1.0 - x) * lax.log(1.0 - p)
+
+
+def cdf(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
+    """
+    The cuplative distribution function (CDF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the CDF.
+    - `p`: The probability of success (p).
+
+    # Returns
+    The CDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
+
+    return jnp.where(
+        x < 0.0,
+        0.0,
+        jnp.where(x < 1.0, 1.0 - p, 1.0)
+    )
+
+
+def logcdf(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the cuplative distribution function (log CDF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log CDF.
+    - `p`: The probability of success (p).
+
+    # Returns
+    The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
+
+    return jnp.where(
+        x < 0.0,
+        -jnp.inf,
+        jnp.where(x < 1.0, lax.log1p(-p), 0.0)
+    )
+
+
+def ccdf(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
+    """
+    The complementary cuplative distribution function (cCDF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the cCDF.
+    - `p`: The probability of success (p).
+
+    # Returns
+    The cCDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
+
+    return jnp.where(
+        x < 0.0,
+        1.0,
+        jnp.where(x < 1.0, p, 0.0)
+    )
+
+
+def logccdf(
+    x: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the complementary cuplative distribution function (log cCDF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log cCDF.
+    - `p`: The probability of success (p).
+
+    # Returns
+    The log cCDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+    # Cast to Array
+    x, p = jnp.asarray(x), jnp.asarray(p)
+
+    return jnp.where(
+        x < 0.0,
+        0.0,
+        jnp.where(x < 1.0, lax.log(p), -jnp.inf)
+    )
+
+
+class Bernoulli(Distribution):
+    """
+    A Bernoulli distribution.
+
+    # Attributes
+    - `p`: The probability of success (p).
+    """
+
+    p: Node[Real[Array, "..."]]
+
+
+    def __init__(
+        self,
+        p: Real[ArrayLike, "..."] | Node[Real[ArrayLike, "..."]],
+    ):
+        # Initialize probability of success parameter (p)
+        if isinstance(p, Node):
+            if isinstance(p.obj, ArrayLike):
+                self.p = p # type: ignore
+        else:
+            self.p = Observed(jnp.asarray(p))
+
+
+    def logprob(self, node: Node) -> Scalar:
+        obj, p = (
+            node.obj,
+            self.p.obj,
+        )
+
+        # Filter out irrelevant values
+        obj, _ = eqx.partition(obj, node._filter_spec)
+
+        # Helper function for the single-leaf log-probability evaluation
+        def leaf_logprob(x: Float[ArrayLike, ""]) -> Scalar:
+            return logprob(x, p).sum()
+
+        # Compute log probabilities across leaves
+        eval_obj = jt.map(leaf_logprob, obj)
+
+        # Compute total sum
+        total = jt.reduce_associative(lambda x,y: x + y, eval_obj, identity=0.0)
+
+        return jnp.asarray(total)
+
+    def sample(self, shape: int | Tuple[int, ...], key: PRNGKeyArray = jr.key(0)):
+        # Coerce to tuple
+        if isinstance(shape, int):
+            shape = (shape, )
 
-    return lax.log(p) * x + (1 - x) * lax.log(1 - p)
+        return jr.bernoulli(key, self.p.obj, shape)

bayinx/dists/binomial.py

@@ -0,0 +1,215 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import jax.scipy.special as jsp
+import jax.tree as jt
+from jaxtyping import Array, ArrayLike, Float, Integer, PRNGKeyArray, Real, Scalar
+
+from bayinx.core.distribution import Distribution
+from bayinx.core.node import Node
+from bayinx.nodes import Observed
+
+
+def log_binom_coeff(n: ArrayLike, x: ArrayLike) -> Array:
+    n, x = jnp.asarray(n), jnp.asarray(x)
+    return jsp.gammaln(n + 1) - jsp.gammaln(x + 1) - jsp.gammaln(n - x + 1)
+
+
+def prob(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The probability mass function (PMF) for a Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the PMF (number of successes).
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The PMF evaluated at `x`. The output will have the broadcasted shapes of `x`, `n`, and `p`.
+    """
+    # Cast to Array
+    x, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return lax.exp(logprob(x, n, p))
+
+
+def logprob(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the probability mass function (log PMF) for a Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PMF (number of successes).
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The log PMF evaluated at `x`. The output will have the broadcasted shapes of `x`, `n`, and `p`.
+    """
+    # Cast to Array
+    k, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return log_binom_coeff(n, k) + k * lax.log(p) + (n - k) * lax.log1p(-p)
+
+
+def cdf(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The cumulative density function (CDF) for a Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the CDF (P(X <= x)).
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `n`, and `p`.
+    """
+    # Cast to Array
+    x, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return jsp.betainc(n - x, x + 1, 1.0 - p)
+
+
+def logcdf(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the cumulative density function (log CDF) for a Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log CDF (ln P(X <= x)).
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `n`, and `p`.
+    """
+    # Cast to Array
+    x, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return lax.log(cdf(x, n, p))
+
+
+def ccdf(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for a Binomial distribution (P(X > k)).
+
+    # Parameters
+    - `k`: Where to evaluate the cCDF.
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The cCDF evaluated at `k`. The output will have the broadcasted shapes of `k`, `n`, and `p`.
+    """
+
+    # Cast to Array
+    x, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return jsp.betainc(x + 1, n - x, p)
+
+
+def logccdf(
+    x: Integer[ArrayLike, "..."],
+    n: Integer[ArrayLike, "..."],
+    p: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the complementary cumulative density function (log cCDF) for a Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log cCDF.
+    - `n`: The number of trials.
+    - `p`: The probability of success.
+
+    # Returns
+    The log cCDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `n`, and `p`.
+    """
+    # Cast to Array
+    x, n, p = jnp.asarray(x), jnp.asarray(n), jnp.asarray(p)
+
+    return lax.log(ccdf(x, n, p))
+
+
+class Binomial(Distribution):
+    """
+    A Binomial distribution.
+
+    # Attributes
+    - `n`: The number of trials parameter (integer).
+    - `p`: The probability of success parameter (float, [0, 1]).
+    """
+
+    n: Node[Integer[Array, "..."]]
+    p: Node[Real[Array, "..."]]
+
+
+    def __init__(
+        self,
+        n: Integer[ArrayLike, "..."] | Node[Integer[Array, "..."]],
+        p: Real[ArrayLike, "..."] | Node[Real[Array, "..."]]
+    ):
+        # Initialize number of trials (n)
+        if isinstance(n, Node):
+            if isinstance(n.obj, ArrayLike):
+                self.n = n # type: ignore
+        else:
+            self.n = Observed(jnp.asarray(n))
+
+        # Initialize probability of success parameter (p)
+        if isinstance(p, Node):
+            if isinstance(p.obj, ArrayLike):
+                self.p = p # type: ignore
+        else:
+            self.p = Observed(jnp.asarray(p))
+
+    def logprob(self, node: Node) -> Scalar:
+        obj, n, p = (
+            node.obj,
+            self.n.obj,
+            self.p.obj
+        )
+
+        # Filter out irrelevant values
+        obj, _ = eqx.partition(obj, node._filter_spec)
+
+        # Helper function for the single-leaf log-probability evaluation
+        def leaf_logprob(k: Float[ArrayLike, "..."]) -> Scalar:
+            return logprob(k, n, p).sum()
+
+        # Compute log probabilities across leaves
+        eval_obj = jt.map(leaf_logprob, obj)
+
+        # Compute total sum
+        total = jt.reduce_associative(lambda x,y: x + y, eval_obj, identity=0.0)
+
+        return jnp.asarray(total)
+
+    def sample(self, shape: int | Tuple[int, ...], key: PRNGKeyArray = jr.key(0)):
+        # Coerce to tuple
+        if isinstance(shape, int):
+            shape = (shape, )
+
+        # Use jr.binomial for sampling (returns integer array)
+        # Note: jr.binomial accepts n as int, p as float, and shape
+        return jr.binomial(key, self.n.obj, self.p.obj, shape=shape)

bayinx/dists/exponential.py

@@ -0,0 +1,211 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import jax.tree as jt
+from jaxtyping import Array, ArrayLike, PRNGKeyArray, Real, Scalar
+
+from bayinx.core.distribution import Distribution
+from bayinx.core.node import Node
+from bayinx.nodes import Observed
+
+
+def prob(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The probability density function (PDF) for an Exponential distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the PDF (must be >= 0).
+    - `lam`: The rate parameter (lambda), must be > 0.
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    result = lam * lax.exp(-lam * x)
+
+    # Handle values outside of support
+    result = lax.select(x >= 0, result, jnp.array(0.0))
+
+    return result
+
+
+def logprob(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for an Exponential distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PDF (must be >= 0).
+    - `lam`: The rate parameter (lambda), must be > 0.
+
+    # Returns
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    result = lax.log(lam) - lam * x
+
+    # Handle values outside of support
+    result = lax.select(x >= 0, result, -jnp.inf)
+
+    return result
+
+def cdf(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The cumulative density function (CDF) for an Exponential distribution (P(X <= x)).
+
+    # Parameters
+    - `x`: Where to evaluate the CDF.
+    - `lam`: The rate parameter (lambda).
+
+    # Returns
+    The CDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    result = 1.0 - lax.exp(-lam * x)
+
+    # Handle values outside of support
+    result = lax.select(x >= 0, result, jnp.array(0.0))
+
+    return result
+
+def logcdf(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the cumulative density function (log CDF) for an Exponential distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log CDF (ln P(X <= x)).
+    - `lam`: The rate parameter (lambda).
+
+    # Returns
+    The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    result = lax.log1p(-lax.exp(-lam * x))
+
+    # Handle values outside of support (x < 0)
+    result = lax.select(x >= 0, result, -jnp.inf)
+
+    return result
+
+
+def ccdf(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for an Exponential distribution (P(X > x)).
+
+    # Parameters
+    - `x`: Where to evaluate the cCDF.
+    - `lam`: The rate parameter (lambda).
+
+    # Returns
+    The cCDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    result = lax.exp(-lam * x)
+
+    # Handle values outside of support
+    result = lax.select(x >= 0, result, jnp.array(1.0))
+
+    return result
+
+
+def logccdf(
+    x: Real[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the complementary cumulative density function (log cCDF) for an Exponential distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log cCDF.
+    - `lam`: The rate parameter (lambda).
+
+    # Returns
+    The log cCDF evaluated at `x`. The output will have the broadcasted shapes of `x` and `lam`.
+    """
+    # Cast to Array
+    x, lam = jnp.asarray(x), jnp.asarray(lam)
+
+    # log(cCDF(x)) = -lambda * x for x >= 0
+    log_ccdf_val = -lam * x
+
+    # Handle values outside of support (x < 0), log(P(X > x)) = log(1.0) = 0.0
+    return lax.select(x >= 0, log_ccdf_val, jnp.array(0.0))
+
+
+class Exponential(Distribution):
+    """
+    An Exponential distribution.
+
+    # Attributes
+    - `lam`: The rate parameter (lambda), must be positive.
+    """
+
+    lam: Node[Real[Array, "..."]]
+
+
+    def __init__(
+        self,
+        lam: Real[ArrayLike, "..."] | Node[Real[Array, "..."]],
+    ):
+        # Initialize rate parameter (lambda)
+        if isinstance(lam, Node):
+            if isinstance(lam.obj, ArrayLike):
+                self.lam = lam # type: ignore
+        else:
+            self.lam = Observed(jnp.asarray(lam))
+
+
+    def logprob(self, node: Node) -> Scalar:
+        obj, lam = (
+            node.obj,
+            self.lam.obj,
+        )
+
+        # Filter out irrelevant values
+        obj, _ = eqx.partition(obj, node._filter_spec)
+
+        # Helper function for the single-leaf log-probability evaluation
+        def leaf_logprob(x: Real[ArrayLike, "..."]) -> Scalar:
+            return logprob(x, lam).sum()
+
+        # Compute log probabilities across leaves
+        eval_obj = jt.map(leaf_logprob, obj)
+
+        # Compute total sum
+        total = jt.reduce_associative(lambda x,y: x + y, eval_obj, identity=0.0)
+
+        return jnp.asarray(total)
+
+    def sample(self, shape: int | Tuple[int, ...], key: PRNGKeyArray = jr.key(0)):
+        # Coerce to tuple
+        if isinstance(shape, int):
+            shape = (shape, )
+
+        return jr.exponential(key, shape=shape) *  1.0 / self.lam.obj