Merge branch 'dev' into point-estimation

Author: han-ol

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/adapters/adapter.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, MutableSequence, Sequence
 
 import numpy as np
 from keras.saving import (
@@ -26,17 +26,16 @@
     ToArray,
     Transform,
 )
-
 from .transforms.filter_transform import Predicate
 
 
 @serializable(package="bayesflow.adapters")
-class Adapter:
+class Adapter(MutableSequence[Transform]):
     def __init__(self, transforms: Sequence[Transform] | None = None):
         if transforms is None:
             transforms = []
 
-        self.transforms = transforms
+        self.transforms = list(transforms)
 
     @staticmethod
     def create_default(inference_variables: Sequence[str]) -> "Adapter":
@@ -77,12 +76,70 @@ def __call__(self, data: dict[str, any], *, inverse: bool = False, **kwargs) ->
         return self.forward(data, **kwargs)
 
     def __repr__(self):
-        return f"Adapter([{' -> '.join(map(repr, self.transforms))}])"
+        result = ""
+        for i, transform in enumerate(self):
+            result += f"{i}: {transform!r}"
+            if i != len(self) - 1:
+                result += " -> "
 
-    def add_transform(self, transform: Transform):
-        self.transforms.append(transform)
+        return f"Adapter([{result}])"
+
+    # list methods
+
+    def append(self, value: Transform) -> "Adapter":
+        self.transforms.append(value)
         return self
 
+    def __delitem__(self, key: int | slice):
+        del self.transforms[key]
+
+    def extend(self, values: Sequence[Transform]) -> "Adapter":
+        if isinstance(values, Adapter):
+            values = values.transforms
+
+        self.transforms.extend(values)
+
+        return self
+
+    def __getitem__(self, item: int | slice) -> "Adapter":
+        if isinstance(item, int):
+            return self.transforms[item]
+
+        return Adapter(self.transforms[item])
+
+    def insert(self, index: int, value: Transform | Sequence[Transform]) -> "Adapter":
+        if isinstance(value, Adapter):
+            value = value.transforms
+
+        if isinstance(value, Sequence):
+            # convenience: Adapters are always flat
+            self.transforms = self.transforms[:index] + list(value) + self.transforms[index:]
+        else:
+            self.transforms.insert(index, value)
+
+        return self
+
+    def __setitem__(self, key: int | slice, value: Transform | Sequence[Transform]) -> "Adapter":
+        if isinstance(value, Adapter):
+            value = value.transforms
+
+        if isinstance(key, int) and isinstance(value, Sequence):
+            if key < 0:
+                key += len(self.transforms)
+
+            key = slice(key, key + 1)
+
+        self.transforms[key] = value
+
+        return self
+
+    def __len__(self):
+        return len(self.transforms)
+
+    # adapter methods
+
+    add_transform = append
+
     def apply(
         self,
         *,

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/mc_calibration.py

@@ -68,41 +68,42 @@ def mc_calibration(
 
     # Gather plot data and metadata into a dictionary
     plot_data = prepare_plot_data(
-        estimates=pred_models,
-        ground_truths=true_models,
+        targets=pred_models,
+        references=true_models,
         variable_names=model_names,
         num_col=num_col,
         num_row=num_row,
         figsize=figsize,
+        default_name="M",
     )
 
     # Compute calibration
     cal_errors, true_probs, pred_probs = expected_calibration_error(
-        plot_data["ground_truths"], plot_data["estimates"], num_bins
+        plot_data["references"], plot_data["targets"], num_bins
     )
 
     for j, ax in enumerate(plot_data["axes"].flat):
         # Plot calibration curve
-        ax[j].plot(pred_probs[j], true_probs[j], "o-", color=color)
+        ax.plot(pred_probs[j], true_probs[j], "o-", color=color)
 
         # Plot PMP distribution over bins
         uniform_bins = np.linspace(0.0, 1.0, num_bins + 1)
-        norm_weights = np.ones_like(plot_data["estimates"]) / len(plot_data["estimates"])
-        ax[j].hist(plot_data["estimates"][:, j], bins=uniform_bins, weights=norm_weights[:, j], color="grey", alpha=0.3)
+        norm_weights = np.ones_like(plot_data["targets"]) / len(plot_data["targets"])
+        ax.hist(plot_data["targets"][:, j], bins=uniform_bins, weights=norm_weights[:, j], color="grey", alpha=0.3)
 
         # Plot AB line
-        ax[j].plot((0, 1), (0, 1), "--", color="black", alpha=0.9)
+        ax.plot((0, 1), (0, 1), "--", color="black", alpha=0.9)
 
         # Tweak plot
-        ax[j].set_xlim([0 - epsilon, 1 + epsilon])
-        ax[j].set_ylim([0 - epsilon, 1 + epsilon])
-        ax[j].set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
-        ax[j].set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
+        ax.set_xlim([0 - epsilon, 1 + epsilon])
+        ax.set_ylim([0 - epsilon, 1 + epsilon])
+        ax.set_xticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
+        ax.set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
 
         # Add ECE label
         add_metric(
-            ax[j],
-            metric_text=r"$\widehat{{\mathrm{{ECE}}}}$ = {0:.3f}",
+            ax,
+            metric_text=r"$\widehat{{\mathrm{{ECE}}}}$",
             metric_value=cal_errors[j],
             metric_fontsize=metric_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