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

bayinx/flows/fullaffine.py

@@ -0,0 +1,123 @@
+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 FullAffineLayer(FlowLayer):
+    """
+    A full affine flow.
+
+    # Attributes
+    - `params`: The parameters of the full affine flow.
+    - `constraints`: The constraining transformations for the parameters of the full 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(key, (dim, )) / dim**0.5,
+            "scale": jr.normal(key, (dim, dim)) / dim**0.5,
+        }
+
+        # Define constraints
+        if dim == 1:
+            self.constraints = {"scale": jnp.exp}
+        else:
+            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: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        shift: Float[Array, " n_dim"] = params["shift"]
+        scale: Float[Array, "n_dim n_dim"] = params["scale"]
+
+        # Compute forward transformation
+        draws = (scale @ draws.T).T + shift
+
+        return draws
+
+    def __adjust(self, draw: Float[Array, " n_dim"]) -> Float[Array, " n_dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        scale: Float[Array, "n_dim n_dim"] = params["scale"]
+
+        # Compute log-Jacobian adjustments
+        lja: Array = jnp.log(jnp.diag(scale)).sum()
+
+        assert lja.shape == ()
+
+        return lja
+
+    @eqx.filter_jit
+    def adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__adjust, 0)
+        return f(draws)
+
+    def __forward_and_adjust(self, draw: Float[Array, " n_dim"]) -> Tuple[Float[Array, " n_dim"], Scalar]:
+        params = self.transform_params()
+
+        assert len(draw.shape) == 1
+
+        # Extract parameters
+        shift: Float[Array, " n_dim"] = params["shift"]
+        scale: Float[Array, "n_dim n_dim"] = params["scale"]
+
+        # Compute forward transformation
+        draw = (scale @ draw.T).T + shift
+
+        assert len(draw.shape) == 1
+
+        # Compute lja
+        lja: Scalar = jnp.log(jnp.diag(scale)).sum()
+
+        assert lja.shape == ()
+
+        return draw, lja
+
+    @eqx.filter_jit
+    def forward_and_adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Tuple[Float[Array, "n_draws n_dim"], Scalar]:
+        f = jax.vmap(self.__forward_and_adjust, 0)
+        return f(draws)
+
+
+class FullAffine(FlowSpec):
+    key: PRNGKeyArray
+
+    def __init__(self, key: PRNGKeyArray = jr.key(0)):
+        self.key = key
+
+    def construct(self, dim: int) -> FullAffineLayer:
+        return FullAffineLayer(dim, self.key)

bayinx/flows/lowrankaffine.py

@@ -0,0 +1,165 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jax.lax import scan
+from jaxtyping import Array, Float, PRNGKeyArray, Scalar
+
+from bayinx.core.flow import FlowLayer, FlowSpec
+
+
+def _mat_op(
+    carry: Float[Array, " rank"],
+    vars: Tuple[
+        Float[Array, " rank"],
+        Float[Array, " rank"],
+        Scalar,
+        Scalar,
+        Scalar
+    ]
+) -> Tuple[Float[Array, " rank"], Scalar]:
+    """
+    The implicit matrix operations compute:
+
+        x = Lz + b = (D + W)z + b = Dz + (UV^T * M)z + b
+
+    Where W = UV^T * M is the low-rank representation of the strictly lower triangular matrix W,
+        parameterized by U and V, and an implicit mask M that zeros out upper triangular elements.
+
+    In the implementation only U, V, diag(D), and b are needed to complete the matrix operations.
+    """
+    # Unpack state
+    U_r, V_r, z_i, d_i, b_i = vars
+    h_i = carry
+
+    # Compute partial product
+    h_next = h_i + z_i * V_r
+
+    # Compute i-th element of transformed draw:
+    x_i = b_i + z_i * d_i + U_r.dot(h_i)
+
+    return h_next, x_i
+
+class LowRankAffineLayer(FlowLayer):
+    """
+    A low-rank affine flow.
+
+    # Attributes
+    - `params`: A dictionary containing the shift and low-rank representation of the scale parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    - `static`: Whether the flow layer is frozen (parameters are not subject to further optimization).
+    - `rank`: Rank of the scale transformation.
+    """
+
+    rank: int
+
+    def __init__(self, dim: int, rank: int, key: PRNGKeyArray = jr.key(0)):
+        """
+        Initializes a low-rank affine flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        - `rank`: The rank of the (implicit) scale matrix.
+        """
+        self.static = False
+        self.rank = rank
+
+        # Split key
+        k1, k2, k3, k4 = jr.split(key, 4)
+
+        # Initialize parameters
+        self.params = {
+            "shift": jr.normal(k1, (dim, )) / dim**0.5,
+            "diag_scale": jr.normal(k2, (dim, )) / dim**0.5,
+            "offdiag_scale": (
+                jr.normal(k3, (dim, rank)) / dim**0.5,
+                jr.normal(k4, (dim, rank)) / dim**0.5
+            )
+        }
+
+        # Define constraints
+        self.constraints = {"diag_scale": jnp.exp}
+
+
+    def __forward(self, draw: Float[Array, " n_dim"]) -> Float[Array, " n_dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        shift: Array = params["shift"]
+        diag: Array = params["diag_scale"]
+        U, V = params["offdiag_scale"]
+
+        # Compute forward transformation
+        _, draw = scan(
+            f=_mat_op,
+            init=jnp.zeros((self.rank, )),
+            xs= (U, V, draw, diag, shift)
+        )
+
+        return draw
+
+    @eqx.filter_jit
+    def forward(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__forward, 0)
+        return f(draws)
+
+    def __adjust(self, draw: Float[Array, " n_dim"]) -> Scalar:
+        params = self.transform_params()
+
+        diag: Array = params["diag_scale"]
+
+        # Compute log-Jacobian adjustment
+        lja: Scalar = jnp.log(diag).sum()
+
+        assert lja.shape == ()
+
+        return lja
+
+    @eqx.filter_jit
+    def adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__adjust, 0)
+        return f(draws)
+
+    def __forward_and_adjust(self, draw: Float[Array, " n_dim"]) -> Tuple[Float[Array, " n_dim"], Scalar]:
+        params = self.transform_params()
+
+        assert len(draw.shape) == 1
+
+        # Extract parameters
+        shift: Array = params["shift"]
+        diag: Array = params["diag_scale"]
+        U, V = params["offdiag_scale"]
+
+        # Compute log-Jacobian adjustment
+        lja: Scalar = jnp.log(diag).sum()
+
+        # Compute forward transformation
+        _, draw = scan(
+            f=_mat_op,
+            init=jnp.zeros((self.rank, )),
+            xs= (U, V, draw, diag, shift)
+        )
+
+        assert len(draw.shape) == 1
+        assert lja.shape == ()
+
+        return draw, lja
+
+    @eqx.filter_jit
+    def forward_and_adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Tuple[Float[Array, "n_draws n_dim"], Scalar]:
+        f = jax.vmap(self.__forward_and_adjust, 0)
+        return f(draws)
+
+class LowRankAffine(FlowSpec):
+    rank: int
+
+    def __init__(self, rank: int):
+        self.rank = rank
+
+    def construct(self, dim: int) -> LowRankAffineLayer:
+        if (self.rank > (dim - 1)/2):
+            raise ValueError(f"Rank {self.rank} is large, consider using a full affine flow instead.")
+
+        return LowRankAffineLayer(dim, self.rank)

bayinx/flows/planar.py

@@ -0,0 +1,155 @@
+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
+
+
+def _h(x: Array) -> Array:
+    return jnp.tanh(x)
+
+def _dh(x: Array) -> Array:
+    return 1.0 - jnp.tanh(x)**2
+
+
+class PlanarLayer(FlowLayer):
+    """
+    A Planar flow.
+
+    # Attributes
+    - `params`: The parameters of the Planar flow (w, u_hat, b).
+    - `constraints`: The constraining transformations for the parameters.
+    - `static`: Whether the flow layer is frozen.
+    """
+
+    params: Dict[str, Array]
+    constraints: Dict[str, Callable[[Array], Array]]
+    static: bool
+
+    def __init__(self, dim: int, key: PRNGKeyArray = jr.key(0)):
+        """
+        Initializes a Planar flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.static = False
+        # Split key
+        k1, k2, k3 = jr.split(key, 3)
+
+        # Initialize parameters
+        self.params = {
+            "w": jr.normal(k1, (dim,)) / dim**0.5,
+            "u_hat": jr.normal(k2, (dim,)) / dim**0.5,
+            "b": jr.normal(k3, ()) / dim**0.5,
+        }
+
+        self.constraints = {}
+
+    def __forward(self, draw: Float[Array, " n_dim"]) -> Float[Array, " n_dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        w: Float[Array, " n_dim"] = params["w"]
+        u: Float[Array, " n_dim"] = params["u"]
+        b: Scalar = params["b"]
+
+        # Compute inner term
+        a = draw.dot(w) + b
+
+        # Compute nonlinear stretch
+        h = _h(a)
+
+        # Compute forward transformation
+        draw = draw + u * h
+
+        return draw
+
+    @eqx.filter_jit
+    def forward(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__forward, 0)
+        return f(draws)
+
+    def __adjust(self, draw: Float[Array, "n_dim"]) -> Scalar: # noqa
+        params = self.transform_params()
+
+        # Extract parameters
+        w: Float[Array, " n_dim"] = params["w"]
+        u: Float[Array, " n_dim"] = params["u"]
+        b: Scalar = params["b"]
+
+        # Compute inner term
+        a = draw.dot(w) + b
+
+        # Compute log-Jacobian adjustment
+        lja: Scalar = jnp.log(1.0 + u.dot(_dh(a) * w))
+
+        assert lja.shape == ()
+        return lja
+
+    @eqx.filter_jit
+    def adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__adjust, 0)
+        return f(draws)
+
+    def __forward_and_adjust(self, draw: Float[Array, " n_dim"]) -> Tuple[Float[Array, " n_dim"], Scalar]:
+        params = self.transform_params()
+
+        # Extract parameters
+        w: Float[Array, " n_dim"] = params["w"]
+        u: Float[Array, " n_dim"] = params["u"]
+        b: Scalar = params["b"]
+
+        # Compute inner term
+        a = draw.dot(w) + b
+
+        # Compute forward transformation
+        draw = draw + u * _h(a)
+
+        # Compute log-Jacobian adjustment
+        lja: Scalar = jnp.log(1.0 + u.dot(_dh(a) * w))
+
+        assert len(draw.shape) == 1
+        assert lja.shape == ()
+
+        return draw, lja
+
+    @eqx.filter_jit
+    def forward_and_adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Tuple[Float[Array, "n_draws n_dim"], Scalar]:
+        f = jax.vmap(self.__forward_and_adjust, 0)
+        return f(draws)
+
+
+    def transform_params(self) -> Dict[str, Array]:
+        """
+        Applies the affine constraint to u_hat to compute the constrained parameter u,
+        and returns all parameters.
+
+        # Returns
+        A dictionary of parameters including the computed, constrained 'u'.
+        """
+        constrained_params = self.constrain_params()
+
+        # Extract parameters
+        w: Array = constrained_params["w"]
+        u_hat: Array = constrained_params["u_hat"]
+
+        # Compute the constrained u
+        u = u_hat + (-1.0 - jnp.dot(w, u_hat)) / (jnp.sum(w**2) + 1e-6) * w
+        constrained_params["u"] = u
+
+        return constrained_params
+
+class Planar(FlowSpec):
+    """
+    A specification for the Planar flow.
+    """
+    def __init__(self):
+        pass
+
+    def construct(self, dim: int) -> PlanarLayer:
+        return PlanarLayer(dim)

