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

bayinx/dists/posnormal.py

@@ -1,260 +0,0 @@
-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/dists/uniform.py

@@ -1,75 +0,0 @@
-import jax.lax as _lax
-import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float
-
-
-def prob(
-    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
-) -> Float[Array, "..."]:
-    """
-    The probability density function (PDF) for a Uniform distribution.
-
-    # Parameters
-    - `x`:      Value(s) at which to evaluate the PDF.
-    - `lb`:     The lower bound parameter(s).
-    - `ub`:     The upper bound parameter(s).
-
-    # Returns
-    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
-    """
-
-    return 1.0 / (ub - lb)  # pyright: ignore
-
-
-def logprob(
-    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
-) -> Float[Array, "..."]:
-    """
-    The log of the probability density function (log PDF) for a Uniform distribution.
-
-    # Parameters
-    - `x`:      Value(s) at which to evaluate the PDF.
-    - `lb`:     The lower bound parameter(s).
-    - `ub`:     The upper bound parameter(s).
-
-    # Returns
-    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
-
-
-def uprob(
-    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
-) -> Float[Array, "..."]:
-    """
-    The unnormalized probability density function (uPDF) for a Uniform distribution.
-
-    # Parameters
-    - `x`:      Value(s) at which to evaluate the PDF.
-    - `lb`:     The lower bound parameter(s).
-    - `ub`:     The upper bound parameter(s).
-
-    # Returns
-    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))
-
-
-def ulogprob(
-    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.
-
-    # Parameters
-    - `x`:      Value(s) at which to evaluate the PDF.
-    - `lb`:     The lower bound parameter(s).
-    - `ub`:     The upper bound parameter(s).
-
-    # Returns
-    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))

bayinx/mhx/__init__.py

@@ -1 +0,0 @@
-

bayinx/mhx/vi/__init__.py

@@ -1,5 +0,0 @@
-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/flows/__init__.py

@@ -1,3 +0,0 @@
-from bayinx.mhx.vi.flows.fullaffine import FullAffine as FullAffine
-from bayinx.mhx.vi.flows.planar import Planar as Planar
-from bayinx.mhx.vi.flows.radial import Radial as Radial

bayinx/mhx/vi/flows/fullaffine.py

