Add Rational Quadratic Spline Transforms to Normalizing Flows (#291)

* Splines draft

* update keras requirement

* small improvements to error messages

* add rq spline function

* add spline transform

* update searchsorted utils for jax
also add padd util

* update tests

* add assert_allclose util for improved messages

* parametrize transform for flow tests

* update jacobian, jacobian trace, vjp, jvp, and corresponding usages and tests

* fix imports, remove old jacobian and jvp, fix application in free form flow

* improve logdet computation in free form flows

* Fix comparison for symbolic tensors under tf

* Add splines to twomoons notebook

* improve pad utility

* fix missing left edge in spline

* fix inside mask edge case

* explicitly set bias initializer

* add better expand utility

* small clean up, renaming

* fix indexing, fix inside check

* dump

* fix sign of log jacobian for inverse pass in rq spline

* fix parameter splitting for spline transform

* improve readability

* fix scale and shift trailing dimension

* fix inverse pass return value

* correctly choose bins once for each dimension, even for multi-dimensional inputs

* run formatter

* reduce searchsorted log spam

* log backend used at setup

* remove maximum message cache size

* Improve warning message for jax searchsorted

* Fix spline parameter binning for compiled contexts

* update inverse transform same as forward

* Update TwoMoons notebook with splines WIP [skip ci]

* fix spline inverse call for out of bounds values

* Add working splines

---------

Co-authored-by: stefanradev93 <[email protected]>

Author: 2 people

bayesflow/__init__.py

@@ -36,6 +36,10 @@ def setup():
 
         torch.autograd.set_grad_enabled(False)
 
+    from bayesflow.utils import logging
+
+    logging.info(f"Using backend {keras.backend.backend()!r}")
+
 
 # call and clean up namespace
 setup()

bayesflow/diagnostics/plots/calibration_ecdf.py

@@ -176,7 +176,7 @@ def calibration_ecdf(
         titles = ["Stacked ECDFs"]
 
     for ax, title in zip(plot_data["axes"].flat, titles):
-        ax.fill_between(z, L, H, color=fill_color, alpha=0.2, label=rf"{int((1-alpha) * 100)}$\%$ Confidence Bands")
+        ax.fill_between(z, L, H, color=fill_color, alpha=0.2, label=rf"{int((1 - alpha) * 100)}$\%$ Confidence Bands")
         ax.legend(fontsize=legend_fontsize)
         ax.set_title(title, fontsize=title_fontsize)
 

bayesflow/diagnostics/plots/mmd_hypothesis_test.py

@@ -79,7 +79,7 @@ def fill_area_under_kde(kde_object, x_start, x_end=None, **kwargs):
 
     mmd_critical = ops.quantile(mmd_null, 1 - alpha_level)
     fill_area_under_kde(
-        kde, mmd_critical, color=alpha_color, alpha=0.5, label=rf"{int(alpha_level*100)}% rejection area"
+        kde, mmd_critical, color=alpha_color, alpha=0.5, label=rf"{int(alpha_level * 100)}% rejection area"
     )
 
     if truncate_v_lines_at_kde:

bayesflow/networks/consistency_models/continuous_consistency_model.py

@@ -249,7 +249,7 @@ def f_teacher(x, t):
             ops.cos(t) * ops.sin(t) * self.sigma_data,
         )
 
-        teacher_output, cos_sin_dFdt = jvp(f_teacher, primals, tangents)
+        teacher_output, cos_sin_dFdt = jvp(f_teacher, primals, tangents, return_output=True)
         teacher_output = ops.stop_gradient(teacher_output)
         cos_sin_dFdt = ops.stop_gradient(cos_sin_dFdt)
 

bayesflow/networks/coupling_flow/couplings/single_coupling.py

@@ -1,5 +1,4 @@
 import keras
-
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
@@ -24,6 +23,7 @@ def __init__(self, subnet: str | type = "mlp", transform: str = "affine", **kwar
 
         output_projector_kwargs = kwargs.get("output_projector_kwargs", {})
         output_projector_kwargs.setdefault("kernel_initializer", "zeros")
+        output_projector_kwargs.setdefault("bias_initializer", "zeros")
         self.output_projector = keras.layers.Dense(units=None, **output_projector_kwargs)
 
         # serialization: store all parameters necessary to call __init__

bayesflow/networks/coupling_flow/transforms/_rational_quadratic.py

@@ -0,0 +1,81 @@
+from typing import TypedDict
+
+import keras
+
+from bayesflow.types import Tensor
+
+
+class Edges(TypedDict):
+    left: Tensor
+    right: Tensor
+    bottom: Tensor
+    top: Tensor
+
+
+class Derivatives(TypedDict):
+    left: Tensor
+    right: Tensor
+
+
+def _rational_quadratic_spline(
+    x: Tensor, edges: Edges, derivatives: Derivatives, inverse: bool = False
+) -> (Tensor, Tensor):
+    # rename variables to match the paper:
+
+    # $x^{(k)}$
+    xk = edges["left"]
+
+    # $x^{(k+1)}$
+    xkp = edges["right"]
+
+    # $y^{(k)}$
+    yk = edges["bottom"]
+
+    # $y^{(k+1)}$
+    ykp = edges["top"]
+
+    # $delta^{(k)}$
+    dk = derivatives["left"]
+
+    # $delta^{(k+1)}$
+    dkp = derivatives["right"]
+
+    # commonly used values
+    dx = xkp - xk
+    dy = ykp - yk
+    sk = dy / dx
+
+    if not inverse:
+        xi = (x - xk) / dx
+
+        # Eq. 4 in the paper
+        numerator = dy * (sk * xi**2 + dk * xi * (1 - xi))
+        denominator = sk + (dkp + dk - 2 * sk) * xi * (1 - xi)
+        result = yk + numerator / denominator
+    else:
+        # rename for clarity
+        y = x
+
+        # Eq. 6-8 in the paper
+        a = dy * (sk - dk) + (y - yk) * (dkp + dk - 2 * sk)
+        b = dy * dk - (y - yk) * (dkp + dk - 2 * sk)
+        c = -sk * (y - yk)
+
+        # Eq. 29 in the appendix of the paper
+        discriminant = b**2 - 4 * a * c
+
+        # the discriminant must be positive, even when the spline is called out of bounds
+        discriminant = keras.ops.maximum(discriminant, 0)
+
+        xi = 2 * c / (-b - keras.ops.sqrt(discriminant))
+        result = xi * dx + xk
+
+    # Eq 5 in the paper
+    numerator = sk**2 * (dkp * xi**2 + 2 * sk * xi * (1 - xi) + dk * (1 - xi) ** 2)
+    denominator = (sk + (dkp + dk - 2 * sk) * xi * (1 - xi)) ** 2
+    log_jac = keras.ops.log(numerator) - keras.ops.log(denominator)
+
+    if inverse:
+        log_jac = -log_jac
+
+    return result, log_jac