[MRG] Fix ot.bregman.barycenter with non-uniform weights (default sinkhorn method)

RELEASES.md

@@ -18,6 +18,7 @@
 
 #### Closed issues
 
+- Fix `ot.bregman.barycenter` with the default `method="sinkhorn"` (`ot.bregman.barycenter_sinkhorn`) converging to a wrong barycenter when the weights are not uniform. The scalings started from the unweighted geometric mean, which set their weighted product to a value other than 1 for the whole run; they now start from the weighted one. Results with uniform weights are unchanged (PR #920, Issue #919)
 - Fix `ot.semidiscrete_wasserstein2_unif_circle` pairing sorted sample locations with unsorted weights, which broke permutation and rotation invariance for nonuniform weights (Issue #921)
 - Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans was taken from `perms.shape[1]` (the number of samples) instead of `perms.shape[0]`, so merging fewer plans than samples read past the array and merging more silently dropped plans. A `perms` outside the documented `(n_plans, n_samples)` layout, an empty plan stack and rows that are not permutations are now rejected (PR #911, Issue #910)
 - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264)

ot/bregman/_barycenter.py

@@ -226,7 +226,7 @@ def barycenter_sinkhorn(
 
     UKv = nx.dot(K, (A.T / nx.sum(K, axis=0)).T)
 
-    u = (geometricMean(UKv) / UKv.T).T
+    u = (geometricBar(weights, UKv) / UKv.T).T
 
     for ii in range(numItermax):
         UKv = u * nx.dot(K.T, A / nx.dot(K, u))

test/test_bregman.py

@@ -603,6 +603,31 @@ def test_barycenter(nx, method, verbose, warn):
         ot.bregman.barycenter(A_nx, M_nx, reg, log=True)
 
 
+@pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"])
+def test_barycenter_unequal_weights(nx, method):
+    # the barycenter of two 1D Gaussians lies at the weighted mean of their means
+    n_bins = 100
+    a1 = ot.datasets.make_1D_gauss(n_bins, m=20, s=5)
+    a2 = ot.datasets.make_1D_gauss(n_bins, m=60, s=8)
+    A = np.vstack((a1, a2)).T
+    M = ot.utils.dist0(n_bins)
+    M /= M.max()
+    weights = np.array([0.8, 0.2])
+    reg = 1e-3
+
+    A_nx, M_nx, weights_nx = nx.from_numpy(A, M, weights)
+    if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log":
+        pytest.skip("Log-domain barycenter is not implemented for Jax and TF")
+
+    bary = nx.to_numpy(
+        ot.bregman.barycenter(A_nx, M_nx, reg, weights_nx, method=method)
+    )
+    bary_log = ot.bregman.barycenter(A, M, reg, weights, method="sinkhorn_log")
+
+    np.testing.assert_allclose(np.sum(np.arange(n_bins) * bary), 28.0, atol=1e-2)
+    np.testing.assert_allclose(bary, bary_log, atol=1e-6)
+
+
 def test_free_support_sinkhorn_barycenter():
     measures_locations = [
         np.array([-1.0]).reshape((1, 1)),  # First dirac support