add test for assymetric cost barycenters

Author: Hicham Janati

test/test_bregman.py

@@ -490,6 +490,43 @@ def test_barycenter(nx, method, verbose, warn):
         ot.bregman.barycenter(A_nx, M_nx, reg, log=True)
 
 
+
+@pytest.mark.parametrize("method, verbose, warn",
+                         product(["sinkhorn", "sinkhorn_stabilized", "sinkhorn_log"],
+                                 [True, False], [True, False]))
+def test_barycenter_assymetric_cost(nx, method, verbose, warn):
+    n_bins = 20  # nb bins
+
+    # Gaussian distributions
+    A = ot.datasets.make_1D_gauss(n_bins, m=30, s=10)  # m= mean, s= std
+
+    # creating matrix A containing all distributions
+    A = A[:, None]
+
+    # assymetric loss matrix + normalization
+    rng = np.random.RandomState(42)
+    M = rng.randn(n_bins, n_bins) ** 2
+    M /= M.max()
+
+    A_nx, M_nx = nx.from_numpy(A, M)
+    reg = 1e-2
+
+    if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log":
+        with pytest.raises(NotImplementedError):
+            ot.bregman.barycenter(A_nx, M_nx, reg, method=method)
+    else:
+        # wasserstein
+        bary_wass_np = ot.bregman.barycenter(A, M, reg, method=method, verbose=verbose, warn=warn)
+        bary_wass, _ = ot.bregman.barycenter(A_nx, M_nx, reg, method=method, log=True)
+        bary_wass = nx.to_numpy(bary_wass)
+
+        np.testing.assert_allclose(1, np.sum(bary_wass))
+        np.testing.assert_allclose(bary_wass, bary_wass_np)
+
+        ot.bregman.barycenter(A_nx, M_nx, reg, log=True)
+
+
+
 @pytest.mark.parametrize("method, verbose, warn",
                          product(["sinkhorn", "sinkhorn_log"],
                                  [True, False], [True, False]))