diff --git a/RELEASES.md b/RELEASES.md index 063701229..3001fba53 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/bregman/_barycenter.py b/ot/bregman/_barycenter.py index 9f388989c..cdd6642b8 100644 --- a/ot/bregman/_barycenter.py +++ b/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) diff --git a/test/test_bregman.py b/test/test_bregman.py index 17b400306..48995fa38 100644 --- a/test/test_bregman.py +++ b/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