add unit tests for log_sinkhorn_plan

Author: daniel-habermann

tests/test_utils/test_optimal_transport.py

@@ -41,8 +41,7 @@ def test_transport_cost_improves(method):
     assert after_cost < before_cost
 
 
-# @pytest.mark.skip(reason="too unreliable")
-@pytest.mark.parametrize("method", ["sinkhorn"])
+@pytest.mark.parametrize("method", ["log_sinkhorn", "sinkhorn"])
 def test_assignment_is_optimal(method):
     y = keras.random.normal((16, 2), seed=0)
     p = keras.random.shuffle(keras.ops.arange(keras.ops.shape(y)[0]), seed=0)
@@ -72,8 +71,8 @@ def test_assignment_aligns_with_pot():
     M = x[:, None] - y[None, :]
     M = keras.ops.norm(M, axis=-1)
 
-    pot_plan = sinkhorn_log(a, b, M, reg=1e-3, numItermax=10_000, stopThr=1e-99)
-    pot_assignments = keras.random.categorical(pot_plan, num_samples=1, seed=0)
+    pot_plan = sinkhorn_log(a, b, M, reg=1e-3, stopThr=1e-7)
+    pot_assignments = keras.random.categorical(keras.ops.log(pot_plan), num_samples=1, seed=0)
     pot_assignments = keras.ops.squeeze(pot_assignments, axis=-1)
 
     _, _, assignments = optimal_transport(x, y, regularization=1e-3, seed=0, max_steps=10_000, return_assignments=True)
@@ -107,8 +106,8 @@ def test_sinkhorn_plan_aligns_with_pot():
     b = keras.ops.ones(20) / 20
     M = euclidean(x1, x2)
 
-    pot_result = sinkhorn(a, b, M, 0.1)
-    our_result = sinkhorn_plan(x1, x2, regularization=0.1, atol=1e-8, rtol=1e-8)
+    pot_result = sinkhorn(a, b, M, 0.1, stopThr=1e-8)
+    our_result = sinkhorn_plan(x1, x2, regularization=0.1, rtol=1e-7)
 
     assert_allclose(pot_result, our_result)
 
@@ -128,3 +127,56 @@ def test_sinkhorn_plan_matches_analytical_result():
     expected = keras.ops.outer(marginal_x1, marginal_x2)
 
     assert_allclose(result, expected)
+
+
+def test_log_sinkhorn_plan_correct_marginals():
+    from bayesflow.utils.optimal_transport.log_sinkhorn import log_sinkhorn_plan
+
+    x1 = keras.random.normal((10, 2), seed=0)
+    x2 = keras.random.normal((20, 2), seed=1)
+
+    assert keras.ops.all(
+        keras.ops.isclose(keras.ops.logsumexp(log_sinkhorn_plan(x1, x2), axis=0), -keras.ops.log(20), atol=1e-3)
+    )
+    assert keras.ops.all(
+        keras.ops.isclose(keras.ops.logsumexp(log_sinkhorn_plan(x1, x2), axis=1), -keras.ops.log(10), atol=1e-3)
+    )
+
+
+def test_log_sinkhorn_plan_aligns_with_pot():
+    try:
+        from ot.bregman import sinkhorn_log
+    except (ImportError, ModuleNotFoundError):
+        pytest.skip("Need to install POT to run this test.")
+
+    from bayesflow.utils.optimal_transport.log_sinkhorn import log_sinkhorn_plan
+    from bayesflow.utils.optimal_transport.euclidean import euclidean
+
+    x1 = keras.random.normal((100, 3), seed=0)
+    x2 = keras.random.normal((200, 3), seed=1)
+
+    a = keras.ops.ones(100) / 100
+    b = keras.ops.ones(200) / 200
+    M = euclidean(x1, x2)
+
+    pot_result = keras.ops.log(sinkhorn_log(a, b, M, 0.1, stopThr=1e-7))  # sinkhorn_log returns probabilities
+    our_result = log_sinkhorn_plan(x1, x2, regularization=0.1)
+
+    assert_allclose(pot_result, our_result)
+
+
+def test_log_sinkhorn_plan_matches_analytical_result():
+    from bayesflow.utils.optimal_transport.log_sinkhorn import log_sinkhorn_plan
+
+    x1 = keras.ops.ones(16)
+    x2 = keras.ops.ones(64)
+
+    marginal_x1 = keras.ops.ones(16) / 16
+    marginal_x2 = keras.ops.ones(64) / 64
+
+    result = keras.ops.exp(log_sinkhorn_plan(x1, x2, regularization=0.1))
+
+    # If x1 and x2 are identical, the optimal plan is simply the outer product of the marginals
+    expected = keras.ops.outer(marginal_x1, marginal_x2)
+
+    assert_allclose(result, expected)