Merge branch 'dev' of https://github.com/stefanradev93/BayesFlow into dev

Author: stefanradev93

bayesflow/utils/optimal_transport/sinkhorn.py

@@ -16,7 +16,7 @@ def sinkhorn(
     cost: str | Tensor = "euclidean",
     seed: int = None,
     regularization: float = 1.0,
-    max_steps: int = 10000,
+    max_steps: int | None = 10_000,
     tolerance: float = 1e-6,
     numpy: bool = False,
 ) -> (Tensor, Tensor):
@@ -88,7 +88,7 @@ def sinkhorn_indices(
     cost: str | Tensor = "euclidean",
     seed: int = None,
     regularization: float = 1.0,
-    max_steps: int = 1000,
+    max_steps: int | None = 10_000,
     tolerance: float = 1e-6,
     numpy: bool = False,
 ) -> Tensor | np.ndarray:
@@ -111,7 +111,7 @@ def sinkhorn_indices(
         Default: 1.0
 
     :param max_steps: Maximum number of iterations.
-        Default: 1000
+        Default: 10_000
 
     :param tolerance: Absolute tolerance for convergence.
         Default: 1e-6
@@ -164,15 +164,13 @@ def sinkhorn_plan(
 
     :param regularization: Regularization parameter.
         Controls the standard deviation of the Gaussian kernel.
-        Default: 1.0
 
     :param max_steps: Maximum number of iterations.
-        Default: 1000
 
     :param tolerance: Absolute tolerance for convergence.
-        Default: 1e-6
 
     :param numpy: Whether to use numpy or keras backend.
+        Default: False
 
     :return: Tensor of shape (n, m)
         The transport probabilities.

tests/test_utils/test_dispatch.py

@@ -0,0 +1,202 @@
+import keras
+import pytest
+
+# Import the dispatch functions
+from bayesflow.utils import find_network, find_permutation, find_pooling, find_recurrent_net
+
+# --- Tests for find_network.py ---
+
+
+class DummyMLP:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+def test_find_network_with_string(monkeypatch):
+    # Monkeypatch the MLP entry in bayesflow.networks
+    monkeypatch.setattr("bayesflow.networks.MLP", DummyMLP)
+
+    net = find_network("mlp", 1, key="value")
+    assert isinstance(net, DummyMLP)
+    assert net.args == (1,)
+    assert net.kwargs == {"key": "value"}
+
+
+def test_find_network_with_type():
+    class CustomNet:
+        def __init__(self, x):
+            self.x = x
+
+    net = find_network(CustomNet, 42)
+    assert isinstance(net, CustomNet)
+    assert net.x == 42
+
+
+def test_find_network_with_keras_layer():
+    layer = keras.layers.Dense(10)
+    returned = find_network(layer)
+    assert returned is layer
+
+
+def test_find_network_invalid_type():
+    with pytest.raises(TypeError):
+        find_network(123)
+
+
+# --- Tests for find_permutation.py ---
+
+
+class DummyRandomPermutation:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+class DummySwap:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+class DummyOrthogonalPermutation:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+def test_find_permutation_random(monkeypatch):
+    type("dummy_mod", (), {"RandomPermutation": DummyRandomPermutation})
+    monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.RandomPermutation", DummyRandomPermutation)
+    perm = find_permutation("random", 99, flag=True)
+    assert isinstance(perm, DummyRandomPermutation)
+    assert perm.args == (99,)
+    assert perm.kwargs == {"flag": True}
+
+
+@pytest.mark.parametrize(
+    "name,dummy_cls",
+    [("swap", DummySwap), ("learnable", DummyOrthogonalPermutation), ("orthogonal", DummyOrthogonalPermutation)],
+)
+def test_find_permutation_by_name(monkeypatch, name, dummy_cls):
+    # Inject dummy classes for each permutation type
+    if name == "swap":
+        monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.Swap", dummy_cls)
+    else:
+        monkeypatch.setattr("bayesflow.networks.coupling_flow.permutations.OrthogonalPermutation", dummy_cls)
+    perm = find_permutation(name, "a", b="c")
+    assert isinstance(perm, dummy_cls)
+    assert perm.args == ("a",)
+    assert perm.kwargs == {"b": "c"}
+
+
+def test_find_permutation_with_keras_layer():
+    layer = keras.layers.Activation("relu")
+    perm = find_permutation(layer)
+    assert perm is layer
+
+
+def test_find_permutation_with_none():
+    res = find_permutation(None)
+    assert res is None
+
+
+def test_find_permutation_invalid_type():
+    with pytest.raises(TypeError):
+        find_permutation(3.14)
+
+
+# --- Tests for find_pooling.py ---
+
+
+def dummy_pooling_constructor(*args, **kwargs):
+    return {"args": args, "kwargs": kwargs}
+
+
+def test_find_pooling_mean():
+    pooling = find_pooling("mean")
+    # Check that a keras Lambda layer is returned
+    assert isinstance(pooling, keras.layers.Lambda)
+    # Test that the lambda function produces a mean when applied to a sample tensor.
+    import numpy as np
+
+    sample = np.array([[1, 2], [3, 4]])
+    # Keras Lambda layers expect tensors via call(), here we simply call the layer's function.
+    result = pooling.call(sample)
+    np.testing.assert_allclose(result, sample.mean(axis=-2))
+
+
+@pytest.mark.parametrize("name,func", [("max", keras.ops.max), ("min", keras.ops.min)])
+def test_find_pooling_max_min(name, func):
+    pooling = find_pooling(name)
+    assert isinstance(pooling, keras.layers.Lambda)
+    import numpy as np
+
+    sample = np.array([[1, 2], [3, 4]])
+    result = pooling.call(sample)
+    np.testing.assert_allclose(result, func(sample, axis=-2))
+
+
+def test_find_pooling_learnable(monkeypatch):
+    # Monkey patch the PoolingByMultiHeadAttention in its module
+    class DummyPoolingAttention:
+        def __init__(self, *args, **kwargs):
+            self.args = args
+            self.kwargs = kwargs
+
+    monkeypatch.setattr("bayesflow.networks.transformers.pma.PoolingByMultiHeadAttention", DummyPoolingAttention)
+    pooling = find_pooling("learnable", 7, option="test")
+    assert isinstance(pooling, DummyPoolingAttention)
+    assert pooling.args == (7,)
+    assert pooling.kwargs == {"option": "test"}
+
+
+def test_find_pooling_with_constructor():
+    # Passing a type should result in an instance.
+    class DummyPooling:
+        def __init__(self, data):
+            self.data = data
+
+    pooling = find_pooling(DummyPooling, "dummy")
+    assert isinstance(pooling, DummyPooling)
+    assert pooling.data == "dummy"
+
+
+def test_find_pooling_with_keras_layer():
+    layer = keras.layers.ReLU()
+    pooling = find_pooling(layer)
+    assert pooling is layer
+
+
+def test_find_pooling_invalid_type():
+    with pytest.raises(TypeError):
+        find_pooling(123)
+
+
+# --- Tests for find_recurrent_net.py ---
+
+
+def test_find_recurrent_net_lstm():
+    constructor = find_recurrent_net("lstm")
+    assert constructor is keras.layers.LSTM
+
+
+def test_find_recurrent_net_gru():
+    constructor = find_recurrent_net("gru")
+    assert constructor is keras.layers.GRU
+
+
+def test_find_recurrent_net_with_keras_layer():
+    layer = keras.layers.SimpleRNN(5)
+    net = find_recurrent_net(layer)
+    assert net is layer
+
+
+def test_find_recurrent_net_invalid_name():
+    with pytest.raises(ValueError):
+        find_recurrent_net("invalid_net")
+
+
+def test_find_recurrent_net_invalid_type():
+    with pytest.raises(TypeError):
+        find_recurrent_net(3.1415)

tests/utils/ecdf.py renamed to tests/test_utils/test_ecdf.py

@@ -15,7 +15,7 @@ def test_data():
 
 def test_fractional_ranks(test_data):
     post_samples, prior_samples, _ = test_data
-    # Compute expected result manually
+    # Compute the expected result manually
     expected = np.mean(post_samples < prior_samples[:, np.newaxis, :], axis=1)
     result = fractional_ranks(post_samples, prior_samples)
     np.testing.assert_almost_equal(result, expected, decimal=6)

tests/test_utils/test_optimal_transport.py

@@ -0,0 +1,53 @@
+import keras
+import pytest
+
+from bayesflow.utils import optimal_transport
+from tests.utils import assert_allclose
+
+
+@pytest.mark.jax
+def test_jit_compile():
+    import jax
+
+    x = keras.random.normal((128, 8), seed=0)
+    y = keras.random.normal((128, 8), seed=1)
+
+    ot = jax.jit(optimal_transport, static_argnames=["regularization", "seed"])
+    ot(x, y, regularization=1.0, seed=0, max_steps=10)
+
+
+def test_shapes():
+    x = keras.random.normal((128, 8), seed=0)
+    y = keras.random.normal((128, 8), seed=1)
+
+    ox, oy = optimal_transport(x, y, regularization=1.0, seed=0, max_steps=10)
+
+    assert keras.ops.shape(ox) == keras.ops.shape(x)
+    assert keras.ops.shape(oy) == keras.ops.shape(y)
+
+
+def test_transport_cost_improves():
+    x = keras.random.normal((32, 8), seed=0)
+    y = keras.random.normal((32, 8), seed=1)
+
+    before_cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))
+
+    x, y = optimal_transport(x, y, regularization=0.1, seed=0, max_steps=None)
+
+    after_cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))
+
+    assert after_cost < before_cost
+
+
+def test_assignment_is_optimal():
+    x = keras.ops.stack([keras.ops.linspace(-1, 1, 10), keras.ops.linspace(-1, 1, 10)])
+    y = keras.ops.copy(x)
+
+    # we could shuffle x and y, but flipping is a more reliable permutation
+    y = keras.ops.flip(y, axis=0)
+
+    x, y = optimal_transport(x, y, regularization=1e-3, seed=0, max_steps=1000)
+
+    cost = keras.ops.sum(keras.ops.norm(x - y, axis=-1))
+
+    assert_allclose(cost, 0.0)

tests/utils/__init__.py

@@ -1,6 +1,5 @@
 from .assertions import *
 from .callbacks import *
 from .check_combinations import *
-from .ecdf import *
 from .jupyter import *
 from .ops import *