bayinx/flows/radial.py

@@ -0,0 +1 @@
+# WIP

bayinx/flows/sylvester.py

@@ -0,0 +1,225 @@
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jax.lax import scan
+from jaxtyping import Array, Float, PRNGKeyArray, Scalar
+
+from bayinx.core.flow import FlowLayer, FlowSpec
+
+
+def _h(x: Array) -> Array:
+    """Non-linearity for Sylvester flow."""
+    return jnp.tanh(x)
+
+
+def _dh(x: Array) -> Array:
+    """Derivative of the non-linearity."""
+    return 1.0 - jnp.tanh(x)**2
+
+
+def _construct_orthogonal(vectors: Float[Array, " dim rank"]) -> Float[Array, " dim rank"]:
+    """
+    Constructs a D x M orthogonal matrix (Q) using Householder reflections.
+    """
+    D, M = vectors.shape
+
+    # Initialize Q as the thin identity matrix (first M columns of I_D)
+    Q_initial = jnp.eye(D, M)
+
+    def apply_reflection(Q_current, v_k):
+        v_norm_sq = jnp.sum(v_k**2)
+        tau = 2.0 / v_norm_sq
+
+        a_T = jnp.dot(v_k, Q_current)
+
+        update = jnp.outer(v_k, tau * a_T)
+
+        Q_next = Q_current - update
+
+        return Q_next, None
+
+    Q_final, _ = scan(
+        f=apply_reflection,
+        init=Q_initial,
+        xs=vectors.T
+    )
+
+    return Q_final
+
+
+class SylvesterLayer(FlowLayer):
+    """
+    A Sylvester flow.
+
+    # Attributes
+    - `params`: Dictionary containing raw parameters.
+    - `constraints`: Dictionary of constraining transformations.
+    - `static`: Whether the flow layer is frozen.
+    - `rank`: The rank (M) of the transformation.
+    """
+
+    rank: int
+    params: Dict[str, Array]
+    constraints: Dict[str, Callable[[Array], Array]]
+    static: bool
+
+    def __init__(self, dim: int, rank: int, key: PRNGKeyArray = jr.key(0)):
+        """
+        Initializes the Sylvester flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space (D).
+        - `rank`: The number of hidden units/reflections (M).
+        """
+        self.static = False
+        self.rank = rank
+
+        k1, k2, k3, k4 = jr.split(key, 4)
+
+        # Initialize parameters
+        # hvecs: D x M matrix where each column is a Householder vector
+        self.params = {
+            "hvecs": jr.normal(k1, (dim, rank)) / dim**0.5,
+            "r1": jr.normal(k2, (rank, rank)) / rank**0.5,
+            "r2": jr.normal(k3, (rank, rank)) / rank**0.5,
+            "b": jr.normal(k4, (rank,)) / rank**0.5,
+        }
+
+        # Constraint for Upper Triangular matrices with positive diagonal
+        def constrain_triangular(matrix: Array):
+            # Extract diagonal and apply exponential to ensure invertibility
+            diag: Array = jnp.exp(jnp.diag(matrix))
+            # Return upper triangular matrix with modified diagonal
+            return jnp.fill_diagonal(jnp.triu(matrix), diag, inplace=False)
+
+        self.constraints = {
+            "r1": constrain_triangular,
+            "r2": constrain_triangular
+        }
+
+    def transform_params(self) -> Dict[str, Array]:
+        """
+        Applies constraints and constructs the orthogonal matrix Q.
+        """
+        constrained = self.constrain_params()
+
+        # Construct orthogonal matrix
+        q = _construct_orthogonal(constrained["hvecs"])
+        constrained["q"] = q
+
+        return constrained
+
+    def __forward(self, draw: Float[Array, " n_dim"]) -> Float[Array, " n_dim"]:
+        params = self.transform_params()
+
+        # Extract parameters
+        Q: Float[Array, "dim rank"] = params["q"]
+        R1: Float[Array, "rank rank"] = params["r1"]
+        R2: Float[Array, "rank rank"] = params["r2"]
+        b: Float[Array, " rank"] = params["b"]
+
+        # Compute inner terms
+        y = R2.dot(Q.T.dot(draw)) + b
+        h_y = _h(y)
+
+        # Compute forward transform
+        draw = draw + Q.dot(R1.dot(h_y))
+
+        return draw
+
+    @eqx.filter_jit
+    def forward(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__forward, 0)
+        return f(draws)
+
+    def __adjust(self, draw: Float[Array, " n_dim"]) -> Scalar:
+        params = self.transform_params()
+
+        # Extract parameters
+        Q: Float[Array, "dim rank"] = params["q"]
+        R1: Float[Array, "rank rank"] = params["r1"]
+        R2: Float[Array, "rank rank"] = params["r2"]
+        b: Float[Array, " rank"] = params["b"]
+
+        # Recompute the argument to the nonlinearity
+        term = R2.dot(Q.T.dot(draw)) + b
+        diag_h_prime = _dh(term)
+
+        # Diagonal of R1 and R2
+        d_r1 = jnp.diag(R1)
+        d_r2 = jnp.diag(R2)
+
+        # Diagonal term of the matrix (I + diag(h') R2 R1)
+        # diag_term_i = 1 + h'_i * (R2_ii * R1_ii)
+        diag_term = 1.0 + diag_h_prime * d_r2 * d_r1
+
+        lja = jnp.sum(jnp.log(diag_term))
+
+        assert lja.shape == ()
+
+        return lja
+
+    @eqx.filter_jit
+    def adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Float[Array, "n_draws n_dim"]:
+        f = jax.vmap(self.__adjust, 0)
+        return f(draws)
+
+    def __forward_and_adjust(self, draw: Float[Array, " n_dim"]) -> Tuple[Float[Array, " n_dim"], Scalar]:
+        params = self.transform_params()
+
+        Q: Float[Array, "dim rank"] = params["q"]
+        R1: Float[Array, "rank rank"] = params["r1"]
+        R2: Float[Array, "rank rank"] = params["r2"]
+        b: Float[Array, " rank"] = params["b"]
+
+        # --- Forward ---
+        q_z = jnp.dot(Q.T, draw)
+        arg = jnp.dot(R2, q_z) + b
+
+        # h(arg)
+        h_y = _h(arg)
+
+        # Update
+        term = jnp.dot(R1, h_y)
+        draw_new = draw + jnp.dot(Q, term)
+
+        # --- Adjust ---
+        # Derivative h'(arg)
+        diag_h_prime = _dh(arg)
+
+        # Diagonals of triangular matrices
+        d_r1 = jnp.diag(R1)
+        d_r2 = jnp.diag(R2)
+
+        # Log-det of triangular matrix
+        diag_term = 1.0 + diag_h_prime * d_r2 * d_r1
+        lja = jnp.sum(jnp.log(jnp.abs(diag_term)))
+
+        assert len(draw_new.shape) == 1
+        assert lja.shape == ()
+
+        return draw_new, lja
+
+    @eqx.filter_jit
+    def forward_and_adjust(self, draws: Float[Array, "n_draws n_dim"]) -> Tuple[Float[Array, "n_draws n_dim"], Scalar]:
+        f = jax.vmap(self.__forward_and_adjust, 0)
+        return f(draws)
+
+
+class Sylvester(FlowSpec):
+    """
+    A specification for the Orthogonal Sylvester flow.
+    """
+    rank: int
+
+    def __init__(self, rank: int):
+        self.rank = rank
+
+    def construct(self, dim: int) -> SylvesterLayer:
+        if self.rank > dim:
+            raise ValueError(f"Rank {self.rank} cannot be greater than dimension {dim}.")
+
+        return SylvesterLayer(dim, self.rank)

bayinx/nodes/__init__.py

@@ -0,0 +1,3 @@
+from .continuous import Continuous as Continuous
+from .observed import Observed as Observed
+from .stochastic import Stochastic as Stochastic