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

bayinx/dists/normal.py

@@ -1,16 +1,25 @@
+from typing import Tuple
+
+import equinox as eqx
 import jax.lax as lax
 import jax.numpy as jnp
-import jax.scipy.special as jss
-from jaxtyping import Array, ArrayLike, Float
+import jax.random as jr
+import jax.scipy.special as jsp
+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
 
-__PI = 3.141592653589793
+PI = 3.141592653589793
 
 
 def prob(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
     """
     The probability density function (PDF) for a Normal distribution.
 
@@ -25,14 +34,14 @@ def prob(
     # 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 * 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, "..."],
-) -> Float[Array, "..."]:
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Normal distribution.
 
@@ -47,92 +56,155 @@ def logprob(
     # 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)
+    return -lax.log(lax.sqrt(2.0 * PI)) - lax.log(sigma) - 0.5 * lax.square((x - mu) / sigma)
 
 
-def uprob(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
+def cdf(
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
     """
-    The unnormalized probability density function (uPDF) for a Normal distribution.
+    The cumulative density function (CDF) for a Normal distribution.
 
     # Parameters
-    - `x`: Where to evaluate the PDF.
+    - `x`: Where to evaluate the CDF.
     - `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`.
+    The CDF 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
+    return jsp.ndtr((x - mu) / sigma)
 
 
-def ulogprob(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
+def logcdf(
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
     """
-    The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
+    The log of the cumulative density function (log CDF) for a Normal distribution.
 
     # Parameters
-    - `x`: Where to evaluate the PDF.
+    - `x`: Where to evaluate the log CDF.
     - `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`.
+    The log CDF 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)
+    return jsp.log_ndtr((x - mu) / sigma)
 
 
-def cdf(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
+def ccdf(
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for a 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 Array
     x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
 
-    return jss.ndtr((x - mu) / sigma)
+    return jsp.ndtr((mu - x) / sigma)
 
 
-def logcdf(
-    x: Float[ArrayLike, "..."],
-    mu: Float[ArrayLike, "..."],
-    sigma: Float[ArrayLike, "..."],
-) -> Float[Array, "..."]:
+def logccdf(
+    x: Real[ArrayLike, "..."],
+    mu: Real[ArrayLike, "..."],
+    sigma: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the complementary cumulative density function (log cCDF) for a 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 Array
     x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
 
-    return jss.log_ndtr((x - mu) / sigma)
+    return jsp.log_ndtr((mu - x) / 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)
+class Normal(Distribution):
+    """
+    A normal distribution.
 
-    return jss.ndtr((mu - x) / sigma)
+    # Attributes
+    - `mu`: The mean/location parameter.
+    - `sigma`: The standard-deviation/scale parameter.
+    """
 
+    mu: Node[Real[Array, "..."]]
+    sigma: Node[Real[Array, "..."]]
 
-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)
+    def __init__(
+        self,
+        mu: Real[ArrayLike, "..."] | Node[Real[Array, "..."]],
+        sigma: Real[ArrayLike, "..."] | Node[Real[Array, "..."]]
+    ):
+        # Initialize mean/location parameter (mu)
+        if isinstance(mu, Node):
+            if isinstance(mu.obj, ArrayLike):
+                self.mu = mu # type: ignore
+        else:
+            self.mu = Observed(jnp.asarray(mu))
+
+        # Initialize dispersion/scale parameter (sigma)
+        if isinstance(sigma, Node):
+            if isinstance(sigma.obj, ArrayLike):
+                self.sigma = sigma # type: ignore
+        else:
+            self.sigma = Observed(jnp.asarray(sigma))
+
+
+    def logprob(self, node: Node) -> Scalar:
+        obj, mu, sigma = (
+            node.obj,
+            self.mu.obj,
+            self.sigma.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, mu, sigma).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.normal(key, shape) * self.sigma.obj + self.mu.obj

bayinx/dists/poisson.py

@@ -0,0 +1,203 @@
+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, Integer, PRNGKeyArray, Real, Scalar
+
+from bayinx.core.distribution import Distribution
+from bayinx.core.node import Node
+from bayinx.nodes import Observed
+
+
+def prob(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The probability mass function (PMF) for a Poisson distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the PMF.
+    - `lam`: The rate parameter (lambda), representing the average number of events.
+
+    # Returns
+    The PMF 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)
+
+    return lax.exp(logprob(x, lam))
+
+
+def logprob(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the probability mass function (log PMF) for a Poisson distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PMF.
+    - `lam`: The rate parameter (lambda), representing the average number of events.
+
+    # Returns
+    The log PMF 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)
+
+    return x * lax.log(lam) - lam - jsp.gammaln(x + 1)
+
+
+def cdf(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The cumulative density function (CDF) for a Poisson 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 = jsp.gammainc(x + 1.0, lam)
+    result = lax.select(x < 0, jnp.array(0.0), result)
+
+    return result
+
+
+def logcdf(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the cumulative density function (log CDF) for a Poisson 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.log(jsp.gammainc(x + 1.0, lam))
+
+    # Handle values outside of support
+    result = lax.select(x < 0, -jnp.inf, result)
+
+    return result
+
+
+def ccdf(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for a Poisson 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 = jsp.gammaincc(x + 1.0, lam)
+
+    # Handle values outside of support
+    result = lax.select(x < 0, jnp.array(1.0), result)
+
+    return result
+
+
+def logccdf(
+    x: Integer[ArrayLike, "..."],
+    lam: Real[ArrayLike, "..."],
+) -> Real[Array, "..."]:
+    """
+    The log of the complementary cumulative density function (log cCDF) for a Poisson 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)
+
+    result = lax.log(jsp.gammaincc(x + 1.0, lam))
+
+    # Handle values outside of support
+    result = lax.select(x < 0, jnp.array(0.0), result)
+
+    return result
+
+
+class Poisson(Distribution):
+    """
+    A Poisson distribution.
+
+    # Attributes
+    - `lam`: The rate parameter (lambda), representing the average number of events in an interval.
+    """
+
+    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: Integer[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.poisson(key, self.lam.obj, shape=shape)

bayinx/flows/__init__.py

@@ -0,0 +1,5 @@
+from bayinx.flows.diagaffine import DiagAffine as DiagAffine
+from bayinx.flows.fullaffine import FullAffine as FullAffine
+from bayinx.flows.lowrankaffine import LowRankAffine as LowRankAffine
+from bayinx.flows.planar import Planar as Planar
+from bayinx.flows.sylvester import Sylvester as Sylvester

bayinx/flows/diagaffine.py

@@ -0,0 +1,120 @@
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jaxtyping import Array, Float, PRNGKeyArray, Scalar
+
+from bayinx.core.flow import FlowLayer, FlowSpec
+
+
+class DiagAffineLayer(FlowLayer):
+    """
+    A diagonal (element-wise) affine flow.
+
+    # Attributes
+    - `params`: The parameters of the diagonal affine flow.
+    - `constraints`: The constraining transformations for the parameters of the diagonal affine flow.
+    - `static`: Whether the flow layer is frozen (parameters are not subject to further optimization).
+    """
+
+    params: Dict[str, Array]
+    constraints: Dict[str, Callable[[Array], Array]]
+    static: bool
+
+    def __init__(self, dim: int, key: PRNGKeyArray):
+        """
+        Initializes a full affine flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.static = False
+        # Split key
+        k1, k2 = jr.split(key)
+
+        # Initialize parameters
+        self.params = {
+            "shift": jr.normal(k1, (dim,)) / dim**0.5,
+            "scale": jr.normal(k2, (dim,)) / dim**0.5,
+        }
+
+        # Define constraints
+        self.constraints = {"scale": jnp.exp}
+
+
+    def __forward(self, draw: Float[Array, " D"]) -> Float[Array, " D"]:
+        params = self.transform_params()
+
+        assert len(draw.shape) == 1
+
+        # Extract parameters
+        shift: Float[Array, " dim"] = params["shift"]
+        scale: Float[Array, " dim"] = params["scale"]
+
+        # Compute forward transformation
+        draw = draw * scale + shift
+
+        return draw
+
+    @eqx.filter_jit
+    def forward(self, draws: Float[Array, "draws dim"]) -> Float[Array, "draws dim"]:
+        f = jax.vmap(self.__forward, 0)
+        return f(draws)
+
+    def __adjust(self, draw: Float[Array, " dim"]) -> Float[Array, " dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        scale: Float[Array, " dim"] = params["scale"]
+
+        # Compute log-Jacobian adjustment
+        lja: Array = jnp.log(scale).sum()
+
+        assert lja.shape == ()
+
+        return lja
+
+    @eqx.filter_jit
+    def adjust(self, draws: Float[Array, "draws dim"]) -> Float[Array, "draws dim"]:
+        f = jax.vmap(self.__adjust, 0)
+        return f(draws)
+
+    def __forward_and_adjust(self, draw: Float[Array, " dim"]) -> Tuple[Float[Array, " dim"], Scalar]:
+        params = self.transform_params()
+
+        assert len(draw.shape) == 1
+
+        # Extract parameters
+        shift: Float[Array, " dim"] = params["shift"] # noqa
+        scale: Float[Array, " dim"] = params["scale"] # noqa
+
+        # Compute forward transformation
+        draw = scale * draw + shift
+
+        assert len(draw.shape) == 1
+
+        # Compute log-Jacobian adjustment
+        lja: Scalar = jnp.log(scale).sum()
+
+        assert lja.shape == ()
+
+        return draw, lja
+
+    @eqx.filter_jit
+    def forward_and_adjust(self, draws: Float[Array, "draws dim"]) -> Tuple[Float[Array, "draws dim"], Scalar]:
+        f = jax.vmap(self.__forward_and_adjust, 0)
+        return f(draws)
+
+
+class DiagAffine(FlowSpec):
+    """
+    A specification for the diagonal affine flow.
+    """
+    key: PRNGKeyArray
+    def __init__(self, key: PRNGKeyArray = jr.key(0)):
+        self.key = key
+
+    def construct(self, dim: int) -> DiagAffineLayer:
+        return DiagAffineLayer(dim, self.key)