[MRG] Fix sinkhorn and stabilized barycenters with non-uniform weights

RELEASES.md

@@ -15,6 +15,7 @@
 
 #### Closed issues
 
+- Fix `ot.bregman.barycenter` with non-uniform `weights`: `method="sinkhorn"` started from the unweighted geometric mean and converged to a different barycenter than `"sinkhorn_log"`, and `method="sinkhorn_stabilized"` lost the scaling of each histogram when absorbing, returning a vector that did not sum to one for small `reg` (PR #870)
 - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)
 - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859)
 - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)

ot/bregman/_barycenter.py

@@ -226,7 +226,9 @@ def barycenter_sinkhorn(
 
     UKv = nx.dot(K, (A.T / nx.sum(K, axis=0)).T)
 
-    u = (geometricMean(UKv) / UKv.T).T
+    # the updates below keep the weighted geometric mean of u unchanged, so
+    # it has to start from the weighted one for the barycenter to be optimal
+    u = (geometricBar(weights, UKv) / UKv.T).T
 
     for ii in range(numItermax):
         UKv = u * nx.dot(K.T, A / nx.dot(K, u))
@@ -573,16 +575,22 @@ def barycenter_stabilized(
         Kv = nx.dot(K, v)
         u = A / Kv
         Ktu = nx.dot(K.T, u)
-        q = geometricBar(weights, Ktu)
+        # K has absorbed the scalings exp(alpha / reg) and exp(beta / reg), so
+        # Ktu is K^T u of the original kernel multiplied by exp(beta / reg)
+        q = nx.exp(nx.dot(nx.log(Ktu), weights) - beta / reg * nx.sum(weights))
         Q = q[:, None]
         v = Q / Ktu
         absorbing = False
         if nx.any(u > tau) or nx.any(v > tau):
             absorbing = True
-            alpha += reg * nx.log(nx.max(u, 1))
-            beta += reg * nx.log(nx.max(v, 1))
+            max_u = nx.max(u, 1)
+            max_v = nx.max(v, 1)
+            alpha += reg * nx.log(max_u)
+            beta += reg * nx.log(max_v)
             K = nx.exp((alpha[:, None] + beta[None, :] - M) / reg)
-            v = nx.ones(tuple(v.shape), type_as=v)
+            # keep the scalings of every histogram, only move their common
+            # part into K
+            v = v / max_v[:, None]
         Kv = nx.dot(K, v)
         if (
             nx.any(Ktu == 0.0)

test/test_bregman.py

@@ -603,6 +603,33 @@ def test_barycenter(nx, method, verbose, warn):
         ot.bregman.barycenter(A_nx, M_nx, reg, log=True)
 
 
+
+@pytest.mark.parametrize("reg", [1e-3, 1e-2])
+def test_barycenter_non_uniform_weights(reg):
+    # the three solvers compute the same barycenter; with non-uniform weights
+    # "sinkhorn" started from a uniform geometric mean and
+    # "sinkhorn_stabilized" lost the per-histogram scalings when absorbing
+    n_bins = 50
+    a1 = ot.datasets.make_1D_gauss(n_bins, m=15, s=5)
+    a2 = ot.datasets.make_1D_gauss(n_bins, m=35, s=5)
+    A = np.vstack((a1, a2)).T
+    M = ot.utils.dist0(n_bins)
+    M /= M.max()
+    weights = np.array([0.2, 0.8])
+
+    bars = {
+        method: ot.bregman.barycenter(
+            A, M, reg, weights, method=method, stopThr=1e-10, numItermax=20000
+        )
+        for method in ["sinkhorn", "sinkhorn_stabilized", "sinkhorn_log"]
+    }
+    for method, bar in bars.items():
+        np.testing.assert_allclose(bar.sum(), 1, err_msg=method)
+        np.testing.assert_allclose(
+            bar, bars["sinkhorn_log"], atol=1e-7, err_msg=method
+        )
+
+
 def test_free_support_sinkhorn_barycenter():
     measures_locations = [
         np.array([-1.0]).reshape((1, 1)),  # First dirac support