Refactor links

* rename link_functions to links
* separated activation function Ordered from OrderedQuantiles, one for
  generality, the other for automatic smart anchor selection based on
  quantile levels
* introduce link function for learnable positive semi-definite matrices
* link tests

Author: han-ol

bayesflow/link_functions/__init__.py

@@ -0,0 +1,3 @@
+from .ordered import Ordered
+from .ordered_quantiles import OrderedQuantiles
+from .positive_semi_definite import PositiveSemiDefinite

bayesflow/link_functions/ordered_quantiles.py

@@ -0,0 +1,40 @@
+import keras
+
+from bayesflow.utils import keras_kwargs
+
+
+class Ordered(keras.Layer):
+    def __init__(self, axis: int, anchor_index: int, **kwargs):
+        super().__init__(**keras_kwargs(kwargs))
+        self.axis = axis
+        self.anchor_index = anchor_index
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        print("build Ordered()")
+
+        self.group_indeces = dict(
+            below=list(range(0, self.anchor_index)),
+            above=list(range(self.anchor_index + 1, input_shape[self.axis])),
+        )
+
+    def call(self, inputs):
+        # Divide in anchor, below and above
+        below_inputs = keras.ops.take(inputs, self.group_indeces["below"], axis=self.axis)
+        anchor_input = keras.ops.take(inputs, self.anchor_index, axis=self.axis)
+        anchor_input = keras.ops.expand_dims(anchor_input, axis=self.axis)
+        above_inputs = keras.ops.take(inputs, self.group_indeces["above"], axis=self.axis)
+
+        # Apply softplus for positivity and cumulate to ensure ordered quantiles
+        below = keras.activations.softplus(below_inputs)
+        above = keras.activations.softplus(above_inputs)
+
+        below = anchor_input - keras.ops.flip(keras.ops.cumsum(below, axis=self.axis), self.axis)
+        above = anchor_input + keras.ops.cumsum(above, axis=self.axis)
+
+        # Concatenate and reshape back
+        x = keras.ops.concatenate([below, anchor_input, above], self.axis)
+        return x
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

bayesflow/links/__init__.py

@@ -0,0 +1,31 @@
+import keras
+
+from bayesflow.utils import keras_kwargs
+
+from collections.abc import Sequence
+
+from .ordered import Ordered
+
+
+class OrderedQuantiles(Ordered):
+    def __init__(self, q: Sequence[float] = None, axis: int = None, anchor_index: int = None, **kwargs):
+        super().__init__(axis, anchor_index, **keras_kwargs(kwargs))
+        self.q = q
+
+    def build(self, input_shape):
+        if 1 < len(input_shape) <= 3:
+            self.axis = -2
+            if self.q is None:
+                # choose the middle of the specified axis as anchor index
+                num_quantile_levels = input_shape[self.axis]
+                self.anchor_index = num_quantile_levels // 2
+            else:
+                # choose quantile level closest to median as anchor index
+                self.anchor_index = keras.ops.argmin(keras.ops.abs(keras.ops.convert_to_tensor(self.q) - 0.5))
+
+            super().build(input_shape)
+
+        else:
+            raise AssertionError(
+                f"Cannot resolve which axis should be ordered automatically from input shape {input_shape}."
+            )

bayesflow/links/ordered.py

@@ -0,0 +1,14 @@
+import keras
+
+from bayesflow.utils import keras_kwargs
+
+
+class PositiveSemiDefinite(keras.Layer):
+    def __init__(self, **kwargs):
+        super().__init__(**keras_kwargs(kwargs))
+
+    def call(self, inputs):
+        return keras.ops.einsum("...ij,...kj->...ik", inputs, inputs)
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

bayesflow/links/ordered_quantiles.py

@@ -0,0 +1,53 @@
+import numpy as np
+import keras
+import pytest
+
+
+@pytest.fixture()
+def batch_size():
+    return 16
+
+
+@pytest.fixture()
+def num_quantiles():
+    return 19
+
+
+@pytest.fixture()
+def quantiles_np(num_quantiles):
+    return np.linspace(0, 1, num_quantiles + 2)[1:-1]
+
+
+@pytest.fixture()
+def quantiles_py(quantiles_np):
+    return list(quantiles_np)
+
+
+@pytest.fixture()
+def quantiles_keras(quantiles_np):
+    return keras.ops.convert_to_tensor(quantiles_np)
+
+
+@pytest.fixture()
+def none():
+    return None
+
+
+@pytest.fixture(params=["quantiles_np", "quantiles_py", "quantiles_keras", "none"], scope="function")
+def quantiles(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture()
+def num_variables():
+    return 10
+
+
+@pytest.fixture()
+def unordered(batch_size, num_quantiles, num_variables):
+    return keras.random.normal((batch_size, num_quantiles, num_variables))
+
+
+@pytest.fixture()
+def random_matrix_batch(batch_size, num_variables):
+    return keras.random.normal((batch_size, num_variables, num_variables))

bayesflow/links/positive_semi_definite.py

@@ -0,0 +1,49 @@
+import numpy as np
+import pytest
+
+
+def check_ordering(output, axis):
+    assert np.all(np.diff(output, axis=axis) > 0), f"is not ordered along specified axis: {axis}."
+    for i in range(output.ndim):
+        if i != axis % output.ndim:
+            assert not np.all(
+                np.diff(output, axis=i) > 0
+            ), f"is ordered along axis which is not meant to be ordered: {i}."
+
+
+@pytest.mark.parametrize("axis", [0, 1, 2])
+def test_ordering(axis, unordered):
+    from bayesflow.links import Ordered
+
+    activation = Ordered(axis=axis, anchor_index=5)
+
+    output = activation(unordered)
+
+    check_ordering(output, axis)
+
+
+def test_quantile_ordering(quantiles, unordered):
+    from bayesflow.links import OrderedQuantiles
+
+    activation = OrderedQuantiles(q=None)
+
+    activation.build(unordered.shape)
+    axis = activation.axis
+
+    output = activation(unordered)
+
+    check_ordering(output, axis)
+
+
+def test_positive_semi_definite(random_matrix_batch):
+    from bayesflow.links import PositiveSemiDefinite
+
+    activation = PositiveSemiDefinite()
+
+    output = activation(random_matrix_batch)
+
+    eigenvalues = np.linalg.eig(output).eigenvalues
+
+    assert np.all(eigenvalues.real > 0) and np.all(
+        np.isclose(eigenvalues.imag, 0)
+    ), "output is not positive semi-definite."

tests/test_links/__init__.py

@@ -9,7 +9,7 @@ def batch_size():
 
 @pytest.fixture()
 def num_variables():
-    return 4
+    return 10
 
 
 @pytest.fixture()