@@ -1,75 +0,0 @@
-from functools import partial
-from typing import Tuple
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array, Scalar
-
-from bayinx.core import Flow
-
-
-class FullAffine(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.
-    """
-
-    def __init__(self, dim: int):
-        """
-        Initializes a full affine flow.
-
-        # Parameters
-        - `dim`: The dimension of the parameter space.
-        """
-        self.params = {
-            "shift": jnp.zeros(dim),
-            "scale": jnp.zeros((dim, dim)),
-        }
-
-        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.transform_params()
-
-        # Extract parameters
-        shift: Array = params["shift"]
-        scale: Array = params["scale"]
-
-        # Compute forward transformation
-        draws = draws @ scale + shift
-
-        return draws
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0))
-    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_params()
-
-        # Extract parameters
-        shift: Array = params["shift"]
-        scale: Array = params["scale"]
-
-        # Compute forward transformation
-        draws = draws @ scale + shift
-
-        # Compute laj
-        laj: Scalar = jnp.log(jnp.diag(scale)).sum()
-
-        return draws, laj

bayinx/mhx/vi/flows/planar.py

@@ -1,74 +0,0 @@
-from functools import partial
-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, Scalar
-
-from bayinx.core import Flow
-
-
-class Planar(Flow):
-    """
-    A planar flow.
-
-    # Attributes
-    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
-    - `constraints`: A dictionary of constraining transformations.
-    """
-
-    params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Array], Array]]
-
-    def __init__(self, dim: int, key=jr.PRNGKey(0)):
-        """
-        Initializes a planar flow.
-
-        # Parameters
-        - `dim`: The dimension of the parameter space.
-        """
-        self.params = {
-            "u": jnp.zeros(dim),
-            "w": jnp.ones(dim),
-            "b": jnp.zeros(1),
-        }
-        self.constraints = {}
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0))
-    def forward(self, draws: Array) -> Array:
-        params = self.transform_params()
-
-        # Extract parameters
-        w: Array = params["w"]
-        u: Array = params["u"]
-        b: Array = params["b"]
-
-        # Compute forward transformation
-        draws = draws + u * jnp.tanh(draws.dot(w) + b)
-
-        return draws
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0))
-    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_params()
-
-        # Extract parameters
-        w: Array = params["w"]
-        u: Array = params["u"]
-        b: Array = params["b"]
-
-        # Compute shared intermediates
-        x: Array = draws.dot(w) + b
-
-        # Compute forward transformation
-        draws = draws + u * jnp.tanh(x)
-
-        # Compute laj
-        h_prime: Scalar = 1.0 - jnp.square(jnp.tanh(x))
-        laj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
-
-        return draws, laj

bayinx/mhx/vi/flows/radial.py

@@ -1,94 +0,0 @@
-from functools import partial
-from typing import Callable, Dict, Tuple
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-import jax.random as jr
-from jax.numpy.linalg import norm
-from jaxtyping import Array, Float, Scalar
-
-from bayinx.core import Flow
-
-
-class Radial(Flow):
-    """
-    A radial flow.
-
-    # Attributes
-    - `params`: A dictionary containing the JAX Arrays representing the flow 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, key=jr.PRNGKey(0)):
-        """
-        Initializes a planar flow.
-
-        # Parameters
-        - `dim`: The dimension of the parameter space.
-        """
-        self.params = {
-            "alpha": jnp.array(1.0),
-            "beta": jnp.array(1.0),
-            "center": jnp.ones(dim),
-        }
-        self.constraints = {"beta": jnp.exp}
-
-    @partial(jax.vmap, in_axes=(None, 0))
-    @eqx.filter_jit
-    def forward(self, draws: Array) -> Array:
-        """
-        Applies the forward radial transformation for each draw.
-
-        # Parameters
-        - `draws`: Draws from some layer of a normalizing flow.
-
-        # Returns
-        The transformed samples.
-        """
-        params = self.transform_params()
-
-        # Extract parameters
-        alpha = params["alpha"]
-        beta = params["beta"]
-        center = params["center"]
-
-        # Compute distance to center per-draw
-        r: Array = norm(draws - center)
-
-        # Apply forward transformation
-        draws = draws + (beta / (alpha + r)) * (draws - center)
-
-        return draws
-
-    @partial(jax.vmap, in_axes=(None, 0))
-    @eqx.filter_jit
-    def adjust_density(self, draws: Array) -> Tuple[Array, Scalar]:
-        params = self.transform_params()
-
-        # Extract parameters
-        alpha = params["alpha"]
-        beta = params["beta"]
-        center = params["center"]
-
-        # Compute distance to center per-draw
-        r: Array = norm(draws - center)
-
-        # Compute shared intermediates
-        x: Array = beta / (alpha + r)
-
-        # Apply forward transformation
-        draws = draws + (x) * (draws - center)
-
-        # Compute density adjustment
-        laj = jnp.log(
-            jnp.abs(
-                (1.0 + alpha * beta / (alpha + r) ** 2.0)
-                * (1.0 + x) ** (center.size - 1.0)
-            )
-        )
-
-        return draws, laj

bayinx/mhx/vi/flows/sylvester.py

@@ -1,19 +0,0 @@
-from typing import Callable, Dict
-
-from jaxtyping import Array, Float
-
-from bayinx.core import Flow
-
-
-# TODO
-class Sylvester(Flow):
-    """
-    A sylvester flow.
-
-    # Attributes
-    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
-    - `constraints`: A dictionary of constraining transformations.
-    """
-
-    params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]