Merge pull request #206 from aseyboldt/add-bijections-only

Closes #70

flowjax/bijections/__init__.py

@@ -24,7 +24,10 @@
     NumericalInverse,
     Permute,
     Reshape,
+    Sandwich,
 )
+from .utils import EmbedCondition, Flip, Identity, Invert, Permute, Reshape, Sandwich
+from .orthogonal import Householder, DiscreteCosine
 
 __all__ = [
     "AdditiveCondition",
@@ -34,9 +37,11 @@
     "Chain",
     "Concatenate",
     "Coupling",
+    "DiscreteCosine",
     "EmbedCondition",
     "Exp",
     "Flip",
+    "Householder",
     "Identity",
     "Invert",
     "LeakyTanh",
@@ -48,6 +53,7 @@
     "Planar",
     "RationalQuadraticSpline",
     "Reshape",
+    "Sandwich",
     "Scale",
     "Scan",
     "Sigmoid",

flowjax/bijections/affine.py

@@ -200,7 +200,7 @@ def __init__(
         self.cond_shape = cond_shape
 
     def transform_and_log_det(self, x, condition=None):
-        return x + self.module(condition), jnp.array(0)
+        return x + self.module(condition), jnp.zeros(())
 
     def inverse_and_log_det(self, y, condition=None):
-        return y - self.module(condition), jnp.array(0)
+        return y - self.module(condition), jnp.zeros(())

flowjax/bijections/bijection.py

@@ -15,6 +15,7 @@
 from equinox import AbstractVar
 from jaxtyping import Array, ArrayLike
 from paramax import unwrap
+import jax
 
 from flowjax.utils import _get_ufunc_signature, arraylike_to_array
 

flowjax/bijections/chain.py

@@ -3,6 +3,8 @@
 from collections.abc import Sequence
 
 from paramax import AbstractUnwrappable, unwrap
+import jax.numpy as jnp
+from jax import Array
 
 from flowjax.bijections.bijection import AbstractBijection
 from flowjax.utils import check_shapes_match, merge_cond_shapes
@@ -36,15 +38,19 @@ def __init__(
         self.cond_shape = merge_cond_shapes([unwrap(b).cond_shape for b in unwrapped])
         self.bijections = tuple(bijections)
 
-    def transform_and_log_det(self, x, condition=None):
-        log_abs_det_jac = 0
+    def transform_and_log_det(
+        self, x: Array, condition: Array | None = None
+    ) -> tuple[Array, Array]:
+        log_abs_det_jac = jnp.zeros(())
         for bijection in self.bijections:
             x, log_abs_det_jac_i = bijection.transform_and_log_det(x, condition)
             log_abs_det_jac += log_abs_det_jac_i.sum()
         return x, log_abs_det_jac
 
-    def inverse_and_log_det(self, y, condition=None):
-        log_abs_det_jac = 0
+    def inverse_and_log_det(
+        self, y: Array, condition: Array | None = None
+    ) -> tuple[Array, Array]:
+        log_abs_det_jac = jnp.zeros(())
         for bijection in reversed(self.bijections):
             y, log_abs_det_jac_i = bijection.inverse_and_log_det(y, condition)
             log_abs_det_jac += log_abs_det_jac_i.sum()

flowjax/bijections/coupling.py

@@ -8,6 +8,7 @@
 import equinox as eqx
 import jax.nn as jnn
 import jax.numpy as jnp
+import jax
 import paramax
 from jaxtyping import PRNGKeyArray
 

flowjax/bijections/jax_transforms.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable
 
 import equinox as eqx
+from jax import Array
 import jax.numpy as jnp
 from jax.lax import scan
 from jax.tree_util import tree_leaves, tree_map

flowjax/bijections/orthogonal.py

@@ -0,0 +1,71 @@
+from paramax import AbstractUnwrappable, Parameterize
+from flowjax.bijections.bijection import AbstractBijection
+from jax import Array
+import jax.numpy as jnp
+import jax.nn as jnn
+from jax.scipy import fft
+
+
+class Householder(AbstractBijection):
+    """A Householder reflection bijection.
+
+    This bijection implements a Householder reflection, which is a linear
+    transformation that reflects vectors across a hyperplane defined by a normal
+    vector (params). The transformation is its own inverse and volume-preserving
+    (determinant = ±1).
+
+    Given a unit vector v, the transformation is:
+    x → x - 2(x·v)v
+
+    Attributes:
+        shape: Shape of the input/output vectors
+        cond_shape: Shape of conditional inputs (None as this bijection is unconditional)
+        params: Normal vector defining the reflection hyperplane. The vector is
+            normalized in the transformation, so scaling params will have no effect
+            on the bijection.
+    """
+    shape: tuple[int, ...]
+    unit_vec: Array | AbstractUnwrappable
+    cond_shape = None
+
+    def __init__(self, params: Array):
+        self.shape = (params.shape[-1],)
+        self.unit_vec = Parameterize(lambda x: x / jnp.linalg.norm(x), params)
+
+    def _householder(self, x: Array) -> Array:
+        return x - 2 * self.unit_vec * (x @ self.unit_vec)
+
+    def transform_and_log_det(self, x: jnp.ndarray, condition: Array | None = None):
+        return self._householder(x), jnp.zeros(())
+
+    def inverse_and_log_det(self, y: Array, condition: Array | None = None):
+        return self._householder(y), jnp.zeros(())
+
+
+class DiscreteCosine(AbstractBijection):
+    """Discrete Cosine Transform (DCT) bijection.
+
+    This bijection applies the DCT or its inverse along a specified axis.
+
+    Attributes:
+        shape: Shape of the input/output arrays
+        cond_shape: Shape of conditional inputs (None as this bijection is unconditional)
+        axis: Axis along which to apply the DCT
+        norm: Normalization method, fixed to 'ortho' to ensure bijectivity
+    """
+
+    shape: tuple[int, ...]
+    cond_shape = None
+    axis: int
+
+    def __init__(self, shape, *, axis: int = -1):
+        self.shape = shape
+        self.axis = axis
+
+    def transform_and_log_det(self, x: jnp.ndarray, condition: Array | None = None):
+        y = fft.dct(x, norm="ortho", axis=self.axis)
+        return y, jnp.zeros(())
+
+    def inverse_and_log_det(self, y: Array, condition: Array | None = None):
+        x = fft.idct(y, norm="ortho", axis=self.axis)
+        return x, jnp.zeros(())

flowjax/bijections/softplus.py

@@ -2,11 +2,15 @@
 
 from typing import ClassVar
 
+import jax
 import jax.numpy as jnp
-from jax.nn import softplus
+from jax.nn import softplus, soft_sign
+from jaxtyping import Array, ArrayLike
+from paramax import AbstractUnwrappable, Parameterize, unwrap
+from paramax.utils import inv_softplus
 
 from flowjax.bijections.bijection import AbstractBijection
-
+from flowjax.utils import arraylike_to_array
 
 class SoftPlus(AbstractBijection):
     r"""Transforms to positive domain using softplus :math:`y = \log(1 + \exp(x))`."""

flowjax/bijections/utils.py

@@ -10,7 +10,8 @@
 from jaxtyping import Array, Int
 
 from flowjax.bijections.bijection import AbstractBijection
-from flowjax.utils import arraylike_to_array
+from flowjax.bijections.chain import Chain
+from flowjax.utils import arraylike_to_array, check_shapes_match, merge_cond_shapes
 
 
 class Invert(AbstractBijection):
@@ -79,10 +80,10 @@ def __init__(self, permutation: Int[Array | np.ndarray, "..."]):
         )
 
     def transform_and_log_det(self, x, condition=None):
-        return x[self.permutation], jnp.array(0)
+        return x[self.permutation], jnp.array(0.0)
 
     def inverse_and_log_det(self, y, condition=None):
-        return y[self.inverse_permutation], jnp.array(0)
+        return y[self.inverse_permutation], jnp.array(0.0)
 
 
 class Flip(AbstractBijection):
@@ -95,11 +96,15 @@ class Flip(AbstractBijection):
     shape: tuple[int, ...] = ()
     cond_shape: ClassVar[None] = None
 
-    def transform_and_log_det(self, x, condition=None):
-        return jnp.flip(x), jnp.array(0)
+    def transform_and_log_det(
+        self, x: Array, condition: Array | None = None
+    ) -> tuple[Array, Array]:
+        return jnp.flip(x), jnp.zeros(())
 
-    def inverse_and_log_det(self, y, condition=None):
-        return jnp.flip(y), jnp.array(0)
+    def inverse_and_log_det(
+        self, y: Array, condition: Array | None = None
+    ) -> tuple[Array, Array]:
+        return jnp.flip(y), jnp.zeros(())
 
 
 class Indexed(AbstractBijection):
@@ -300,3 +305,44 @@ def inverse_and_log_det(self, y, condition=None):
         x = self.inverter(self.bijection, y, condition)
         _, log_det = self.bijection.transform_and_log_det(x, condition)
         return x, -log_det
+
+
+class Sandwich(AbstractBijection):
+    """A bijection that composes bijections in a nested structure: g⁻¹ ∘ f ∘ g.
+
+    The Sandwich bijection creates a new transformation by "sandwiching" one
+    bijection between the forward and inverse applications of another. Given
+    bijections f and g, it computes:
+        Forward:  x → g⁻¹(f(g(x)))
+        Inverse:  y → g⁻¹(f⁻¹(g(y)))
+
+    This composition pattern is useful for:
+    - Creating symmetries in the transformation
+    - Applying a transformation in a different coordinate system
+    - Building more complex bijections from simpler ones
+
+    Attributes:
+        shape: Shape of the input/output arrays
+        cond_shape: Shape of conditional inputs
+        outer: Transformation g applied first and inverted last
+        inner: Transformation f applied in the middle
+    """
+    shape: tuple[int, ...]
+    cond_shape: tuple[int, ...] | None
+    outer: AbstractBijection
+    inner: AbstractBijection
+
+    def __init__(self, outer: AbstractBijection, inner: AbstractBijection):
+        check_shapes_match([outer.shape, inner.shape])
+        self.cond_shape = merge_cond_shapes([outer.cond_shape, inner.cond_shape])
+        self.shape = inner.shape
+        self.outer = outer
+        self.inner = inner
+
+    def transform_and_log_det(self, x: Array, condition=None) -> tuple[Array, Array]:
+        chain = Chain([self.outer, self.inner, Invert(self.outer)])
+        return chain.transform_and_log_det(x, condition)
+
+    def inverse_and_log_det(self, y: Array, condition=None) -> tuple[Array, Array]:
+        chain = Chain([self.outer, self.inner, Invert(self.outer)])
+        return chain.inverse_and_log_det(y, condition)

tests/test_bijections/test_bijections.py

@@ -7,6 +7,8 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+import numpy as np
+from scipy import stats
 
 from flowjax.bijections import (
     AbstractBijection,
@@ -16,9 +18,11 @@
     Chain,
     Concatenate,
     Coupling,
+    DiscreteCosine,
     EmbedCondition,
     Exp,
     Flip,
+    Householder,
     Identity,
     Indexed,
     LeakyTanh,
@@ -30,6 +34,7 @@
     Power,
     RationalQuadraticSpline,
     Reshape,
+    Sandwich,
     Scale,
     Scan,
     Sigmoid,
@@ -217,6 +222,12 @@
             partial(bisection_search, lower=-1, upper=1, atol=1e-7),
         ),
     ),
+    "Sandwich": lambda: Sandwich(
+        Exp(),
+        Affine(0.1, 0.5),
+    ),
+    "DiscreteCosine": lambda: DiscreteCosine(shape=(3, 4)),
+    "Householder": lambda: Householder(jnp.ones(3)),
 }