[MRG] Proximal point solver for batched optimal transport (#832)

* correction SGOT cost matrix, added to contributor list, move sgot example to other

* updated graphs

* fix plots

* update example

* fix releases.md

* reformating

* updating format

* updated format

* updated format

* added PR and references

* aligned usage of solve and solve_batch + added batch tests

* fixed example

* Apply suggestions from code review

Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>

---------

Co-authored-by: Sienna O'Shea <osheasienna@gmail.com>
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>

Author: 3 people

README.md

@@ -471,4 +471,8 @@ Artificial Intelligence.
 
 \[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.
+\[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.
+
+\[92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). 
+[A fast proximal point method for computing exact wasserstein distance.](https://proceedings.mlr.press/v115/xie20b/xie20b.pdf) In Uncertainty in artificial intelligence (pp. 433-453). PMLR.
+

RELEASES.md

@@ -30,14 +30,15 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s
 - Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723)
 - Add cost functions between linear operators following  
   [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920),  
-  implemented in `ot.sgot` (PR #792)
+  implemented in `ot.sgot` (PR #792, PR #830)
 - Add batch FUGW loss to `ot.batch` and fix issues in some default parameters in the batch module (PR #775)
 - Wrapper for barycenter solvers with free support `ot.solvers.bary_free_support` (PR #730)
 - Build wheels on ubuntu ARM to avoid QEMU emulation (PR #818)
 - Add new methods to compute the linear transport map and the related 2-Wasserstein distance betweeen high-dimensional (HD) Gaussian distributions as described in [88], implemented in  `ot.gaussian.bures_wasserstein_mapping_hd` and `ot.gaussian.bures_wasserstein_distance_hd`, respectively. Two additional methods estimate the same quantities from the source and destination observed data and are implemented in `ot.gaussian.empirical_bures_wasserstein_mapping_hd` and `ot.gaussian.empirical_bures_wasserstein_distance_hd`, respectively (PR #814)
 - 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)
+- Implemented batch proximal point solver for OT problems `ot.batch.proximal_bregman_log_plan_batch` function and updated wrapper functions `ot.batch.solve_batch` and `ot.batch.solve_sample_batch` (PR #832)
 - Implement debiased OT solvers in `ot.solve_sample`.
 
 

examples/backends/plot_ot_batch.py

@@ -13,6 +13,7 @@
 """
 
 # Author: Paul Krzakala <paul.krzakala@gmail.com>
+#         Thibaut Germain <thibaut.germain.pro@gmail.com>
 # License: MIT License
 
 # sphinx_gallery_thumbnail_number = 1
@@ -75,25 +76,52 @@
 #    This is simple but inefficient for large batches.
 #
 # Instead, you can use :func:`ot.batch.solve_batch`, which solves all
-# problems in parallel.
+# problems in parallel. Several methods are available: ["sinkhorn", "log_sinkhorn"]
+# which solve regularized OT problems, and ["proximal"] which
+# solves regularized and unregularized OT problem using a proximal point scheme.
+# By default, the method is set to "auto" which automatically selects the appropriate
+# method based on the value of `reg`. If `reg` is None or 0, the proximal point method
+# is used. If `reg` is greater than 0, the Sinkhorn algorithm is used.
+
+max_iter = 10000
+tol = 1e-4
+
+# Classical OT problem
+## Naive approach
+results_values_list = []
+for i in range(n_problems):
+    res = ot.solve(M_list[i], max_iter=max_iter, tol=tol)
+    results_values_list.append(res.value_linear)
 
-reg = 1.0
-max_iter = 100
-tol = 1e-3
+## Batched approach
+results_batch = ot.solve_batch(M=M_batch, max_iter=max_iter, tol=tol)
+results_values_batch = results_batch.value_linear
 
-# Naive approach
+exact_validated = np.allclose(
+    np.array(results_values_list), results_values_batch, atol=tol * 10
+)
+
+# Entropic regularized OT problem
+## Naive approach
+reg = 1.0
 results_values_list = []
 for i in range(n_problems):
     res = ot.solve(M_list[i], reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy")
     results_values_list.append(res.value_linear)
 
-# Batched approach
+## Batched approach
 results_batch = ot.solve_batch(
     M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy"
 )
 results_values_batch = results_batch.value_linear
 
-assert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10)
+entropic_validated = np.allclose(
+    np.array(results_values_list), results_values_batch, atol=tol * 10
+)
+
+print(
+    f"Exact solve vs proximal batch close: {exact_validated} \nSinkhorn solve vs Sinkhorn solve_batch close: {entropic_validated}"
+)
 
 #############################################################################
 #

ot/batch/__init__.py

@@ -5,6 +5,7 @@
 
 # Author: Remi Flamary <remi.flamary@unice.fr>
 #         Paul Krzakala <paul.krzakala@gmail.com>
+#         Thibaut Germain <thibaut.germain.pro@gmail.com>
 #
 # License: MIT License
 
@@ -25,6 +26,7 @@
     bregman_log_projection_batch,
     bregman_projection_batch,
     entropy_batch,
+    proximal_bregman_log_plan_batch,
 )
 
 __all__ = [
@@ -40,4 +42,5 @@
     "loss_quadratic_batch",
     "loss_quadratic_samples_batch",
     "tensor_batch",
+    "proximal_bregman_log_plan_batch",
 ]