From 37719e6666f645729be618fdba01c1eae7136928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 06:32:22 +0200 Subject: [PATCH 01/28] add deiased --- ot/solvers/_linear.py | 75 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 74 insertions(+), 1 deletion(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index a550df315..404f566d9 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -619,6 +619,7 @@ def solve_sample( verbose=False, grad="autodiff", random_state=None, + debiased=False, ): r"""Solve the discrete optimal transport problem using the samples in the source and target domains. @@ -679,7 +680,14 @@ def solve_sample( method : str, optional Method for solving the problem, this can be used to select the solver for unbalanced problems (see :any:`ot.solve`), or to select a specific - large scale solver. + large scale solver. Available methods are: + - "1d" for 1D OT solver done in parallel for each dimension. + - "gaussian" for Gaussian OT solver that estimates mean and + covariance and solve the closed form solution) + - "lowrank" for low-rank Sinkhorn solver (see :any:`ot.lowrank`) + - "nystroem" for Nystroem Sinkhorn solver + - "factored" for factored OT solver + - "geomloss" for GeomLoss Sinkhorn solver n_threads : int, optional Number of OMP threads for exact OT solver, by default 1 max_iter : int, optional @@ -706,6 +714,13 @@ def solve_sample( detached. This is useful for memory saving when only the value is needed. random_state : int, optional The random state for sampling the components in each distribution for method='nystroem'. + debiased : bool, optional + Whether to use the debiased version of the OT problem, by default False + if True, the value returned is the Sinkhorn divergence but the plan is + still the Sinkhorn plan. The results for all pairs of problems and + provided in the log dictionary. if debiased='split', then X_a and X_b + are split into two halves and the debiased value is computed using the + two halves as proposed in minibtach OT. Returns ------- @@ -960,6 +975,64 @@ def solve_sample( """ + if debiased: + dict_params = dict( + metric=metric, + reg=reg, + c=c, + reg_type=reg_type, + unbalanced=unbalanced, + unbalanced_type=unbalanced_type, + lazy=lazy, + batch_size=batch_size, + method=method, + n_threads=n_threads, + max_iter=max_iter, + plan_init=plan_init, + rank=rank, + scaling=scaling, + potentials_init=potentials_init, + X_init=X_init, + tol=tol, + verbose=verbose, + grad=grad, + random_state=random_state, + debiased=False, + ) + + if isinstance(debiased, str) and debiased.lower() == "split": + raise NotImplementedError( + "Debiased OT with split is not implemented yet. Please use debiased=True for now." + ) + else: + # standard debiasing à la sinkhorn divergence + + res11 = solve_sample(X_a, X_a, a=a, b=a, **dict_params) + + res22 = solve_sample(X_b, X_b, a=b, b=b, **dict_params) + + res12 = solve_sample(X_a, X_b, a=a, b=b, **dict_params) + value = res12.value - 0.5 * (res11.value + res22.value) + value_linear = res12.value_linear - 0.5 * ( + res11.value_linear + res22.value_linear + ) + + log = {"res11": res11, "res22": res22, "res12": res12} + + res = OTResult( + value=value, + value_linear=value_linear, + plan=res12.plan, + potentials=res12.potentials, + sparse_plan=res12.sparse_plan, + lazy_plan=res12.lazy_plan, + status=res12.status, + log=log, + backend=res12.backend, + ) + # return debiased result + return res + if method is not None and method.lower() in lst_method_lazy: lazy0 = lazy lazy = True From fcc354853127a209a4ec99d2dac4dbf099dd7496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 09:45:17 +0200 Subject: [PATCH 02/28] add tests --- RELEASES.md | 2 +- ot/solvers/_linear.py | 107 ++++++++++++++++++++++++++++++++++-------- test/test_solvers.py | 25 ++++++++++ 3 files changed, 113 insertions(+), 21 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index f10d7b62b..bce595807 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -38,7 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Update the geomloss wrapper to the new version and API (PR #826) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) - +- Implement debiased OT solvers in `ot.solve_sample` #### Closed issues diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 404f566d9..c32622641 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -619,7 +619,7 @@ def solve_sample( verbose=False, grad="autodiff", random_state=None, - debiased=False, + debias=False, ): r"""Solve the discrete optimal transport problem using the samples in the source and target domains. @@ -714,7 +714,7 @@ def solve_sample( detached. This is useful for memory saving when only the value is needed. random_state : int, optional The random state for sampling the components in each distribution for method='nystroem'. - debiased : bool, optional + debias : bool, optional Whether to use the debiased version of the OT problem, by default False if True, the value returned is the Sinkhorn divergence but the plan is still the Sinkhorn plan. The results for all pairs of problems and @@ -975,7 +975,7 @@ def solve_sample( """ - if debiased: + if debias: dict_params = dict( metric=metric, reg=reg, @@ -997,39 +997,106 @@ def solve_sample( verbose=verbose, grad=grad, random_state=random_state, - debiased=False, + debias=False, ) - if isinstance(debiased, str) and debiased.lower() == "split": - raise NotImplementedError( - "Debiased OT with split is not implemented yet. Please use debiased=True for now." + nx = get_backend(X_a, X_b, a, b) + + if isinstance(debias, str) and debias.lower() == "split": + # split the samples into two halves with each half the mass + + n_a = X_a.shape[0] + n_b = X_b.shape[0] + + if a is None: + a = nx.ones(n_a, type_as=X_a) / n_a + if b is None: + b = nx.ones(n_b, type_as=X_b) / n_b + + # find the split indices + acs = nx.cumsum(a) + bcs = nx.cumsum(b) + + thr_a = 0.5 * nx.sum(a) + thr_b = 0.5 * nx.sum(b) + + idx_a = nx.searchsorted(acs, thr_a) + idx_b = nx.searchsorted(bcs, thr_b) + + # split the samples and weights + X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] + X_b1, X_b2 = X_b[: idx_b + 1], X_b[idx_b:] + + # compute weights for each half and adjust the last/first weight to sum to 0.5 + a1, a2 = a[: idx_a + 1], a[idx_a:] + a1[-1] = a1[-1] * nx.detach((thr_a - nx.sum(a1[:-1])) / a1[-1]) + a2[0] = a2[0] * nx.detach((thr_a - nx.sum(a2[1:])) / a2[0]) + b1, b2 = b[: idx_b + 1], b[idx_b:] + b1[-1] = b1[-1] * nx.detach((thr_b - nx.sum(b1[:-1])) / b1[-1]) + b2[0] = b2[0] * nx.detach((thr_b - nx.sum(b2[1:])) / b2[0]) + + # compute the four OT problems + resaa = solve_sample(X_a1, X_a2, a=a1, b=a2, **dict_params) + resbb = solve_sample(X_b1, X_b2, a=b1, b=b2, **dict_params) + resab1 = solve_sample(X_a1, X_b1, a=a1, b=b1, **dict_params) + resab2 = solve_sample(X_a2, X_b2, a=a2, b=b2, **dict_params) + + # compute debiased values + value = 0.5 * (resab1.value + resab2.value) - 0.5 * ( + resaa.value + resbb.value ) + value_linear = 0.5 * (resab1.value_linear + resab2.value_linear) - 0.5 * ( + resaa.value_linear + resbb.value_linear + ) + + log = { + "res_aa": resaa, + "res_bb": resbb, + "res_ab1": resab1, + "res_ab2": resab2, + } + + res = OTResult( + value=value, + value_linear=value_linear, + plan=None, # no plan for debiased version + potentials=None, # no potentials for debiased version + sparse_plan=None, # no sparse plan for debiased version + lazy_plan=None, # no lazy plan for debiased version + status="Debiased", + log=log, + backend=nx, + ) + else: # standard debiasing à la sinkhorn divergence - res11 = solve_sample(X_a, X_a, a=a, b=a, **dict_params) + resaa = solve_sample(X_a, X_a, a=a, b=a, **dict_params) - res22 = solve_sample(X_b, X_b, a=b, b=b, **dict_params) + resbb = solve_sample(X_b, X_b, a=b, b=b, **dict_params) - res12 = solve_sample(X_a, X_b, a=a, b=b, **dict_params) - value = res12.value - 0.5 * (res11.value + res22.value) - value_linear = res12.value_linear - 0.5 * ( - res11.value_linear + res22.value_linear + resab = solve_sample(X_a, X_b, a=a, b=b, **dict_params) + + # compute debiased values + value = resab.value - 0.5 * (resaa.value + resbb.value) + value_linear = resab.value_linear - 0.5 * ( + resaa.value_linear + resbb.value_linear ) - log = {"res11": res11, "res22": res22, "res12": res12} + log = {"res_aa": resaa, "res_bb": resbb, "res_ab": resab} res = OTResult( value=value, value_linear=value_linear, - plan=res12.plan, - potentials=res12.potentials, - sparse_plan=res12.sparse_plan, - lazy_plan=res12.lazy_plan, - status=res12.status, + plan=resab.plan, + potentials=resab.potentials, + sparse_plan=resab.sparse_plan, + lazy_plan=resab.lazy_plan, + status=resab.status, log=log, - backend=res12.backend, + backend=nx, ) + # return debiased result return res diff --git a/test/test_solvers.py b/test/test_solvers.py index 505cdb3af..30ef9b071 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -797,6 +797,31 @@ def test_solve_sample_methods(nx, method_params): np.testing.assert_allclose(sol2.value, 0, atol=1e-10) +@pytest.mark.parametrize("debias", [True, False, "split"]) +def test_solve_sample_debias(nx, debias): + n_samples_s = 10 + n_samples_t = 9 + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + xb, yb, ab, bb = nx.from_numpy(x, y, a, b) + + sol = ot.solve_sample(x, y, reg=10, debias=debias) + solb = ot.solve_sample(xb, yb, ab, bb, reg=10, debias=debias) + + # check some attributes (no need ) + assert_allclose_sol(sol, solb) + + sol2 = ot.solve_sample(x, x, reg=10, debias=True) + + np.testing.assert_allclose(sol2.value, 0, atol=1e-10) + + @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): n_samples_s = 20 From f208f9f333564d61e78728f94f9940252c81e6a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 09:50:43 +0200 Subject: [PATCH 03/28] add references --- README.md | 18 ++++++++++-------- RELEASES.md | 2 +- ot/solvers/_linear.py | 11 +++++++++-- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 4ce4082ab..950a509e9 100644 --- a/README.md +++ b/README.md @@ -455,18 +455,20 @@ Artificial Intelligence. \[82] Bonet, C., Nadjahi, K., Séjourné, T., Fatras, K., & Courty, N. (2024). [Slicing Unbalanced Optimal Transport](https://openreview.net/forum?id=AjJTg5M0r8). Transactions on Machine Learning Research. -[83] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920). +\[83] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920). -[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15. +\[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15. -[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385. +\[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385. -[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661. +\[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661. -[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations. +\[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations. -[88] Bouveyron, C. & Corneli, M. (2026). [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88. +\[88] Bouveyron, C. & Corneli, M. (2026). [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88. -[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. +\[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. -[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. +\[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. + +\[91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & Courty, N. (2021). [Minibatch optimal transport distances; analysis and applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. \ No newline at end of file diff --git a/RELEASES.md b/RELEASES.md index bce595807..340d01c79 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -38,7 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Update the geomloss wrapper to the new version and API (PR #826) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) -- Implement debiased OT solvers in `ot.solve_sample` +- Implement debiased OT solvers in `ot.solve_sample`. #### Closed issues diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index c32622641..afb774e94 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -716,11 +716,11 @@ def solve_sample( The random state for sampling the components in each distribution for method='nystroem'. debias : bool, optional Whether to use the debiased version of the OT problem, by default False - if True, the value returned is the Sinkhorn divergence but the plan is + if True, the value returned is the Sinkhorn divergence [23] but the plan is still the Sinkhorn plan. The results for all pairs of problems and provided in the log dictionary. if debiased='split', then X_a and X_b are split into two halves and the debiased value is computed using the - two halves as proposed in minibtach OT. + two halves as proposed in minibatch OT [91]. Returns ------- @@ -951,6 +951,10 @@ def solve_sample( Optimal Transport. Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics (AISTATS). + ..[23] Aude, G., Peyré, G., Cuturi, M., Learning Generative Models with + Sinkhorn Divergences, Proceedings of the Twenty-First International + Conference on Artificial Intelligence and Statistics, (AISTATS) 21, 2018 + .. [34] Feydy, J., Séjourné, T., Vialard, F. X., Amari, S. I., Trouvé, A., & Peyré, G. (2019, April). Interpolating between optimal transport and MMD using Sinkhorn divergences. In The 22nd International Conference @@ -972,6 +976,9 @@ def solve_sample( .. [80] Altschuler, J., Bach, F., Rudi, A., Niles-Weed, J. (2019). Massively scalable Sinkhorn distances via the Nyström method. NeurIPS. + .. [91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & + Courty, N. (2021). [Minibatch optimal transport distances; analysis and + applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. """ From dc2eb2e48d097af8b65aa10b5220d08d36d8ba1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 11:23:29 +0200 Subject: [PATCH 04/28] fix tojax and tf? --- ot/solvers/_linear.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index afb774e94..ee04b9816 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1036,11 +1036,16 @@ def solve_sample( # compute weights for each half and adjust the last/first weight to sum to 0.5 a1, a2 = a[: idx_a + 1], a[idx_a:] - a1[-1] = a1[-1] * nx.detach((thr_a - nx.sum(a1[:-1])) / a1[-1]) - a2[0] = a2[0] * nx.detach((thr_a - nx.sum(a2[1:])) / a2[0]) + a1_1 = a1[idx_a] * nx.detach((thr_a - nx.sum(a1[:idx_a])) / a1[-1]) + a2_0 = a2[0] * nx.detach((thr_a - nx.sum(a2[1:])) / a2[0]) + a1 = nx.concatenate([a1[:idx_a], nx.reshape(a1_1, (1,))]) + a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a2[1:]]) + b1, b2 = b[: idx_b + 1], b[idx_b:] - b1[-1] = b1[-1] * nx.detach((thr_b - nx.sum(b1[:-1])) / b1[-1]) - b2[0] = b2[0] * nx.detach((thr_b - nx.sum(b2[1:])) / b2[0]) + b1_1 = b1[idx_b] * nx.detach((thr_b - nx.sum(b1[:idx_b])) / b1[-1]) + b2_0 = b2[0] * nx.detach((thr_b - nx.sum(b2[1:])) / b2[0]) + b1 = nx.concatenate([b1[:idx_b], nx.reshape(b1_1, (1,))]) + b2 = nx.concatenate([nx.reshape(b2_0, (1,)), b2[1:]]) # compute the four OT problems resaa = solve_sample(X_a1, X_a2, a=a1, b=a2, **dict_params) From af8144402fafe15547d57b61dafab79ae1e54012 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 13:00:20 +0200 Subject: [PATCH 05/28] try again tf --- ot/solvers/_linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index ee04b9816..4e1544d18 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1036,13 +1036,13 @@ def solve_sample( # compute weights for each half and adjust the last/first weight to sum to 0.5 a1, a2 = a[: idx_a + 1], a[idx_a:] - a1_1 = a1[idx_a] * nx.detach((thr_a - nx.sum(a1[:idx_a])) / a1[-1]) + a1_1 = a1[idx_a] * nx.detach((thr_a - nx.sum(a1[:idx_a])) / a1[idx_a]) a2_0 = a2[0] * nx.detach((thr_a - nx.sum(a2[1:])) / a2[0]) a1 = nx.concatenate([a1[:idx_a], nx.reshape(a1_1, (1,))]) a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a2[1:]]) b1, b2 = b[: idx_b + 1], b[idx_b:] - b1_1 = b1[idx_b] * nx.detach((thr_b - nx.sum(b1[:idx_b])) / b1[-1]) + b1_1 = b1[idx_b] * nx.detach((thr_b - nx.sum(b1[:idx_b])) / b1[idx_b]) b2_0 = b2[0] * nx.detach((thr_b - nx.sum(b2[1:])) / b2[0]) b1 = nx.concatenate([b1[:idx_b], nx.reshape(b1_1, (1,))]) b2 = nx.concatenate([nx.reshape(b2_0, (1,)), b2[1:]]) From 1b63e633d7c84be576f5539e55b9d25e1c7b494d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 15:36:19 +0200 Subject: [PATCH 06/28] tests? --- ot/solvers/_linear.py | 76 +++++++++++++++++++++++-------------------- 1 file changed, 41 insertions(+), 35 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 4e1544d18..3abd0ca7f 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1011,41 +1011,8 @@ def solve_sample( if isinstance(debias, str) and debias.lower() == "split": # split the samples into two halves with each half the mass - - n_a = X_a.shape[0] - n_b = X_b.shape[0] - - if a is None: - a = nx.ones(n_a, type_as=X_a) / n_a - if b is None: - b = nx.ones(n_b, type_as=X_b) / n_b - - # find the split indices - acs = nx.cumsum(a) - bcs = nx.cumsum(b) - - thr_a = 0.5 * nx.sum(a) - thr_b = 0.5 * nx.sum(b) - - idx_a = nx.searchsorted(acs, thr_a) - idx_b = nx.searchsorted(bcs, thr_b) - - # split the samples and weights - X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] - X_b1, X_b2 = X_b[: idx_b + 1], X_b[idx_b:] - - # compute weights for each half and adjust the last/first weight to sum to 0.5 - a1, a2 = a[: idx_a + 1], a[idx_a:] - a1_1 = a1[idx_a] * nx.detach((thr_a - nx.sum(a1[:idx_a])) / a1[idx_a]) - a2_0 = a2[0] * nx.detach((thr_a - nx.sum(a2[1:])) / a2[0]) - a1 = nx.concatenate([a1[:idx_a], nx.reshape(a1_1, (1,))]) - a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a2[1:]]) - - b1, b2 = b[: idx_b + 1], b[idx_b:] - b1_1 = b1[idx_b] * nx.detach((thr_b - nx.sum(b1[:idx_b])) / b1[idx_b]) - b2_0 = b2[0] * nx.detach((thr_b - nx.sum(b2[1:])) / b2[0]) - b1 = nx.concatenate([b1[:idx_b], nx.reshape(b1_1, (1,))]) - b2 = nx.concatenate([nx.reshape(b2_0, (1,)), b2[1:]]) + X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_samples_ratio(X_a, a, nx) + X_b1, X_b2, b1, b2, sel_b1, sel_b2 = split_samples_ratio(X_b, b, nx) # compute the four OT problems resaa = solve_sample(X_a1, X_a2, a=a1, b=a2, **dict_params) @@ -1066,6 +1033,14 @@ def solve_sample( "res_bb": resbb, "res_ab1": resab1, "res_ab2": resab2, + "sel_a1": sel_a1, + "sel_a2": sel_a2, + "sel_b1": sel_b1, + "sel_b2": sel_b2, + "a1": a1, + "a2": a2, + "b1": b1, + "b2": b2, } res = OTResult( @@ -1415,3 +1390,34 @@ def solve_sample( log=log, ) return res + + +def split_samples_ratio(X_a, a, nx, ratio=0.5): + "returns a split of the samples and weights according to a ratio" + + n_a = X_a.shape[0] + + if a is None: + a = nx.ones(n_a, type_as=X_a) / n_a + + # find the split indices + acs = nx.cumsum(a) + + thr_a = ratio * nx.sum(a) + + idx_a = nx.searchsorted(acs, thr_a, side="right") + + # split the samples and weights + X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] + + # compute weights for each half and adjust the last/first weight to sum to 0.5 + a01, a02 = a[:idx_a], a[idx_a + 1 :] + v_idx = a[idx_a] + a1_1 = nx.maximum(v_idx * nx.detach((thr_a - nx.sum(a01)) / v_idx), 0) + a2_0 = nx.maximum(v_idx * nx.detach((nx.sum(a) - thr_a - nx.sum(a02)) / v_idx), 0) + a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) + a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) + sel_a1 = slice(0, idx_a + 1) + sel_a2 = slice(idx_a, n_a) + + return X_a1, X_a2, a1, a2, sel_a1, sel_a2 From 73fd3b1a923554382bf0dc4b5f7b094d247f3239 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 16:44:39 +0200 Subject: [PATCH 07/28] add exmaple and remove tf implementation --- examples/plot_debias_sink_div.py | 252 +++++++++++++++++++++++++++++++ ot/solvers/_linear.py | 31 +++- test/test_solvers.py | 1 + 3 files changed, 281 insertions(+), 3 deletions(-) create mode 100644 examples/plot_debias_sink_div.py diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py new file mode 100644 index 000000000..e77ebcd11 --- /dev/null +++ b/examples/plot_debias_sink_div.py @@ -0,0 +1,252 @@ +# -*- coding: utf-8 -*- +""" +====================================== +Sinkhorn Divergence and Debiased OT solvers +====================================== + +This example shows how to use the debiased OT solvers in `ot.solve_sample` to +compute Sinkhorn divergences and debiased Minibtach solutions. The debiased OT solvers +can be used with balanced and unbalanced OT problems, and with different +regularization types (entropic, L2, group lasso). +""" + +# Author: Remi Flamary +# +# License: MIT License +# sphinx_gallery_thumbnail_number = 3 + +# %% + +import numpy as np +import matplotlib.pylab as pl +import ot +import ot.plot +from ot.datasets import make_1D_gauss as gauss + +############################################################################## +# Generate data +# ------------- + + +# %% +def sample_ball(n, radius=1.0, center=(0.0, 0.0)): + np.random.seed(0) + theta = 2 * np.pi * np.random.rand(n) + r = radius * np.sqrt(np.random.rand(n)) + + x = r * np.cos(theta) + center[0] + y = r * np.sin(theta) + center[1] + + return np.stack((x, y), axis=1) + + +def sample_two_balls(n, radius=1.0, sep=1): + assert n % 2 == 0, "n must be even" + centers = ((-sep, -sep), (sep, sep)) + n_half = n // 2 + X1 = sample_ball(n_half, radius, centers[0]) + X2 = sample_ball(n_half, radius, centers[1]) + + perm = np.random.permutation(n_half * 2) + X = np.vstack((X1, X2)) + X = X[perm] + return X + + +n = 50 + +x1 = sample_ball(n, radius=1.0, center=(0, 0)) +x2 = sample_two_balls(n, radius=1.0, sep=1.5) + +pl.figure(1, figsize=(5, 5)) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", alpha=0.7) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", alpha=0.7) +pl.legend() +pl.title("Two distributions") +ax = pl.axis() + + +############################################################################## +# Compute Sinkhorn divergence and visualize plans +# ----------------------------------------------- +# The Sinkhorn divergence is computed by setting the `debias` parameter to +# `True` in the `ot.solve_sample` function. The resulting value is the Sinkhorn +# divergence. The Sinkhorn divergences is computed as: +# +# .. math:: +# S_\epsilon(\mu, \nu) = OT_\epsilon(\mu, \nu) - \frac{1}{2} OT_\epsilon(\mu, \mu) - \frac{1}{2} OT_\epsilon(\nu, \nu) +# +# The entropic OT plans for each of those terms can be accessed in the `log` +# attribute of the result, and can be visualized using the +# `ot.plot.plot2D_samples_mat` function. + +res = ot.solve_sample(x1, x2, reg=0.1, debias=True) + +print("Sinkhorn divergence: ", res.value) + +plan_11 = res.log["res_aa"].plan +plan_12 = res.log["res_ab"].plan +plan_22 = res.log["res_bb"].plan + +# +pl.figure(2, figsize=(15, 5)) + +pl.subplot(1, 3, 1) +ot.plot.plot2D_samples_mat(x1, x1, plan_11, thr=0.05) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) +pl.axis(ax) +pl.title("Plan between source and source") +pl.subplot(1, 3, 2) +ot.plot.plot2D_samples_mat(x1, x2, plan_12, thr=0.05) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", zorder=2) +pl.axis(ax) +pl.title("Plan between source and target") +pl.subplot(1, 3, 3) +ot.plot.plot2D_samples_mat(x2, x2, plan_22, thr=0.05) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", color="C1", zorder=2) +pl.axis(ax) +pl.title("Plan between target and target") + + +############################################################################## +# Debiased Minibatch OT +# --------------------------------- +# +# Doing OT on minibatches leads to a similar bias than using entropic +# regularization since the average OT plan is densified to to the stochasticity +# of the minibatch sampling. On a given minibatch, the debiased loss can be +# computed by setting the `debias` parameter to `'split'`that split the data +# points in each distributions in two and computes the debias OT loss as: +# +# .. math:: +# \tilde{OT}_m(\mu, \nu) = \frac{1}{2}(\hat{OT}_m(\mu_1, \nu_1) + \hat{OT}_m(\mu_2, \nu_2) - \hat{OT}_m(\nu_1, \nu_2) - \hat{OT}_m(\mu_1, \nu_2)) +# + +# %% solve OT minibtach and visualize the plans + +res = ot.solve_sample(x1, x2, debias="split") + +print("Debiased minibatch OT loss: ", res.value) + +# recover the plans for each of the four terms in the debiased loss +plan_11 = res.log["res_aa"].plan +plan_12 = res.log["res_ab1"].plan +plan_21 = res.log["res_ab2"].plan +plan_22 = res.log["res_bb"].plan +sel_a1 = res.log["sel_a1"] +sel_a2 = res.log["sel_a2"] +sel_b1 = res.log["sel_b1"] +sel_b2 = res.log["sel_b2"] + +nb1 = plan_11.shape[0] +nb2 = plan_22.shape[0] + +pl.figure(4, figsize=(15, 3)) + +pl.subplot(1, 4, 1) +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="$\mu_1$", zorder=2) +pl.scatter( + x1[sel_a2, 0], x1[sel_a2, 1], label=r"$\mu_2$", zorder=2, color="C0", alpha=0.5 +) +pl.scatter(x2[sel_b1, 0], x2[sel_b1, 1], label=r"$\nu_1$", zorder=2, color="C1") +pl.scatter( + x2[sel_b2, 0], x2[sel_b2, 1], label=r"$\nu_2$", zorder=2, color="C1", alpha=0.5 +) +pl.title("Minibatch split") +pl.axis(ax) +pl.legend() + + +pl.subplot(1, 4, 2) +ot.plot.plot2D_samples_mat(x1[sel_a1], x1[sel_a2], plan_11, thr=0.05) +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2) +pl.scatter( + x1[sel_a2, 0], + x1[sel_a2, 1], + label="Source distribution", + zorder=2, + color="C0", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between source and source") +pl.subplot(1, 4, 3) +ot.plot.plot2D_samples_mat(x1[sel_a1], x2[sel_b1], plan_12, thr=0.05) +ot.plot.plot2D_samples_mat(x1[sel_a2], x2[sel_b2], plan_21, thr=0.05, alpha=0.5) + +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2) +pl.scatter( + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" +) +pl.scatter( + x1[sel_a2, 0], + x1[sel_a2, 1], + label="Source distribution", + zorder=2, + color="C0", + alpha=0.5, +) +pl.scatter( + x2[sel_b2, 0], + x2[sel_b2, 1], + label="Target distribution", + zorder=2, + color="C1", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between source and target") + +pl.subplot(1, 4, 4) +ot.plot.plot2D_samples_mat(x2[sel_b1], x2[sel_b2], plan_22, thr=0.05) +pl.scatter( + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" +) +pl.scatter( + x2[sel_b2, 0], + x2[sel_b2, 1], + label="Target distribution", + zorder=2, + color="C1", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between target and target") + + +############################################################################## +# Comparison of the divergences +# ------------------------------------------------- + +# %% move a distribution and compute Sinkhorn divergence and Sinkhorn distance +reg = 0.1 + +sep_list = np.linspace(0, 1.0, 10) +sink_list = [] +sink_div_list = [] +ot_mb_list = [] +ot_mb_sink_list = [] +for sep in sep_list: + x2sep = sample_two_balls(n, radius=1.0, sep=sep) + sink_list.append( + ot.solve_sample( + x1, + x2sep, + reg=reg, + ).value + ) + sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value) + ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value) + ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value) + +pl.figure(3) +pl.plot(sep_list, sink_list, label="Sinkhorn loss", color="C0") +pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") +pl.plot(sep_list, ot_mb_list, label="Minibatch OT", color="C2") +pl.plot(sep_list, ot_mb_sink_list, label="Minibatch Sinkhorn", color="C3") +pl.xlabel("Separation between distributions") +pl.ylabel("Loss/Divergence") +pl.title("Comparison of biased VS debiased OT losses") +pl.grid() +pl.legend() diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 3abd0ca7f..a5c34ce65 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -463,6 +463,13 @@ def solve( if reg_type.lower() == "entropy": value = value_linear + reg * nx.sum(plan * nx.log(plan + 1e-16)) else: + # stabilize kl of 0 mass + if nx.any(a == 0) or nx.any(b == 0): + a = a + 1e-16 * nx.min(a) + b = b + 1e-16 * nx.min(b) + print( + "Warning: a or b has 0 mass, adding small mass to avoid NaN in KL divergence." + ) value = value_linear + reg * nx.kl_div( plan, a[:, None] * b[None, :] ) @@ -1009,6 +1016,11 @@ def solve_sample( nx = get_backend(X_a, X_b, a, b) + if nx.__name__ == "tf": + raise NotImplementedError( + "Debiasing is not implemented for TensorFlow backend." + ) + if isinstance(debias, str) and debias.lower() == "split": # split the samples into two halves with each half the mass X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_samples_ratio(X_a, a, nx) @@ -1415,9 +1427,22 @@ def split_samples_ratio(X_a, a, nx, ratio=0.5): v_idx = a[idx_a] a1_1 = nx.maximum(v_idx * nx.detach((thr_a - nx.sum(a01)) / v_idx), 0) a2_0 = nx.maximum(v_idx * nx.detach((nx.sum(a) - thr_a - nx.sum(a02)) / v_idx), 0) - a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) - a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) - sel_a1 = slice(0, idx_a + 1) + # concat mass or remove samples if mass is 0 + if a1_1 > 0: + a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) + sel_a1 = slice(0, idx_a + 1) + else: + X_a1 = X_a[:idx_a] + sel_a1 = slice(0, idx_a) + a1 = a01 + + if a2_0 > 0: + a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) + sel_a2 = slice(idx_a, n_a) + else: + X_a2 = X_a[idx_a + 1 :] + sel_a2 = slice(idx_a + 1, n_a) + a2 = a02 sel_a2 = slice(idx_a, n_a) return X_a1, X_a2, a1, a2, sel_a1, sel_a2 diff --git a/test/test_solvers.py b/test/test_solvers.py index 30ef9b071..39d96aba2 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -797,6 +797,7 @@ def test_solve_sample_methods(nx, method_params): np.testing.assert_allclose(sol2.value, 0, atol=1e-10) +@pytest.skip_backend("tf", reason="Not implemented for tf backend") @pytest.mark.parametrize("debias", [True, False, "split"]) def test_solve_sample_debias(nx, debias): n_samples_s = 10 From 4ea4922207b3d1b66b77e4f0cc570463c32e8967 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Thu, 2 Jul 2026 16:46:17 +0200 Subject: [PATCH 08/28] chante method name --- examples/plot_debias_sink_div.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py index e77ebcd11..0cf6cb013 100644 --- a/examples/plot_debias_sink_div.py +++ b/examples/plot_debias_sink_div.py @@ -243,8 +243,8 @@ def sample_two_balls(n, radius=1.0, sep=1): pl.figure(3) pl.plot(sep_list, sink_list, label="Sinkhorn loss", color="C0") pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") -pl.plot(sep_list, ot_mb_list, label="Minibatch OT", color="C2") -pl.plot(sep_list, ot_mb_sink_list, label="Minibatch Sinkhorn", color="C3") +pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2") +pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3") pl.xlabel("Separation between distributions") pl.ylabel("Loss/Divergence") pl.title("Comparison of biased VS debiased OT losses") From 00de7bd218e2391e172fb6af5d2026265341281c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 10:07:29 +0200 Subject: [PATCH 09/28] debug split function --- examples/plot_debias_sink_div.py | 3 +-- ot/solvers/_linear.py | 8 +++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py index 0cf6cb013..90dde4b96 100644 --- a/examples/plot_debias_sink_div.py +++ b/examples/plot_debias_sink_div.py @@ -160,11 +160,10 @@ def sample_two_balls(n, radius=1.0, sep=1): pl.subplot(1, 4, 2) ot.plot.plot2D_samples_mat(x1[sel_a1], x1[sel_a2], plan_11, thr=0.05) -pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2) +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], zorder=2) pl.scatter( x1[sel_a2, 0], x1[sel_a2, 1], - label="Source distribution", zorder=2, color="C0", alpha=0.5, diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index a5c34ce65..91d92e85a 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1404,8 +1404,11 @@ def solve_sample( return res -def split_samples_ratio(X_a, a, nx, ratio=0.5): - "returns a split of the samples and weights according to a ratio" +def split_samples_ratio(X_a, a, nx=None, ratio=0.5): + "Split distribution according to a ratio of weights (using point ordering)." + + if nx is None: + nx = get_backend(X_a, a) n_a = X_a.shape[0] @@ -1443,6 +1446,5 @@ def split_samples_ratio(X_a, a, nx, ratio=0.5): X_a2 = X_a[idx_a + 1 :] sel_a2 = slice(idx_a + 1, n_a) a2 = a02 - sel_a2 = slice(idx_a, n_a) return X_a1, X_a2, a1, a2, sel_a1, sel_a2 From b2cecee9865a00695055c68f1d8627449c67f3fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 10:30:00 +0200 Subject: [PATCH 10/28] cretae proper function tets and and debug --- ot/solvers/_linear.py | 52 ++-------------------- ot/utils.py | 101 ++++++++++++++++++++++++++++++++++++++++++ test/test_utils.py | 35 +++++++++++++++ 3 files changed, 139 insertions(+), 49 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 91d92e85a..64268e4ce 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -8,7 +8,7 @@ # # License: MIT License -from ..utils import OTResult, dist +from ..utils import OTResult, dist, split_sample_ratio from ..lp import emd2, emd2_lazy, wasserstein_1d from ..backend import get_backend from ..unbalanced import mm_unbalanced, sinkhorn_knopp_unbalanced, lbfgsb_unbalanced @@ -1023,8 +1023,8 @@ def solve_sample( if isinstance(debias, str) and debias.lower() == "split": # split the samples into two halves with each half the mass - X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_samples_ratio(X_a, a, nx) - X_b1, X_b2, b1, b2, sel_b1, sel_b2 = split_samples_ratio(X_b, b, nx) + X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_sample_ratio(X_a, a, nx=nx) + X_b1, X_b2, b1, b2, sel_b1, sel_b2 = split_sample_ratio(X_b, b, nx=nx) # compute the four OT problems resaa = solve_sample(X_a1, X_a2, a=a1, b=a2, **dict_params) @@ -1402,49 +1402,3 @@ def solve_sample( log=log, ) return res - - -def split_samples_ratio(X_a, a, nx=None, ratio=0.5): - "Split distribution according to a ratio of weights (using point ordering)." - - if nx is None: - nx = get_backend(X_a, a) - - n_a = X_a.shape[0] - - if a is None: - a = nx.ones(n_a, type_as=X_a) / n_a - - # find the split indices - acs = nx.cumsum(a) - - thr_a = ratio * nx.sum(a) - - idx_a = nx.searchsorted(acs, thr_a, side="right") - - # split the samples and weights - X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] - - # compute weights for each half and adjust the last/first weight to sum to 0.5 - a01, a02 = a[:idx_a], a[idx_a + 1 :] - v_idx = a[idx_a] - a1_1 = nx.maximum(v_idx * nx.detach((thr_a - nx.sum(a01)) / v_idx), 0) - a2_0 = nx.maximum(v_idx * nx.detach((nx.sum(a) - thr_a - nx.sum(a02)) / v_idx), 0) - # concat mass or remove samples if mass is 0 - if a1_1 > 0: - a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) - sel_a1 = slice(0, idx_a + 1) - else: - X_a1 = X_a[:idx_a] - sel_a1 = slice(0, idx_a) - a1 = a01 - - if a2_0 > 0: - a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) - sel_a2 = slice(idx_a, n_a) - else: - X_a2 = X_a[idx_a + 1 :] - sel_a2 = slice(idx_a + 1, n_a) - a2 = a02 - - return X_a1, X_a2, a1, a2, sel_a1, sel_a2 diff --git a/ot/utils.py b/ot/utils.py index f79d69809..fcd823415 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -2047,3 +2047,104 @@ def fun_numpy(x): return nx.to_numpy(fun(nx.from_numpy(x))) return fun_numpy + + +def split_sample_ratio( + X_a, a=None, ratio=0.5, random_split=False, random_state=None, nx=None +): + """Split distribution according to a ratio of weights (using point ordering). + + + Parameters + ---------- + X_a : array-like, shape (n_samples_a, dim) + samples in the source domain + a : array-like, shape (dim_a,), optional + Samples weights in the source domain (default is uniform) + nx : backend, optional + Backend for array operations, by default None (auto-detect) + ratio : float, optional + Ratio of the split, by default 0.5 + random_split : bool, optional + Whether to split randomly, by default False + random_state : int, optional + Random state for reproducibility, by default None + + Returns + ------- + + X_a1 : array-like, shape (n_samples_a1, dim) + First half of the samples in the source domain + X_a2 : array-like, shape (n_samples_a2, dim) + Second half of the samples in the source domain + a1 : array-like, shape (dim_a1,) + First half of the weights in the source domain + a2 : array-like, shape (dim_a2,) + Second half of the weights in the source domain + sel_a1 : slice, ndarray-like + Slice or indexes for the first half of the samples in the source domain + sel_a2 : slice, ndarray-like + Slice or indexes for the second half of the samples in the source domain + """ + + if ratio < 0 or ratio > 1: + raise ValueError("ratio should be in [0, 1]") + + if nx is None: + nx = get_backend(X_a, a) + + n_a = X_a.shape[0] + + if a is None: + a = nx.ones(n_a, type_as=X_a) / n_a + + if random_split: + if random_state is not None: + nx.seed(random_state) + perm = nx.randperm(n_a, type_as=X_a) + X_a = X_a[perm] + a = a[perm] + + # find the split indices + acs = nx.cumsum(a) + thr_a = ratio * nx.sum(a) + idx_a = nx.searchsorted(acs, thr_a, side="right") + + # split the samples and weights + X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] + + # compute weights for each half and adjust the last/first weight to sum to 0.5 + a01, a02 = a[:idx_a], a[idx_a + 1 :] + v_idx = a[idx_a] + a1_1 = nx.maximum(v_idx * nx.detach((thr_a - nx.sum(a01)) / v_idx), 0) + a2_0 = nx.maximum(v_idx * nx.detach((nx.sum(a) - thr_a - nx.sum(a02)) / v_idx), 0) + # concat mass or remove samples if mass is 0 + if a1_1 > 0: + a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) + if random_split: + sel_a1 = perm[: idx_a + 1] + else: + sel_a1 = slice(0, idx_a + 1) + else: + X_a1 = X_a[:idx_a] + if random_split: + sel_a1 = perm[:idx_a] + else: + sel_a1 = slice(0, idx_a) + a1 = a01 + + if a2_0 > 0: + a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) + if random_split: + sel_a2 = perm[idx_a:] + else: + sel_a2 = slice(idx_a, n_a) + else: + X_a2 = X_a[idx_a + 1 :] + if random_split: + sel_a2 = perm[idx_a + 1 :] + else: + sel_a2 = slice(idx_a + 1, n_a) + a2 = a02 + + return X_a1, X_a2, a1, a2, sel_a1, sel_a2 diff --git a/test/test_utils.py b/test/test_utils.py index 05885a1c2..d66bb6503 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -957,3 +957,38 @@ def test_datascaler_backend_mismatch_raises(): scaler._nx = object() with pytest.raises(ValueError, match="Backend mismatch"): scaler.transform(X) + + +@pytest.mark.parametrize("ratio", [0.3, 0.5, 0.51, 0.7]) +@pytest.mark.parametrize("n", [10, 50, 100, 101]) +@pytest.mark.parametrize("random", [True, False]) +def test_split_sample_ratio(nx, ratio, n, random): + rng = np.random.RandomState(0) + X = rng.normal(0, 1, (100, 3)) + a = rng.uniform(size=(100,)) + a = a / np.sum(a) + X = nx.from_numpy(X) + a = nx.from_numpy(a) + + seed = 42 + + # test uniform weights + X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( + X, ratio=ratio, random_state=seed, nx=nx + ) + + np.testing.assert_allclose(nx.to_numpy(a1).sum(), ratio, atol=1e-5) + np.testing.assert_allclose(nx.to_numpy(a2).sum(), 1 - ratio, atol=1e-5) + + assert X1.shape[0] == a1.shape[0] + assert X2.shape[0] == a2.shape[0] + + # test non uniform weights + X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( + X, ratio=ratio, a=a, random_state=seed + ) + + with pytest.raises(ValueError): + ot.utils.split_sample_ratio(X, ratio=1.5, a=a, random_state=seed, nx=nx) + + # From 7f5bc6c0e7e5d06723fcf41eb80b70506a53361c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 10:50:26 +0200 Subject: [PATCH 11/28] remoe split random for the moement --- examples/plot_debias_sink_div.py | 7 +++++++ ot/backend.py | 5 ++--- ot/solvers/_linear.py | 2 +- test/test_solvers.py | 11 ++++------- 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py index 90dde4b96..382b8f568 100644 --- a/examples/plot_debias_sink_div.py +++ b/examples/plot_debias_sink_div.py @@ -225,6 +225,7 @@ def sample_two_balls(n, radius=1.0, sep=1): sink_list = [] sink_div_list = [] ot_mb_list = [] +ot_mb_rand = [] ot_mb_sink_list = [] for sep in sep_list: x2sep = sample_two_balls(n, radius=1.0, sep=sep) @@ -237,6 +238,9 @@ def sample_two_balls(n, radius=1.0, sep=1): ) sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value) ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value) + ot_mb_rand.append( + ot.solve_sample(x1, x2sep, debias="split_random", random_state=41).value + ) ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value) pl.figure(3) @@ -244,8 +248,11 @@ def sample_two_balls(n, radius=1.0, sep=1): pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2") pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3") +pl.plot(sep_list, ot_mb_rand, label="Debiased MB OT (random)", color="C4") pl.xlabel("Separation between distributions") pl.ylabel("Loss/Divergence") pl.title("Comparison of biased VS debiased OT losses") pl.grid() pl.legend() + +# %% diff --git a/ot/backend.py b/ot/backend.py index 08ab22bdd..d210576b7 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -2444,7 +2444,6 @@ def randperm(self, size, type_as=None): ) return torch.randperm( n=size, - dtype=type_as.dtype, generator=generator, device=type_as.device, ) @@ -2909,7 +2908,7 @@ def randperm(self, size, type_as=None): return self.rng_.permutation(size) else: with cp.cuda.Device(type_as.device): - return cp.asarray(self.rng_.permutation(size), dtype=type_as.dtype) + return cp.asarray(self.rng_.permutation(size)) def coo_matrix(self, data, rows, cols, shape=None, type_as=None): data = self.from_numpy(data) @@ -3360,7 +3359,7 @@ def randperm(self, size, type_as=None): ) else: return tf.random.experimental.stateless_shuffle( - tf.range(size, dtype=type_as.dtype), seed=local_seed + tf.range(size), seed=local_seed ) def _convert_to_index_for_coo(self, tensor): diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 64268e4ce..70b5b9bda 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1021,7 +1021,7 @@ def solve_sample( "Debiasing is not implemented for TensorFlow backend." ) - if isinstance(debias, str) and debias.lower() == "split": + if isinstance(debias, str) and debias == "split": # split the samples into two halves with each half the mass X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_sample_ratio(X_a, a, nx=nx) X_b1, X_b2, b1, b2, sel_b1, sel_b2 = split_sample_ratio(X_b, b, nx=nx) diff --git a/test/test_solvers.py b/test/test_solvers.py index 39d96aba2..7152ae200 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -799,7 +799,8 @@ def test_solve_sample_methods(nx, method_params): @pytest.skip_backend("tf", reason="Not implemented for tf backend") @pytest.mark.parametrize("debias", [True, False, "split"]) -def test_solve_sample_debias(nx, debias): +@pytest.mark.parametrize("reg", [None, 10]) +def test_solve_sample_debias(nx, debias, reg): n_samples_s = 10 n_samples_t = 9 n_features = 2 @@ -812,16 +813,12 @@ def test_solve_sample_debias(nx, debias): xb, yb, ab, bb = nx.from_numpy(x, y, a, b) - sol = ot.solve_sample(x, y, reg=10, debias=debias) - solb = ot.solve_sample(xb, yb, ab, bb, reg=10, debias=debias) + sol = ot.solve_sample(x, y, reg=reg, debias=debias) + solb = ot.solve_sample(xb, yb, ab, bb, reg=reg, debias=debias) # check some attributes (no need ) assert_allclose_sol(sol, solb) - sol2 = ot.solve_sample(x, x, reg=10, debias=True) - - np.testing.assert_allclose(sol2.value, 0, atol=1e-10) - @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): From 657c5015de7263d42f4ac25145d3029ee770a8e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 13:26:36 +0200 Subject: [PATCH 12/28] add bsp wrapper --- ot/solvers/_bary.py | 4 +- ot/solvers/_linear.py | 191 +++++++++++++++++++++++++++++++++++++++--- test/test_solvers.py | 85 ++++++++++++------- 3 files changed, 238 insertions(+), 42 deletions(-) diff --git a/ot/solvers/_bary.py b/ot/solvers/_bary.py index 7ec7469d4..988a5dc37 100644 --- a/ot/solvers/_bary.py +++ b/ot/solvers/_bary.py @@ -12,7 +12,7 @@ from ..lp import free_support_barycenter_generic_costs from ..backend import get_backend -from ._linear import solve, solve_sample, lst_method_lazy +from ._linear import solve, solve_sample, lst_method_solve_sample def _bary_sample_bcd( @@ -434,7 +434,7 @@ def solve_bary_sample( """ - if method is not None and method.lower() in lst_method_lazy: + if method is not None and method.lower() in lst_method_solve_sample: raise NotImplementedError( f"method {method} operating on lazy tensors is not implemented yet" ) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 70b5b9bda..5ea48967a 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -8,7 +8,7 @@ # # License: MIT License -from ..utils import OTResult, dist, split_sample_ratio +from ..utils import OTResult, dist, split_sample_ratio, apply_scaler from ..lp import emd2, emd2_lazy, wasserstein_1d from ..backend import get_backend from ..unbalanced import mm_unbalanced, sinkhorn_knopp_unbalanced, lbfgsb_unbalanced @@ -18,6 +18,8 @@ empirical_sinkhorn_nystroem2, old_geomloss, ) +from ..sliced import sliced_wasserstein_distance, max_sliced_wasserstein_distance +from ..bsp import compute_bspot_bijection from ..smooth import smooth_ot_dual from ..gaussian import empirical_bures_wasserstein_distance from ..factored import factored_optimal_transport @@ -25,8 +27,13 @@ from ..optim import cg from warnings import warn +lst_valid_methods_solve = [ + "sinkhorn", # sinkhorn for unbalanced + "sinkhorn_log", # sinkhorn for unbalanced +] + -lst_method_lazy = [ +lst_method_solve_sample = [ "1d", "gaussian", "lowrank", @@ -37,8 +44,13 @@ "geomloss_tensorized", "geomloss_online", "geomloss_multiscale", + "sliced", + "max_sliced", + "bsp", ] +all_valid_methods_solve_sample = lst_valid_methods_solve + lst_method_solve_sample + def solve( M, @@ -294,6 +306,13 @@ def solve( if reg is None: reg = 0 + if method is not None and method not in lst_valid_methods_solve: + raise ValueError( + "Unknown method={} for solve, must be in {}".format( + method, lst_valid_methods_solve + ) + ) + # default values for solutions potentials = None value = None @@ -627,6 +646,9 @@ def solve_sample( grad="autodiff", random_state=None, debias=False, + n_projections=50, + projections=None, + scaler=None, ): r"""Solve the discrete optimal transport problem using the samples in the source and target domains. @@ -695,6 +717,11 @@ def solve_sample( - "nystroem" for Nystroem Sinkhorn solver - "factored" for factored OT solver - "geomloss" for GeomLoss Sinkhorn solver + - "sliced" for sliced wasserstein distance (see :any:`ot.sliced`) + - "max_sliced" for max sliced wasserstein distance (see + :any:`ot.max_sliced`) + - "bsp" for BSP-OT solver (only for distribution with same number of + samples and uniform weights) n_threads : int, optional Number of OMP threads for exact OT solver, by default 1 max_iter : int, optional @@ -728,6 +755,12 @@ def solve_sample( provided in the log dictionary. if debiased='split', then X_a and X_b are split into two halves and the debiased value is computed using the two halves as proposed in minibatch OT [91]. + n_projections : int, optional + Number of projections for sliced OT, by default 50 + projections : array-like, shape (n_projections, dim), optional + Projections for sliced OT, by default None (random projections are used) + scaler : callable, optional + Function to scale the input data, by default None (no scaling) Returns ------- @@ -834,6 +867,22 @@ def solve_sample( res = ot.solve_sample(xa, xb, a, b, reg=1.0, reg_type='L2') + - **Debiased OT and Sinkhorn divergence** (when ``debias=True``): + + One can compute the Sinkhorn divergence between two distributions using the + following code: + + .. code-block:: python + + div = ot.solve_sample(xa, xb, a, b, reg=1.0, debias=True).value + + Debiasing for all existing solvers is implemented and in particular one can + recover the minibatch exact OT debiasing [91] using the following code: + + .. code-block:: python + + div = ot.solve_sample(xa, xb, a, b, debias='split').value + - **Unbalanced OT [41]** (when ``unbalanced!=None``): .. math:: @@ -936,6 +985,38 @@ def solve_sample( # recover the squared Wasserstein distances W_dists = res.value + - **Sliced Wasserstein distance** (when ``method='sliced'``): + + This method computes the sliced Wasserstein distance between two + distributions. The sliced Wasserstein distance + is defined as the average of the Wasserstein distances between the projected + distributions on random directions. + + .. code-block:: python + + swd = ot.solve_sample(xa, xb, method='sliced', n_projections=50).value + + - **Max Sliced Wasserstein [34]** (when ``method='max_sliced'``): + + This method computes the max sliced Wasserstein distance between two + distributions. + + .. code-block:: python + + mswd = ot.solve_sample(xa, xb, method='max_sliced', n_projections=50).value + + - **Binary Space Partitioning OT (BSP)** (when ``method='bsp'``): + + This method computes the BSP-OT distance between two distributions and + alignments. The number of partitioning directions can be set with the `n_projections` parameter. + + .. code-block:: python + + res = ot.solve_sample(xa, xb, method='bsp', n_projections=50) + + bsp_ot_distance = res.value + bsp_plan = res.plan + .. _references-solve-sample: References @@ -989,6 +1070,16 @@ def solve_sample( """ + if method is not None and method.lower() not in all_valid_methods_solve_sample: + raise ValueError( + "Unknown method={} for solve_sample. Available methods are: {}".format( + method, all_valid_methods_solve_sample + ) + ) + + if scaler is not None: + X_a, Xb = apply_scaler(X_a, X_b, scaler=scaler) + if debias: dict_params = dict( metric=metric, @@ -1099,7 +1190,7 @@ def solve_sample( # return debiased result return res - if method is not None and method.lower() in lst_method_lazy: + if method is not None and method.lower() in all_valid_methods_solve_sample: lazy0 = lazy lazy = True @@ -1128,14 +1219,7 @@ def solve_sample( return res - elif ( - lazy - and method is None - and (reg is None or reg == 0) - and unbalanced is None - and X_a is not None - and X_b is not None - ): + elif lazy and method is None and (reg is None or reg == 0) and unbalanced is None: # Use lazy EMD solver with coordinates (no regularization, balanced) nx = get_backend(X_a, X_b, a, b) value_linear, log = emd2_lazy( @@ -1171,6 +1255,7 @@ def solve_sample( value_linear = None plan = None lazy_plan = None + sparse_plan = None status = None log = None @@ -1189,6 +1274,89 @@ def solve_sample( value = wasserstein_1d(X_a, X_b, a, b, p=p) value_linear = value + elif method == "sliced": # Sliced Wasserstein distance + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + value, log = sliced_wasserstein_distance( + X_a, + X_b, + a=a, + b=b, + p=p, + n_projections=n_projections, + projections=projections, + seed=random_state, + log=True, + ) + value_linear = value + + elif method == "max_sliced": # Max Sliced Wasserstein distance + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + value, log = max_sliced_wasserstein_distance( + X_a, + X_b, + a=a, + b=b, + p=p, + n_projections=n_projections, + projections=projections, + seed=random_state, + log=True, + ) + value_linear = value + + elif method == "bsp": # BSP-OT solver + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + if (X_a.shape[0] != X_b.shape[0]) or a is not None or b is not None: + raise ValueError( + "BSP-OT solver requires the same number of samples in both distributions and uniform weights (provided as None)." + ) + + if random_state is None: + random_state = 0 + + n = X_a.shape[0] + + value, perm, perms = compute_bspot_bijection( + X_a, X_b, p=p, seed=random_state, n_plans=n_projections + ) + value_linear = value + sparse_plan = nx.coo_matrix( + nx.ones(n, type_as=X_a) / n, + nx.arange(n), + perm, + shape=(n, n), + type_as=X_a, + ) + + log = { + "perm": perm, + "perms": perms, + } + elif method == "gaussian": # Gaussian Bures-Wasserstein if metric.lower() not in ["sqeuclidean"]: raise ( @@ -1396,6 +1564,7 @@ def solve_sample( value=value, lazy_plan=lazy_plan, value_linear=value_linear, + sparse_plan=sparse_plan, plan=plan, status=status, backend=nx, diff --git a/test/test_solvers.py b/test/test_solvers.py index 7152ae200..6f6afdad0 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -13,7 +13,7 @@ import ot from ot.bregman import geomloss from ot.backend import torch -from ot.solvers._linear import lst_method_lazy +from ot.solvers._linear import lst_method_solve_sample lst_reg = [None, 0.1] @@ -35,6 +35,8 @@ {"method": "factored", "rank": 2}, {"method": "lowrank", "rank": 2, "max_iter": 5}, {"method": "nystroem", "rank": 2}, + {"method": "sliced", "n_projections": 10}, + {"method": "max_sliced", "n_projections": 10}, ] lst_parameters_solve_sample_NotImplemented = [ @@ -64,7 +66,7 @@ ] lst_parameters_solve_bary_sample_NotImplemented = [ - {"method": method} for method in lst_method_lazy + {"method": method} for method in lst_method_solve_sample ] + [ {"lazy": True}, # fail lazy {"metric": "cosine"}, # fail on invalid metric @@ -735,33 +737,36 @@ def test_solve_sample_geomloss(nx, metric): sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=False, method="geomloss") assert_allclose_sol(sol0, sol) - sol1 = ot.solve_sample( - xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_tensorized" - ) - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) - - sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_online") - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) - - sol1 = ot.solve_sample( - xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_multiscale" - ) - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) + # commented because geomloss_tensorized and geomloss_online are not + # implemented in geomloss 0.2.0 yet + + # sol1 = ot.solve_sample( + # xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_tensorized" + # ) + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) + + # sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_online") + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) + + # sol1 = ot.solve_sample( + # xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_multiscale" + # ) + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss") np.testing.assert_allclose( @@ -820,6 +825,28 @@ def test_solve_sample_debias(nx, debias, reg): assert_allclose_sol(sol, solb) +def test_solve_sample_bsp(nx): + n_samples_s = 10 + n_samples_t = 10 # same number of samples for source and target + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + + xb, yb = nx.from_numpy(x, y) + + sol = ot.solve_sample(x, y, method="bsp") + solb = ot.solve_sample(xb, yb, method="bsp") + + # check some attributes (no need ) + assert_allclose_sol(sol, solb) + + with pytest.raises(ValueError): + # bsp method requires same number of samples for source and target + ot.solve_sample(x, y[:5], method="bsp") + + @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): n_samples_s = 20 From e1c8834a52c06e08850ea934c1a1ae1cb101fb1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 13:28:47 +0200 Subject: [PATCH 13/28] tets propertly methods --- test/test_solvers.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/test_solvers.py b/test/test_solvers.py index 6f6afdad0..46f17fc8a 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -847,6 +847,24 @@ def test_solve_sample_bsp(nx): ot.solve_sample(x, y[:5], method="bsp") +def test_solvers_bad_method(): + n_samples_s = 20 + n_samples_t = 7 + n_features = 2 + rng = np.random.RandomState(0) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + + C = ot.dist(x, y) + + with pytest.raises(ValueError): + ot.solve_sample(x, y, method="invalid_method") + + with pytest.raises(ValueError): + ot.solve(C, method="invalid_method") + + @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): n_samples_s = 20 From 902aeca788f077cfecfbf2ddd5160a68ec318d68 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 13:39:37 +0200 Subject: [PATCH 14/28] debug mehod value --- ot/solvers/_linear.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 5ea48967a..796b49769 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -30,6 +30,8 @@ lst_valid_methods_solve = [ "sinkhorn", # sinkhorn for unbalanced "sinkhorn_log", # sinkhorn for unbalanced + "mm", # unbalanced + "lbfgsb", # unbalanced ] From 8fee3aecae22ddc28912fcb684593004d0d48b42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 15:56:15 +0200 Subject: [PATCH 15/28] update doc and commit --- examples/plot_quickstart_guide.py | 60 +++++++++++++++++++++++++++++-- ot/solvers/_linear.py | 50 ++++++++++++++++++++++++-- test/test_solvers.py | 1 + test/test_utils.py | 1 + 4 files changed, 107 insertions(+), 5 deletions(-) diff --git a/examples/plot_quickstart_guide.py b/examples/plot_quickstart_guide.py index d5cfac8e6..2acf592bb 100644 --- a/examples/plot_quickstart_guide.py +++ b/examples/plot_quickstart_guide.py @@ -467,8 +467,8 @@ def df(G): # loss_fgw = ot.gromov.fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1) # loss_efgw = ot.gromov.entropic_fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1, epsilon=reg) # -# Large scale OT -# -------------- +# Large scale OT and approximations +# --------------------------------- # # We discuss here strategies to solve large scale OT problems using approximations # of the exact OT problem. @@ -583,6 +583,62 @@ def df(G): print(f"Exact OT loss = {loss:1.3f}") print(f"Bures-Wasserstein distance = {bw_value:1.3f}") +# %% +# One can also use the HD gaussian assumption (low rank covariance + diagonal) +# that has better properties in high dimension. The rank of the covariance +# matrices can be controlled with the :code:`rank` parameter. + +hdbw_value = ot.solve_sample(x1, x2, a, b, method="gaussian_hd", rank=1).value + +print(f"Bures-Wasserstein distance = {bw_value:1.3f}") +print(f"High Dimensional Bures-Wasserstein distance = {hdbw_value}") + +# %% +# Sliced Wasserstein +# ~~~~~~~~~~~~~~~~~~ +# +# The Sliced Wasserstein distance is a Wasserstein distance between +# empirical distributions that is computed by projecting the samples on random +# directions and averaging the Wasserstein distances between the projected +# distributions. It can be used as an approximation of the Wasserstein distance +# between empirical distributions. + +sw_value = ot.solve_sample(x1, x2, a, b, method="sliced", n_projections=10).value +max_sw_value = ot.solve_sample( + x1, x2, a, b, method="max_sliced", n_projections=10 +).value + +print(f"Exact OT loss = {loss:1.3f}") +print(f"Sliced Wasserstein distance = {sw_value:1.3f}") +print(f"Max Sliced Wasserstein distance = {max_sw_value:1.3f}") + +# %% +# Binary Space Partitioning (BSP) OT +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# One can also use the BSP OT approximation that is based on a recursive +# partitioning of the space and computes the Wasserstein distance between the +# empirical distributions by solving small OT problems between the samples in each +# partition. The number of partitions can be controlled with the +# code:`n_projections` parameter. + +# BSP can only find bijections so require same number of points +x1_bsp = np.concatenate([x1, x1], axis=0) + +sol_bsp = ot.solve_sample(x1_bsp, x2, method="bsp", n_projections=10) +bsp_value = sol_bsp.value +bsp_sparse_plan = sol_bsp.sparse_plan + +# sphinx_gallery_start_ignore + +pl.figure(1, (3, 3)) +plot2D_samples_mat(x1_bsp, x2, bsp_sparse_plan) +pl.plot(x1_bsp[:, 0], x1_bsp[:, 1], "ob", label="Source samples", **style) +pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style) +pl.title("BSP OT plan loss={:.3f}".format(bsp_value)) +pl.show() + +# sphinx_gallery_end_ignore # %% # .. note:: # The Gaussian Wasserstein problem can be solved with the classic API using the diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 796b49769..203347b57 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -21,7 +21,10 @@ from ..sliced import sliced_wasserstein_distance, max_sliced_wasserstein_distance from ..bsp import compute_bspot_bijection from ..smooth import smooth_ot_dual -from ..gaussian import empirical_bures_wasserstein_distance +from ..gaussian import ( + empirical_bures_wasserstein_distance, + empirical_bures_wasserstein_distance_hd, +) from ..factored import factored_optimal_transport from ..lowrank import lowrank_sinkhorn from ..optim import cg @@ -38,6 +41,7 @@ lst_method_solve_sample = [ "1d", "gaussian", + "gaussian_hd", "lowrank", "nystroem", "factored", @@ -715,6 +719,7 @@ def solve_sample( - "1d" for 1D OT solver done in parallel for each dimension. - "gaussian" for Gaussian OT solver that estimates mean and covariance and solve the closed form solution) + - "gaussian_hd" for high-dimensional Gaussian OT solver with rank given by `rank` parameter . - "lowrank" for low-rank Sinkhorn solver (see :any:`ot.lowrank`) - "nystroem" for Nystroem Sinkhorn solver - "factored" for factored OT solver @@ -952,7 +957,7 @@ def solve_sample( Corresponds to a low rank approximation of entropic OT (for a squared Euclidean cost) that runs in linear time. - - **Gaussian Bures-Wasserstein [2]** (when ``method='gaussian'``): + - **Gaussian Bures-Wasserstein** (when ``method='gaussian'``): This method computes the Gaussian Bures-Wasserstein distance between two Gaussian distributions estimated from the empirical distributions @@ -974,6 +979,14 @@ def solve_sample( # recover the squared Gaussian Bures-Wasserstein distance BW_dist = res.value + The Gaussian-HD variant of this method can be used for high-dimensional data with the following code: + + .. code-block:: python + + res = ot.solve_sample(xa, xb, method='gaussian_hd', rank=10) + + where the rank parameter controls the rank of the covariance matrices used in the computation of the Bures-Wasserstein distance. + - **Wasserstein 1d [1]** (when ``method='1D'``): This method computes the Wasserstein distance between two 1d distributions @@ -1368,8 +1381,39 @@ def solve_sample( if reg is None: reg = 1e-6 + if len(a.shape) == 1: + a = a.reshape(-1, 1) + if len(b.shape) == 1: + b = b.reshape(-1, 1) + value, log = empirical_bures_wasserstein_distance( - X_a, X_b, reg=reg, log=True + X_a, X_b, reg=reg, ws=a, wt=b, log=True + ) + value = value**2 # return the value (squared bures distance) + value_linear = value # return the value + + elif method == "gaussian_hd": # Gaussian BW for high-dimensional data + if metric.lower() not in ["sqeuclidean"]: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + if reg is None: + reg = 1e-5 + + if len(a.shape) == 1: + a = a.reshape(-1, 1) + if len(b.shape) == 1: + b = b.reshape(-1, 1) + + if rank > X_a.shape[1] or rank > X_b.shape[1]: + warn( + "Rank is larger than the number of features. Using full rank for Gaussian Bures-Wasserstein distance." + ) + rank = min(X_a.shape[1], X_b.shape[1]) + + value, log = empirical_bures_wasserstein_distance_hd( + X_a, X_b, d_intrinsic=rank, ws=a, wt=b, reg=reg, log=True ) value = value**2 # return the value (squared bures distance) value_linear = value # return the value diff --git a/test/test_solvers.py b/test/test_solvers.py index 46f17fc8a..d9bd5e5b2 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -32,6 +32,7 @@ {"method": "1d", "metric": "euclidean"}, {"method": "gaussian"}, {"method": "gaussian", "reg": 1}, + {"method": "gaussian_hd", "rank": 1}, {"method": "factored", "rank": 2}, {"method": "lowrank", "rank": 2, "max_iter": 5}, {"method": "nystroem", "rank": 2}, diff --git a/test/test_utils.py b/test/test_utils.py index d66bb6503..06ffe3133 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -959,6 +959,7 @@ def test_datascaler_backend_mismatch_raises(): scaler.transform(X) +@pytest.skip_backend("tf") @pytest.mark.parametrize("ratio", [0.3, 0.5, 0.51, 0.7]) @pytest.mark.parametrize("n", [10, 50, 100, 101]) @pytest.mark.parametrize("random", [True, False]) From b3ef5b683132a0d310c1f39bccc148dce5b7d41b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 16:03:25 +0200 Subject: [PATCH 16/28] fix function --- ot/solvers/_linear.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 203347b57..ce0fe0416 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1381,9 +1381,9 @@ def solve_sample( if reg is None: reg = 1e-6 - if len(a.shape) == 1: + if a is not None and len(a.shape) == 1: a = a.reshape(-1, 1) - if len(b.shape) == 1: + if b is not None and len(b.shape) == 1: b = b.reshape(-1, 1) value, log = empirical_bures_wasserstein_distance( @@ -1401,9 +1401,9 @@ def solve_sample( if reg is None: reg = 1e-5 - if len(a.shape) == 1: + if a is not None and len(a.shape) == 1: a = a.reshape(-1, 1) - if len(b.shape) == 1: + if b is not None and len(b.shape) == 1: b = b.reshape(-1, 1) if rank > X_a.shape[1] or rank > X_b.shape[1]: From 1a82b7446ab40a3a4f1f29d80047e238d6b80d2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 16:15:11 +0200 Subject: [PATCH 17/28] stabilize hd gaussians --- ot/gaussian.py | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/ot/gaussian.py b/ot/gaussian.py index 59d368afa..cfe944834 100644 --- a/ot/gaussian.py +++ b/ot/gaussian.py @@ -171,15 +171,6 @@ def bures_wasserstein_mapping_hd(ms, mt, Us, Ut, ls, lt, sigma2_s, sigma2_t, log """ nx = get_backend(ms, mt, Us, Ut, ls, lt, sigma2_s, sigma2_t) - print(ms.dtype) - print(mt.dtype) - print(Us.dtype) - print(Ut.dtype) - print(ls.dtype) - print(lt.dtype) - print(sigma2_s.dtype) - print(sigma2_t.dtype) - p = Us.shape[0] # source @@ -429,8 +420,8 @@ def empirical_bures_wasserstein_mapping_hd( Ut = Qt[:, -dt:] lt = a_t - sgm2_t - sgm2_s = nx.unsqueeze(sgm2_s, 0) - sgm2_t = nx.unsqueeze(sgm2_t, 0) + sgm2_s = nx.maximum(nx.unsqueeze(sgm2_s, 0), reg) + sgm2_t = nx.maximum(nx.unsqueeze(sgm2_t, 0), reg) if log: A, b, log = bures_wasserstein_mapping_hd( @@ -901,8 +892,8 @@ def empirical_bures_wasserstein_distance_hd( Ut = Qt[:, -dt:] lt = a_t - sgm2_t - sgm2_s = nx.unsqueeze(sgm2_s, 0) - sgm2_t = nx.unsqueeze(sgm2_t, 0) + sgm2_s = nx.maximum(nx.unsqueeze(sgm2_s, 0), reg) + sgm2_t = nx.maximum(nx.unsqueeze(sgm2_t, 0), reg) if log: W, log = bures_wasserstein_distance_hd( From 13bd0e7098ab341ec50dfa7159a2f7ff48b1f253 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 16:39:00 +0200 Subject: [PATCH 18/28] stabilise really HD BW --- ot/gaussian.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ot/gaussian.py b/ot/gaussian.py index cfe944834..103d102d8 100644 --- a/ot/gaussian.py +++ b/ot/gaussian.py @@ -880,14 +880,14 @@ def empirical_bures_wasserstein_distance_hd( eigs = nx.eigh(Cs) a_s = eigs[0][-ds:] - sgm2_s = (nx.trace(Cs) - nx.sum(a_s)) / (p - ds) + sgm2_s = (nx.trace(Cs) - nx.sum(a_s)) / (p - ds + reg) Qs = eigs[1] Us = Qs[:, -ds:] ls = a_s - sgm2_s eigt = nx.eigh(Ct) a_t = eigt[0][-dt:] - sgm2_t = (nx.trace(Ct) - nx.sum(a_t)) / (p - dt) + sgm2_t = (nx.trace(Ct) - nx.sum(a_t)) / (p - dt + reg) Qt = eigt[1] Ut = Qt[:, -dt:] lt = a_t - sgm2_t From f23973d66efcc0092aaa07bc6355da3d524ea054 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 16:40:02 +0200 Subject: [PATCH 19/28] cleanup tests --- test/test_solvers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/test/test_solvers.py b/test/test_solvers.py index d9bd5e5b2..d27418d86 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -799,7 +799,12 @@ def test_solve_sample_methods(nx, method_params): assert_allclose_sol(sol, solb) sol2 = ot.solve_sample(x, x, **method_params) - if method_params["method"] not in ["factored", "lowrank", "nystroem"]: + if method_params["method"] not in [ + "factored", + "lowrank", + "nystroem", + "gaussian_hd", + ]: np.testing.assert_allclose(sol2.value, 0, atol=1e-10) From 3caf068b53c32b9330173a6c43ddb2f8366325bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 16:58:31 +0200 Subject: [PATCH 20/28] remove non exiting mehod --- examples/plot_debias_sink_div.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py index 382b8f568..d743f4298 100644 --- a/examples/plot_debias_sink_div.py +++ b/examples/plot_debias_sink_div.py @@ -225,7 +225,6 @@ def sample_two_balls(n, radius=1.0, sep=1): sink_list = [] sink_div_list = [] ot_mb_list = [] -ot_mb_rand = [] ot_mb_sink_list = [] for sep in sep_list: x2sep = sample_two_balls(n, radius=1.0, sep=sep) @@ -238,9 +237,7 @@ def sample_two_balls(n, radius=1.0, sep=1): ) sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value) ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value) - ot_mb_rand.append( - ot.solve_sample(x1, x2sep, debias="split_random", random_state=41).value - ) + ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value) pl.figure(3) @@ -248,7 +245,6 @@ def sample_two_balls(n, radius=1.0, sep=1): pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2") pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3") -pl.plot(sep_list, ot_mb_rand, label="Debiased MB OT (random)", color="C4") pl.xlabel("Separation between distributions") pl.ylabel("Loss/Divergence") pl.title("Comparison of biased VS debiased OT losses") From b11b0b1413f0ea8dff26f78822a4dfade7d9b679 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Fri, 3 Jul 2026 17:33:10 +0200 Subject: [PATCH 21/28] fixx scaling split mb --- ot/solvers/_linear.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index ce0fe0416..035814ba8 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1139,12 +1139,8 @@ def solve_sample( resab2 = solve_sample(X_a2, X_b2, a=a2, b=b2, **dict_params) # compute debiased values - value = 0.5 * (resab1.value + resab2.value) - 0.5 * ( - resaa.value + resbb.value - ) - value_linear = 0.5 * (resab1.value_linear + resab2.value_linear) - 0.5 * ( - resaa.value_linear + resbb.value_linear - ) + value = (resab1.value + resab2.value) - (resaa.value + resbb.value) + value_linear = value log = { "res_aa": resaa, From 921e936511dbe8b620a0fd4c479354acb10cbbb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 09:19:27 +0200 Subject: [PATCH 22/28] cleanup doc --- examples/plot_quickstart_guide.py | 10 +++++----- ot/solvers/_linear.py | 11 ++++++----- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/examples/plot_quickstart_guide.py b/examples/plot_quickstart_guide.py index 2acf592bb..6ba3d410a 100644 --- a/examples/plot_quickstart_guide.py +++ b/examples/plot_quickstart_guide.py @@ -594,6 +594,10 @@ def df(G): print(f"High Dimensional Bures-Wasserstein distance = {hdbw_value}") # %% +# .. note:: +# The Gaussian Wasserstein problem can be solved with the classic API using the +# :func:`ot.gaussian.empirical_bures_wasserstein_distance` function. +# # Sliced Wasserstein # ~~~~~~~~~~~~~~~~~~ # @@ -620,7 +624,7 @@ def df(G): # partitioning of the space and computes the Wasserstein distance between the # empirical distributions by solving small OT problems between the samples in each # partition. The number of partitions can be controlled with the -# code:`n_projections` parameter. +# `n_projections` parameter. # BSP can only find bijections so require same number of points x1_bsp = np.concatenate([x1, x1], axis=0) @@ -640,10 +644,6 @@ def df(G): # sphinx_gallery_end_ignore # %% -# .. note:: -# The Gaussian Wasserstein problem can be solved with the classic API using the -# :func:`ot.gaussian.empirical_bures_wasserstein_distance` function. -# # Comparing all OT plans # ---------------------- # diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 035814ba8..5ea6cdcec 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -691,17 +691,17 @@ def solve_sample( c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. - If :math:`\texttt{reg_type}=`'entropy', then :math:`\mathbf{c} = 1_{dim_a} 1_{dim_b}^T`. + If `reg_type='entropy'`, then :math:`\mathbf{c} = 1_{dim_a} 1_{dim_b}^T`. reg_type : str, optional Type of regularization :math:`R` either "KL", "L2", "entropy", by default "KL" unbalanced : float or indexable object of length 1 or 2 Marginal relaxation term. If it is a scalar or an indexable object of length 1, then the same relaxation is applied to both marginal relaxations. - The balanced OT can be recovered using :math:`unbalanced=float("inf")`. + The balanced OT can be recovered using `unbalanced=float("inf")`. For semi-relaxed case, use either - :math:`unbalanced=(float("inf"), scalar)` or - :math:`unbalanced=(scalar, float("inf"))`. + `unbalanced=(float("inf"), scalar)` or + `unbalanced=(scalar, float("inf"))`. If unbalanced is an array, it must have the same backend as input arrays `(a, b, M)`. unbalanced_type : str, optional @@ -716,6 +716,7 @@ def solve_sample( Method for solving the problem, this can be used to select the solver for unbalanced problems (see :any:`ot.solve`), or to select a specific large scale solver. Available methods are: + - "1d" for 1D OT solver done in parallel for each dimension. - "gaussian" for Gaussian OT solver that estimates mean and covariance and solve the closed form solution) @@ -726,7 +727,7 @@ def solve_sample( - "geomloss" for GeomLoss Sinkhorn solver - "sliced" for sliced wasserstein distance (see :any:`ot.sliced`) - "max_sliced" for max sliced wasserstein distance (see - :any:`ot.max_sliced`) + :any:`ot.sliced`) - "bsp" for BSP-OT solver (only for distribution with same number of samples and uniform weights) n_threads : int, optional From d9419f91037c4e3cddcee251421ac21652569e80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 09:20:48 +0200 Subject: [PATCH 23/28] debug refs --- ot/solvers/_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 5ea6cdcec..3412423fa 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -1055,7 +1055,7 @@ def solve_sample( Optimal Transport. Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics (AISTATS). - ..[23] Aude, G., Peyré, G., Cuturi, M., Learning Generative Models with + .. [23] Aude, G., Peyré, G., Cuturi, M., Learning Generative Models with Sinkhorn Divergences, Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, (AISTATS) 21, 2018 From 7bae4a396e6f280b4635f4edfda0ad0133648563 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 11:48:33 +0200 Subject: [PATCH 24/28] better coverage --- ot/solvers/_linear.py | 6 +++--- test/test_solvers.py | 29 +++++++++++++++++++++++++++++ test/test_utils.py | 4 ++-- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 3412423fa..5de1df7a6 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -490,13 +490,13 @@ def solve( else: # stabilize kl of 0 mass if nx.any(a == 0) or nx.any(b == 0): - a = a + 1e-16 * nx.min(a) - b = b + 1e-16 * nx.min(b) + a = a + 1e-15 * nx.min(a) + b = b + 1e-15 * nx.min(b) print( "Warning: a or b has 0 mass, adding small mass to avoid NaN in KL divergence." ) value = value_linear + reg * nx.kl_div( - plan, a[:, None] * b[None, :] + plan, a[:, None] * b[None, :] + 1e-15 ) if grad == "envelope": # set the gradient at convergence diff --git a/test/test_solvers.py b/test/test_solvers.py index d27418d86..84a11c866 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -33,11 +33,14 @@ {"method": "gaussian"}, {"method": "gaussian", "reg": 1}, {"method": "gaussian_hd", "rank": 1}, + {"method": "gaussian_hd", "rank": 3}, {"method": "factored", "rank": 2}, {"method": "lowrank", "rank": 2, "max_iter": 5}, {"method": "nystroem", "rank": 2}, {"method": "sliced", "n_projections": 10}, + {"method": "sliced", "n_projections": 10, "metric": "euclidean"}, {"method": "max_sliced", "n_projections": 10}, + {"method": "max_sliced", "n_projections": 10, "metric": "euclidean"}, ] lst_parameters_solve_sample_NotImplemented = [ @@ -808,6 +811,32 @@ def test_solve_sample_methods(nx, method_params): np.testing.assert_allclose(sol2.value, 0, atol=1e-10) +def test_zero_mass_solvers(): + # test that solvers handle zero mass distributions correctly + n_samples_s = 10 + n_samples_t = 9 + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + # Set one of the distributions to zero mass + a[0] = 0.0 + b[0] = 0.0 + + a /= a.sum() + b /= b.sum() + + res = ot.solve_sample(x, y, a, b) + res_reg = ot.solve_sample(x, y, a, b, reg=1.0) + + assert np.isnan(res.value) == False + assert np.isnan(res_reg.value) == False + + @pytest.skip_backend("tf", reason="Not implemented for tf backend") @pytest.mark.parametrize("debias", [True, False, "split"]) @pytest.mark.parametrize("reg", [None, 10]) diff --git a/test/test_utils.py b/test/test_utils.py index 06ffe3133..b2d6b4508 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -975,7 +975,7 @@ def test_split_sample_ratio(nx, ratio, n, random): # test uniform weights X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( - X, ratio=ratio, random_state=seed, nx=nx + X, ratio=ratio, random_split=random, random_state=seed, nx=nx ) np.testing.assert_allclose(nx.to_numpy(a1).sum(), ratio, atol=1e-5) @@ -986,7 +986,7 @@ def test_split_sample_ratio(nx, ratio, n, random): # test non uniform weights X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( - X, ratio=ratio, a=a, random_state=seed + X, ratio=ratio, a=a, random_split=random, random_state=seed ) with pytest.raises(ValueError): From 00905ccf774742514192cae92e5b5f382bf10e6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 13:04:46 +0200 Subject: [PATCH 25/28] debut randperm jax --- ot/backend.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ot/backend.py b/ot/backend.py index d210576b7..7749c948d 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -1867,9 +1867,7 @@ def randperm(self, size, type_as=None): if not isinstance(size, int): raise ValueError("size must be an integer") if type_as is not None: - return jnp.asarray( - jax.random.permutation(subkey, size), dtype=type_as.dtype - ) + return jnp.asarray(jax.random.permutation(subkey, size)) else: return jax.random.permutation(subkey, size) From 03b7ed8897e9888009bf54406e16b04f7c039d83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 13:18:37 +0200 Subject: [PATCH 26/28] add all tests --- test/test_solvers.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/test/test_solvers.py b/test/test_solvers.py index 84a11c866..5c676e140 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -14,6 +14,7 @@ from ot.bregman import geomloss from ot.backend import torch from ot.solvers._linear import lst_method_solve_sample +from ot.utils import DataScaler lst_reg = [None, 0.1] @@ -49,6 +50,10 @@ "method": "gaussian", "metric": "euclidean", }, # fail gaussian on metric not euclidean + { + "method": "gaussian_hd", + "metric": "euclidean", + }, # fail gaussian on metric not euclidean { "method": "factored", "metric": "euclidean", @@ -67,6 +72,14 @@ "reg": 1, "unbalanced": 1, }, # fail lazy for unbalanced and regularized + { + "method": "sliced", + "metric": "wrong_metric", + }, # fail sliced on metric not euclidean + { + "method": "max_sliced", + "metric": "wrong_metric", + }, # fail sliced on metric not euclidean ] lst_parameters_solve_bary_sample_NotImplemented = [ @@ -881,6 +894,10 @@ def test_solve_sample_bsp(nx): # bsp method requires same number of samples for source and target ot.solve_sample(x, y[:5], method="bsp") + with pytest.raises(ValueError): + # bsp method requires same number of samples for source and target + ot.solve_sample(x, y, method="bsp", metric="wrong_metric") + def test_solvers_bad_method(): n_samples_s = 20 @@ -900,6 +917,32 @@ def test_solvers_bad_method(): ot.solve(C, method="invalid_method") +@pytest.mark.parametrize("norm", ["standard", "minmax", "l2"]) +def test_solve_sample_scaler(nx, norm): + n_samples_s = 20 + n_samples_t = 7 + n_features = 2 + rng = np.random.RandomState(0) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + xb, yb, ab, bb = nx.from_numpy(x, y, a, b) + + scaler0 = DataScaler(norm=norm) + scaler0.fit(x) + + scaler1 = DataScaler(norm=norm) + scaler1.fit(xb) + + sol0 = ot.solve_sample(x, y, a, b, scaler=scaler0) + sol1 = ot.solve_sample(xb, yb, ab, bb, scaler=scaler1) + + assert_allclose_sol(sol0, sol1) + + @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): n_samples_s = 20 From 048b186bf80efd8c7d03d272789b635bbe060072 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 13:26:24 +0200 Subject: [PATCH 27/28] betste test bsp --- test/test_solvers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_solvers.py b/test/test_solvers.py index 5c676e140..6cc8137cc 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -894,7 +894,7 @@ def test_solve_sample_bsp(nx): # bsp method requires same number of samples for source and target ot.solve_sample(x, y[:5], method="bsp") - with pytest.raises(ValueError): + with pytest.raises(NotImplementedError): # bsp method requires same number of samples for source and target ot.solve_sample(x, y, method="bsp", metric="wrong_metric") From 6d82205b9968f04a84b74413ffeae2c8527379f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 14:52:08 +0200 Subject: [PATCH 28/28] fix typos --- examples/plot_debias_sink_div.py | 4 ++-- ot/solvers/_linear.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py index d743f4298..d6f44bafd 100644 --- a/examples/plot_debias_sink_div.py +++ b/examples/plot_debias_sink_div.py @@ -5,7 +5,7 @@ ====================================== This example shows how to use the debiased OT solvers in `ot.solve_sample` to -compute Sinkhorn divergences and debiased Minibtach solutions. The debiased OT solvers +compute Sinkhorn divergences and debiased Minibatch solutions. The debiased OT solvers can be used with balanced and unbalanced OT problems, and with different regularization types (entropic, L2, group lasso). """ @@ -114,7 +114,7 @@ def sample_two_balls(n, radius=1.0, sep=1): # --------------------------------- # # Doing OT on minibatches leads to a similar bias than using entropic -# regularization since the average OT plan is densified to to the stochasticity +# regularization since the average OT plan is densified due to the stochasticity # of the minibatch sampling. On a given minibatch, the debiased loss can be # computed by setting the `debias` parameter to `'split'`that split the data # points in each distributions in two and computes the debias OT loss as: diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 5de1df7a6..ad27c6e91 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -768,7 +768,8 @@ def solve_sample( projections : array-like, shape (n_projections, dim), optional Projections for sliced OT, by default None (random projections are used) scaler : callable, optional - Function to scale the input data, by default None (no scaling) + Function to scale the input data, following :any:`ot.utils.apply_scaler`. + by default None (no scaling). Returns -------