diff --git a/RELEASES.md b/RELEASES.md index 20a45aef2..34390bac9 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/bregman/_barycenter.py b/ot/bregman/_barycenter.py index 9f388989c..5ca1aca87 100644 --- a/ot/bregman/_barycenter.py +++ b/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)) diff --git a/test/test_bregman.py b/test/test_bregman.py index 17b400306..41af1152e 100644 --- a/test/test_bregman.py +++ b/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