diff --git a/CITATION.cff b/CITATION.cff index 513398eb7..d7e4c56ab 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -109,3 +109,4 @@ keywords: - gromov-wasserstein license: MIT version: 0.9.7 +date-released: 2026-07-10 diff --git a/README.md b/README.md index 950a509e9..70092d9ad 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,7 @@ using the following references from the current version and from our [JMLR paper](https://jmlr.org/papers/v22/20-451.html): ``` -Flamary R., Vincent-Cuaz C., Courty N., Gramfort A., Kachaiev O., Quang Tran H., David L., Bonet C., Cassereau N., Gnassounou T., Tanguy E., Delon J., Collas A., Mazelet S., Chapel L., Kerdoncuff T., Yu X., Feickert M., Krzakala P., Liu T., Fernandes Montesuma E. POT Python Optimal Transport (version 0.9.5). URL: https://github.com/PythonOT/POT +Flamary R., Vincent-Cuaz C., Courty N., Gramfort A., Kachaiev O., Quang Tran H., David L., Bonet C., Cassereau N., Gnassounou T., Tanguy E., Delon J., Collas A., Mazelet S., Chapel L., Kerdoncuff T., Yu X., Feickert M., Krzakala P., Liu T., Fernandes Montesuma E., Neike N., Genest B., Coeurjolly D., Germain T., O'Shea S., Corneli M., Genans F. (2026). POT Python Optimal Transport (version 0.9.7). DOI: 10.5281/zenodo.17161062 URL: https://github.com/PythonOT/POT Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, Titouan Vayer, POT Python Optimal Transport library, Journal of Machine Learning Research, 22(78):1−8, 2021. URL: https://pythonot.github.io/ ``` @@ -104,13 +104,17 @@ Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Bo In Bibtex format: ```bibtex -@misc{flamary2024pot, - author = {Flamary, R{\'e}mi and Vincent-Cuaz, C{\'e}dric and Courty, Nicolas and Gramfort, Alexandre and Kachaiev, Oleksii and Quang Tran, Huy and David, Laurène and Bonet, Cl{\'e}ment and Cassereau, Nathan and Gnassounou, Th{\'e}o and Tanguy, Eloi and Delon, Julie and Collas, Antoine and Mazelet, Sonia and Chapel, Laetitia and Kerdoncuff, Tanguy and Yu, Xizheng and Feickert, Matthew and Krzakala, Paul and Liu, Tianlin and Fernandes Montesuma, Eduardo}, - title = {POT Python Optimal Transport (version 0.9.5)}, - url = {https://github.com/PythonOT/POT}, - year = {2024} +@software{flamary2026pot, +author = {Flamary, Rémi and Vincent-Cuaz, Cédric and Courty, Nicolas and Gramfort, Alexandre and Kachaiev, Oleksii and Quang Tran, Huy and David, Laurène and Bonet, Clément and Cassereau, Nathan and Gnassounou, Théo and Tanguy, Eloi and Delon, Julie and Collas, Antoine and Mazelet, Sonia and Chapel, Laetitia and Kerdoncuff, Tanguy and Yu, Xizheng and Feickert, Matthew and Krzakala, Paul and Liu, Tianlin and Fernandes Montesuma, Eduardo and Neike, Nathan and Genest, Baptiste and Coeurjolly, David and Germain, Thibaut and O'Shea, Sienna and Corneli, Marco and Genans, Ferdinand}, +doi = {10.5281/zenodo.17161062}, +month = {7}, +title = {POT Python Optimal Transport}, +version = {0.9.7}, +url = {https://github.com/PythonOT/POT}, +year = {2026} } + @article{flamary2021pot, author = {R{\'e}mi Flamary and Nicolas Courty and Alexandre Gramfort and Mokhtar Z. Alaya and Aur{\'e}lie Boisbunon and Stanislas Chambon and Laetitia Chapel and Adrien Corenflos and Kilian Fatras and Nemo Fournier and L{\'e}o Gautheron and Nathalie T.H. Gayraud and Hicham Janati and Alain Rakotomamonjy and Ievgen Redko and Antoine Rolet and Antony Schutz and Vivien Seguy and Danica J. Sutherland and Romain Tavenard and Alexander Tong and Titouan Vayer}, title = {POT: Python Optimal Transport}, @@ -467,8 +471,12 @@ Artificial Intelligence. \[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](https://www.cs.columbia.edu/~blei/seminar/2020-representation/readings/TippingBishop1999.pdf). 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. -\[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 +\[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. + diff --git a/RELEASES.md b/RELEASES.md index dad6fa306..3e6014315 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,11 +2,43 @@ ## 0.9.7 -This new release adds support for sparse cost matrices and a new lazy exact OT solver that re-computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m). Both implementations are backend-agnostic and preserve gradient computation for automatic differentiation. +This release contains several bug fixes and performance improvements, as well as updates to the documentation and examples and tests. We provide below a summary of the main new features and closed issues. + + +**New exact OT solvers.** This new release adds many new variants and updates for the exact OT solver. First a new lazy exact OT solver that re-computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m) ( available in [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) with `lazy=True)`). Another major feature is the addition of a warmstart feature for the EMD solver, allowing users to provide initial potentials to speed up convergence (with `init_potentials` in [`ot.solve`](https://pythonot.github.io/all.html#ot.solve) and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample)). Finally the release also include a sparse solver and [`ot.solve`](https://pythonot.github.io/all.html#ot.solve) now accept sparse cost matrices and provides an OT plan whose support is included in the support of the cost matrix. All implementations are backend-agnostic and preserve gradient computation for automatic differentiation (but are solved on CPU). + +The computational time for different solvers (on a laptop CPU) are shown below: + +| n| 100| 500| 1000| 5000| +|--------------------------|:------------------------:|:------------------------:|:------------------------:|:------------------------:| +| Solve | 7.259e-04| 1.501e-02| 1.055e-01| 2.343e+00| +| Solve Warm Start | 6.970e-04| 5.413e-03| 2.291e-02| 6.036e-01| +| Solve Lazy | 1.280e-03| 3.098e-02| 1.421e-01| 3.377e+00| +| Solve Sparse 10% | 3.453e-04| 6.716e-04| 1.535e-03| 3.534e-02| + + +**BSPOT solver.** A new solver relying on [Binary Space Partitioning (BSP)](https://pythonot.github.io/master/auto_examples/plot_bsp_ot.html) has been added to compute sparse transport plans between discrete measures in loglinear time which allows for very large problems to be solved. The release also includes a new [`ot.unbalanced.uot_1d`](https://pythonot.github.io/master/gen_modules/ot.unbalanced.html#ot.unbalanced.uot_1d) solver with a Frank-Wolfe solver for unbalanced optimal transport in 1D. + +**Sliced OT pans** Finally we now have a sliced OT plan solver that can be used to compute sliced transport plans ([min-pivot sliced](https://pythonot.github.io/master/gen_modules/ot.sliced.html#ot.sliced.min_sliced_transport_plan) and [expected sliced](https://pythonot.github.io/master/gen_modules/ot.sliced.html#ot.sliced.expected_sliced_plan)) between two measures. + +**OT between dynamical systems.** A novel [Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://pythonot.github.io/master/auto_examples/others/plot_sgot.html) has been implemented in `ot.sgot`. + +**Unified API for barycenter solvers in `ot.solve_bary_sample`.** A new free support solver for barycenter solvers has been added in [`ot.solve_bary_sample`](https://pythonot.github.io/master/all.html#ot.solve_bary_sample). You can see [examples of how to use it here](https://pythonot.github.io/master/auto_examples/barycenters/plot_solve_barycenter_variants.html). + +**OT between high dimensional Gaussian distributions.** New methods to compute the linear transport map and the related 2-Wasserstein distance between high-dimensional (HD) Gaussian distributions have been added in [`ot.gaussian.bures_wasserstein_mapping_hd`](https://pythonot.github.io/master/gen_modules/ot.gaussian.html#id56) and [`ot.gaussian.bures_wasserstein_distance_hd`](https://pythonot.github.io/master/gen_modules/ot.gaussian.html#ot.gaussian.bures_wasserstein_distance_hd), respectively with empirical versions that estimate the distance from sample. + +**Batched solver for exact OT.** The batch implementations have also been improved and you can now solve exact OT problems in parallel on CPU or GPU using the new proximal point solver in functions [`ot.solve_batch`](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.solve_batch) and [`ot.solve_sample_batch`](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.solve_sample_batch) when `reg=0` or no provided. A new batch [loss for Fused unbalanced Gromov-Wasserstein](https://pythonot.github.io/master/gen_modules/ot.batch.html#ot.batch.loss_quadratic_batch) is also now available in `ot.batch`. + + +**New methods in unified API [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample).** The unified API function +[`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) has also been updated to allows solving of specific problems such as BSP-OT, distance between high-dimensional (HD) Gaussian distributions and sliced and max-sliced distances. + + +**Data normalization for sliced and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample) solvers.** The release also includes tools for data normalization and scaling which can not be used in sliced Wasserstein distance computations. A simple normalization class [`ot.utils.DataScaler`](https://pythonot.github.io/master/gen_modules/ot.utils.html#id46) has been added and supports all backends for `'standard'`, `'minmax'`, and `'l2'` methods. Finally an optional `scaler` parameter has been added to [`ot.sliced_wasserstein_distance`](https://pythonot.github.io/master/all.html#ot.sliced_wasserstein_distance), [`ot.max_sliced_wasserstein_distance`](https://pythonot.github.io/master/all.html#ot.max_sliced_wasserstein_distance) and [`ot.solve_sample`](https://pythonot.github.io/all.html#ot.solve_sample). + #### New features -- Fix reference number error introduced in PR #767 (PR #819) - Refactor lazy EMD network simplex storage to avoid dense per-arc cost, endpoint, flow, and state storage where possible, and return sparse lazy transport plans instead of materializing dense plans internally (PR #813) @@ -28,9 +60,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Add optional `scaler` parameter to `sliced_wasserstein_distance` and `max_sliced_wasserstein_distance` (PR #808) - Add SGD based semi-discrete OT solver in `ot.semidiscrete` and a gallery example. (PR #812) - 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) +- 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, 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) @@ -38,11 +68,13 @@ 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) +- 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`. #### Closed issues +- Fix label-aware cost correction in `ot.da` transport classes: the large cost was applied to unlabeled pairs instead of labeled pairs with different labels, so `ys`/`yt` had no effect when all labels were known (PR #833, Issue #664) - Mitigate NaN regime of `entropic_partial_wasserstein` at small `reg` via a new log-domain solver, reachable with `entropic_partial_wasserstein(..., method='sinkhorn_log')` (Issue #723; the default `method='sinkhorn'` path is unchanged — callers opt into the log-domain variant) - Fix NumPy 2.x compatibility in Brenier potential bounds (PR #788) - Fix MSVC Windows build by removing **restrict** keyword (PR #788) @@ -60,6 +92,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Fix entropic regularization in `gcg`(PR #817, Issue #758) - Fix documentation build on master with submodules (PR #818) - Fix failing test for unbalanced solver with generic regularization (PR #824) +- Fix reference number error introduced in PR #767 (PR #819) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Update sgot cost function and example (PR #830) diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index dea2657e1..54573c5a4 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -13,6 +13,7 @@ """ # Author: Paul Krzakala +# Thibaut Germain # 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}" +) ############################################################################# # diff --git a/ot/__init__.py b/ot/__init__.py index 59cd4f2d6..d867daae4 100644 --- a/ot/__init__.py +++ b/ot/__init__.py @@ -94,7 +94,7 @@ from .utils import dist, unif, tic, toc, toq -__version__ = "0.9.7.dev0" +__version__ = "0.9.7" __all__ = [ "emd", diff --git a/ot/batch/__init__.py b/ot/batch/__init__.py index d21019b98..93b998cff 100644 --- a/ot/batch/__init__.py +++ b/ot/batch/__init__.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # 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", ] diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 1a9ec1955..3c07ae639 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # License: MIT License @@ -14,9 +15,17 @@ bregman_log_projection_batch, bregman_projection_batch, entropy_batch, + proximal_bregman_log_plan_batch, ) +solve_batch_method_lst = ["auto", "proximal", "log_sinkhorn", "sinkhorn"] + +solve_batch_reg_type_lst = ["kl", "entropy"] + +solve_batch_grad_lst = ["detach", "autodiff", "last_step", "envelope"] + + def dist_lp_batch(X, Y, p=2, q=1, nx=None): r"""Computes the cost matrix for a batch of samples using the Lp norm. @@ -235,23 +244,38 @@ def dist_batch( def solve_batch( M, - reg, + reg=None, a=None, b=None, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + method="auto", + inner_iter=1, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object. + + The function solves in parallel a batch of optimal transport problems: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{M} \rangle_F + \textit{reg} \cdot R(\mathbf{T}) \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The problem is solved with either a proximal point method :ref:`[92] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function 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. Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix reg : float - Regularization parameter for entropic regularization + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -260,18 +284,22 @@ def solve_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. + inner_iter : int + Number of inner Bregman iterations for the proximal method. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. - 'detach' does not compute the gradients for the Sinkhorn solver. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. + 'detach' does not compute the gradients. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. + Returns ------- @@ -292,19 +320,57 @@ def solve_batch( >>> X = np.random.randn(5, 10, 3) # 5 batches of 10 samples in 3D >>> Y = np.random.randn(5, 15, 3) # 5 batches of 15 samples in 3D >>> M = dist_batch(X, Y, metric="euclidean") # Compute cost matrices + >>> p_result = solve_batch(M) # Uses proximal method >>> reg = 0.1 - >>> result = solve_batch(M, reg) - >>> result.plan.shape # Optimal transport plans for each batch + >>> s_result = solve_batch(M, reg, method="log_sinkhorn") # Uses Sinkhorn method + >>> s_result.plan.shape # Optimal transport plans for each batch (5, 10, 15) - >>> result.value.shape # Optimal transport values for each batch + >>> s_result.value.shape # Optimal transport values for each batch (5,) See Also -------- ot.batch.dist_batch : batched cost matrix computation for computing M. - ot.solve : non-batched version of the OT solver. + ot.solve : non-batched version of the solve_batch function. + + .. _references-batch-solver: + Reference + ---------- + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + + .. [2] M. Cuturi, Sinkhorn Distances : Lightspeed Computation + of Optimal Transport, Advances in Neural Information Processing + Systems (NIPS) 26, 2013 """ + if method not in solve_batch_method_lst: + raise ValueError( + f"Unknown method: {method}. Must be one of {solve_batch_method_lst}." + ) + + if reg_type not in solve_batch_reg_type_lst: + raise ValueError( + f"Unknown reg_type: {reg_type}. Must be one of {solve_batch_reg_type_lst}." + ) + + if grad not in solve_batch_grad_lst: + raise ValueError( + f"Unknown grad: {grad}. Must be one of {solve_batch_grad_lst}." + ) + + if method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0): + raise ValueError( + "Sinkhorn methods require a strictly positive reg parameter. Please provide a valid reg value." + ) + + if method == "auto": + if reg is None or reg == 0: + method = "proximal" + else: + method = "log_sinkhorn" + nx = get_backend(a, b, M) B, n, m = M.shape @@ -314,18 +380,29 @@ def solve_batch( if b is None: b = nx.ones((B, m), type_as=M) / m - if solver == "log_sinkhorn": + if method == "log_sinkhorn": K = -M / reg out = bregman_log_projection_batch( K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad ) - elif solver == "sinkhorn": + if method == "sinkhorn": K = nx.exp(-M / reg) out = bregman_projection_batch( K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad ) - else: - raise ValueError(f"Unknown solver: {solver}") + if method == "proximal": + out = proximal_bregman_log_plan_batch( + M, + a, + b, + nx=nx, + reg=reg, + inner_reg=inner_reg, + max_iter=max_iter, + tol=tol, + inner_iter=inner_iter, + grad=grad, + ) T = out["T"] @@ -336,13 +413,15 @@ def solve_batch( T = nx.detach(T) value_linear = loss_linear_batch(M, T) - if reg_type.lower() == "entropy": + if reg_type == "entropy" and reg is not None: entr = -entropy_batch(T, nx=nx) value = value_linear + reg * entr - elif reg_type.lower() == "kl": + elif reg_type == "kl" and reg is not None: ref = nx.einsum("bi,bj->bij", a, b) kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) value = value_linear + reg * kl + else: + value = value_linear log = {"n_iter": out["n_iters"]} res = OTResult( @@ -360,29 +439,36 @@ def solve_batch( def solve_sample_batch( X_a, X_b, - reg, + reg=None, a=None, b=None, metric="sqeuclidean", p=2, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + method="auto", + inner_iter=1, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object computed from batches of source and target samples. + + The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function 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. Parameters ---------- - M : array-like, shape (B, ns, nt) - Cost matrix - reg : float - Regularization parameter for entropic regularization + X_a : array-like, shape (B, ns, d) + Samples from source distribution + X_b : array-like, shape (B, nt, d) + Samples from target distribution metric : str, optional 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' p : float, optional p-norm for the Minkowski metrics. Default value is 2. + reg : float + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -391,18 +477,21 @@ def solve_sample_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. + inner_iter : int + Number of inner Bregman iterations for the proximal method. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. - 'detach' does not compute the gradients for the Sinkhorn solver. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. + 'detach' does not compute the gradients. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. Returns ------- @@ -418,10 +507,11 @@ def solve_sample_batch( See Also -------- - ot.batch.solve_batch : solver for computing the optimal T from arbitrary cost matrix M. + ot.batch.solve_batch : function for computing the optimal T from arbitrary cost matrix M. """ M = dist_batch(X_a, X_b, metric=metric, p=p) + return solve_batch( M, reg, @@ -429,7 +519,9 @@ def solve_sample_batch( b=b, max_iter=max_iter, tol=tol, - solver=solver, + method=method, + inner_iter=inner_iter, + inner_reg=inner_reg, reg_type=reg_type, grad=grad, ) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index ce5db4664..8bc96c3b9 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -1,3 +1,14 @@ +# -*- coding: utf-8 -*- +""" +Utility functions for batch operations in optimal transport. +""" + +# Author: Remi Flamary +# Paul Krzakala +# Thibaut Germain +# +# License: MIT License + from ot.backend import get_backend @@ -288,3 +299,146 @@ def bregman_log_projection_batch( log_potentials = (u, v) return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} + + +def proximal_bregman_log_plan_batch( + C, + a=None, + b=None, + nx=None, + reg=None, + inner_reg=1e-3, + max_iter=10000, + tol=1e-5, + inner_iter=1, + grad="detach", +): + r""" + Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[92] `. + + This function solves the following optimization problem: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{C} \rangle_F + \textit{reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: + + .. math:: + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + (\textit{reg} + \textit{inner\_reg}) \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + + Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} + \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + + .. math:: + \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) + + \mathbf{v}^{(i+1)} = \log(\mathbf{b}) - \text{LSE}((\mathbf{K}^{(k)})^T + \mathbf{u}^{(i)}) + + where LSE denotes the log-sum-exp operation. + + Parameters + ---------- + C : array-like, shape (B, n, m) + Cost matrix for each problem in the batch. + a : array-like, shape (B, n), optional + Source distribution for each problem. If None, uniform distribution is used. + b : array-like, shape (B, m), optional + Target distribution for each problem. If None, uniform distribution is used. + nx : backend object, optional + Numerical backend to use for computations. If None, the default backend is used. + max_iter : int, optional + Maximum number of iterations. + inner_iter : int, optional + Number of inner iterations for updating the dual variables u and v. Default is 1. + inner_reg : float, optional + Regularization parameter for the Bregman divergence. Default is 1e-3. + tol : float, optional + Tolerance for convergence. The solver stops when the maximum change in + the dual variables is below this value. + grad : str, optional + Gradient computation mode: 'detach', 'autodiff', or 'last_step'. + + Returns + ------- + dict + Dictionary containing: + - 'T' : array-like, shape (B, n, m) + Optimal transport plan for each problem. + - 'log_T' : array-like, shape (B, n, m) + Logarithm of the optimal transport plan for each problem. + - 'potentials' : tuple of array-like, shapes ((B, n), (B, m)) + Log-scaling factors (u, v). + - 'n_iters' : int + Number of iterations performed. + + Example + -------- + >>> import numpy as np + >>> from ot.batch import proximal_bregman_log_plan_batch + >>> # Create batch of affinity matrices + >>> C = np.random.rand(5, 10, 15) # 5 problems, 10x15 cost matrices + >>> result = proximal_bregman_log_plan_batch(C) + >>> T = result['T'] # Shape (5, 10, 15) + + .. _references-OT-bregman-proximal-point: + Reference + ---------- + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + """ + + if nx is None: + nx = get_backend(a, b, C) + + B, n, m = C.shape + + if a is None: + a = nx.ones((B, n)) / n + if b is None: + b = nx.ones((B, m)) / m + + if reg is None: + reg = 0.0 + + u = nx.zeros((B, n), type_as=C) + v = nx.zeros((B, m), type_as=C) + + loga = nx.log(a) + logb = nx.log(b) + + if grad == "detach": + C = nx.detach(C) + elif grad == "last_step": + C_, C = C.clone(), nx.detach(C) + + log_T = nx.zeros(C.shape, type_as=C) + for n_iters in range(max_iter): + K_proj = -(C + inner_reg * log_T) / (reg + inner_reg) + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + + # Check convergence once every 10 iterations + if n_iters % 10 == 0: + T = nx.exp(log_T) + marginal = nx.sum(T, axis=2) + err = nx.max(norm_batch(marginal - a)) + if err < tol: + break + + if grad == "last_step": + K_proj = -(C_ + inner_reg * log_T) / (reg + inner_reg) + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + T = nx.exp(log_T) + log_potentials = (u, v) + + return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} diff --git a/ot/da.py b/ot/da.py index 527999cba..7b4ed7ba2 100644 --- a/ot/da.py +++ b/ot/da.py @@ -48,6 +48,9 @@ joint_OT_mapping_kernel, ) +# value used in ys/yt to mark a sample whose label is unknown +MISSING_LABEL = -1 + def sinkhorn_lpl1_mm( a, @@ -598,23 +601,24 @@ class label if self.limit_max != np.inf: self.limit_max = self.limit_max * nx.max(self.cost_) - # missing_labels is a (ns, nt) matrix of {0, 1} such that - # the cells (i, j) has 0 iff either ys[i] or yt[j] is masked - missing_ys = (ys == -1) + nx.zeros(ys.shape, type_as=ys) - missing_yt = (yt == -1) + nx.zeros(yt.shape, type_as=yt) - missing_labels = missing_ys[:, None] @ missing_yt[None, :] - # labels_match is a (ns, nt) matrix of {True, False} such that - # the cells (i, j) has False if ys[i] != yt[i] - label_match = (ys[:, None] - yt[None, :]) != 0 + # present_labels is a (ns, nt) matrix of {0, 1} such that + # the cell (i, j) is 1 iff both ys[i] and yt[j] are labeled + # (i.e. neither is masked with MISSING_LABEL) + present_ys = (ys != MISSING_LABEL) + nx.zeros(ys.shape, type_as=ys) + present_yt = (yt != MISSING_LABEL) + nx.zeros(yt.shape, type_as=yt) + present_labels = present_ys[:, None] @ present_yt[None, :] + # label_mismatch is a (ns, nt) matrix of {True, False} such that + # the cell (i, j) is True if ys[i] != yt[j] + label_mismatch = (ys[:, None] - yt[None, :]) != 0 # cost correction is a (ns, nt) matrix of {-Inf, float, Inf} such # that he cells (i, j) has -Inf where there's no correction necessary - # by 'correction' we mean setting cost to a large value when - # labels do not match + # by 'correction' we mean setting cost to a large value when two + # labeled samples have different labels # we suppress potential RuntimeWarning caused by Inf multiplication # (as we explicitly cover potential NANs later) with warnings.catch_warnings(): warnings.simplefilter("ignore", category=RuntimeWarning) - cost_correction = label_match * missing_labels * self.limit_max + cost_correction = label_mismatch * present_labels * self.limit_max # this operation is necessary because 0 * Inf = NAN # thus is irrelevant when limit_max is finite cost_correction = nx.nan_to_num(cost_correction, -np.inf) diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index c9fa676c4..32bbbf550 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -36,6 +36,15 @@ linear_circular_ot, ) +from .solver_tree import ( + topological_sort, + tree_wasserstein_distance, +) + +from .tree_barycenter import ( + fixed_support_tree_barycenter, +) + __all__ = [ "emd", "emd2", @@ -60,4 +69,7 @@ "free_support_barycenter_generic_costs", "NorthWestMMGluing", "ot_barycenter_energy", + "topological_sort", + "tree_wasserstein_distance", + "fixed_support_tree_barycenter", ] diff --git a/ot/lp/solver_tree.py b/ot/lp/solver_tree.py new file mode 100644 index 000000000..079227c83 --- /dev/null +++ b/ot/lp/solver_tree.py @@ -0,0 +1,183 @@ +from ..backend import get_backend +import numpy as np +from collections import deque + +""" +Solver for the tree wasserstein distance +""" + +# Author : Ali Boudjema + + +def topological_sort(tree): + r""" + Computes a topological order of the given tree + + Parameters + ----------- + tree: array_like, shape(n) + ancestor of each node in the tree (ancestor of root is root) + """ + + n = tree.shape[0] + + in_degree = np.zeros(n, dtype=int) + + for cur_node in range(n): + if cur_node != tree[cur_node]: + in_degree[tree[cur_node]] += 1 + + queue = deque() + + for cur_node in range(n): + if in_degree[cur_node] == 0: + queue.append(cur_node) + + topo_order = [] + + while queue: + cur_node = queue.popleft() + topo_order.append(cur_node) + + ancestor = tree[cur_node] + + if cur_node != ancestor: + in_degree[ancestor] -= 1 + + if in_degree[ancestor] == 0: + queue.append(ancestor) + + return np.array(topo_order) + + +def tree_wasserstein_distance( + tree, length, u_weights, v_weights, topo_order=None, return_plans=False +): + r""" + Computes the tree wasserstein distance for a given tree between two empirical distributions + + Parameters + ---------- + tree : array_like, shape(n) + parent of each node in the tree (parent of root is root) + length : array_like, shape(n) + length of the edge above each node (length of root is 0) + u_weights : array_like, shape(n) + weights of the first empirical distributions + v_weights : array_like, shape(n) + weights of the second empirical distributions + topo_order : array_like, shape(n), optional + topological order of the tree + return_plans : bool, optional + if True, returns the optimal transport plan between the + two distributions, default is False + + Returns + ------- + cost : float + The tree wasserstein distance + plans : coo_matrix, optional + If return_plans is True, returns a coo_matrix containing the plan + + Reference + --------- + The proof of this algorithm uses the formula (3) in the article + Tree-Sliced Variants of Wasserstein Distances + """ + + n = tree.shape[0] + + assert ( + n == length.shape[0] == u_weights.shape[0] == v_weights.shape[0] + ), "dimension error in the input" + + if topo_order is None: + topo_order = topological_sort(tree) + + nx = get_backend(length, u_weights, v_weights) + + mass_dict = {} + + for cur in range(n): + if u_weights[cur] != v_weights[cur]: + mass_dict[cur] = {cur: u_weights[cur] - v_weights[cur]} + else: + mass_dict[cur] = {} + + source_plan = [] + sink_plan = [] + mass_plan = [] + + virt_size = [len(mass_dict[k]) for k in range(n)] + + cost = 0 + + depth = nx.zeros(n) + + for i in range(n - 2, -1, -1): + cur_node = topo_order[i] + depth[cur_node] = depth[tree[cur_node]] + length[cur_node] + + for cur in topo_order: + dict_cur = mass_dict[cur] + p = tree[cur] + + if cur != p: + dict_p = mass_dict[p] + + if virt_size[cur] > virt_size[p]: + mass_dict[cur], mass_dict[p] = dict_p, dict_cur + dict_cur, dict_p = dict_p, dict_cur + virt_size[cur], virt_size[p] = virt_size[p], virt_size[cur] + + while len(dict_cur) > 0 and len(dict_p) > 0: + node_scur = next(iter(dict_cur)) + amount_scur = dict_cur[node_scur] + + node_sp = next(iter(dict_p)) + amount_sp = dict_p[node_sp] + + if (amount_scur > 0) != (amount_sp > 0): + match_amount = min(abs(amount_scur), abs(amount_sp)) + + source = node_scur if amount_scur > 0 else node_sp + sink = node_sp if amount_scur > 0 else node_scur + + source_plan.append(source) + sink_plan.append(sink) + mass_plan.append(match_amount) + + length_path = depth[source] + depth[sink] - 2 * depth[p] + cost += match_amount * length_path + + if amount_scur > 0: + dict_cur[node_scur] -= match_amount + dict_p[node_sp] += match_amount + else: + dict_cur[node_scur] += match_amount + dict_p[node_sp] -= match_amount + + if dict_cur[node_scur] == 0: + del dict_cur[node_scur] + + if dict_p[node_sp] == 0: + del dict_p[node_sp] + + else: + dict_p[node_scur] = amount_scur + del dict_cur[node_scur] + + if len(dict_p) == 0: + mass_dict[cur], mass_dict[p] = dict_p, dict_cur + dict_cur, dict_p = dict_p, dict_cur + + virt_size[p] += virt_size[cur] + + plans = nx.coo_matrix( + mass_plan, source_plan, sink_plan, shape=(n, n), type_as=length + ) + + if return_plans: + return cost, plans + else: + return cost diff --git a/ot/lp/tree_barycenter.py b/ot/lp/tree_barycenter.py new file mode 100644 index 000000000..6a7811580 --- /dev/null +++ b/ot/lp/tree_barycenter.py @@ -0,0 +1,174 @@ +from ..backend import get_backend +import numpy as np +from ..utils import proj_simplex +from ..utils import list_to_array + +# Author : Ali Boudjema + +# IMPORTANT : ON PREND COMME CONVENTION QUE LES FEUILLES SONT LES PREMIERS SOMMETS DE L'ARBRE + + +def get_B_matrix(tree, length, nb_leafs): + nx = get_backend(length) + + rows = [] + col = [] + data = [] + + for cur_leaf in range(nb_leafs): + cur_node = cur_leaf + + while cur_node != tree[cur_node]: + rows.append(cur_node) + col.append(cur_leaf) + data.append(length[cur_node]) + cur_node = tree[cur_node] + + nb_rows = len(tree) + nb_col = nb_leafs + + B = nx.coo_matrix(data, rows, col, shape=(nb_rows, nb_col), type_as=length) + + return B + + +def get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes): + idx = np.zeros(nb_nodes, dtype=int) + + for node in range(nb_nodes): + idx[node] = np.searchsorted(B_mes_sorted[:, node], cur_B[node]) + 1 + + z = -nb_mes + 2 * idx - 2 + + g = B.T.dot(z) / nb_mes + + return g + + +def pre_process_trees(tree_list, length_list, measures): + nx = get_backend(length_list, measures) + + if tree_list.ndim == 1: + tree_list = nx.reshape(tree_list, (1, *tree_list.shape)) + length_list = nx.reshape(length_list, (1, *length_list.shape)) + measures = nx.reshape(measures, (1, *measures.shape)) + + nb_leafs = measures.shape[2] + + prepared_trees = [] + + for tree, length, mes in zip(tree_list, length_list, measures): + B = get_B_matrix(tree, length, nb_leafs) + nb_mes = mes.shape[0] + + B_mes = list_to_array([B.dot(mes[i]) for i in range(nb_mes)]) + + B_mes_sorted = nx.sort(B_mes, axis=0) + + prepared_trees.append( + { + "B": B, + "B_mes_sorted": B_mes_sorted, + "nb_nodes": tree.shape[0], + "nb_mes": nb_mes, + } + ) + + return prepared_trees + + +def fixed_support_tree_barycenter( + tree_list, length_list, measures, nb_itr=100, step=0.01, tol=1e-5, init_measure=None +): + """ + Computes the Tree-Wasserstein (or Tree-Sliced) barycenter for one or multiple trees, + with the constraint that the support of the barycenter is fixed at the leaves. + It is assumed that the leaves correspond to the first nodes of the tree (indices 0 to k-1). + While the number of leaves (k) must be strictly identical across all structures to ensure + consistent alignment, the total number of nodes (leaves + internal nodes) can + freely vary from one tree to another. + + If a single tree structure is provided (e.g., 1D arrays for tree_list/length_list + and 2D for measures), the function automatically expands their dimensions to 3D + internal structures to handle them uniformly as a multi-tree setting with t=1. + + Parameters + ----------- + tree_list : array_like + A single tree of shape (n_t,) or a list of t trees where each tree has shape (n_t,). + n_t is the total number of nodes in that specific tree. tree[i] contains the index + of the parent of node i (with tree[root] == root). + length_list : array_like + The edge weights corresponding to tree_list. A single array of shape (n_t,) or a list + of t arrays, where the t-th array has shape (n_t,) and contains the length of the edge + connecting node i to its parent. + measures : array_like, shape (t, m, k) or (m,k) + The input probability distributions mapped to the leaves. + k is the fixed number of leaves shared by all trees, and m is the number + of measures. Accepts a 2D array of shape (m, k) for a single tree, or a 3D array + of shape (t, m, k) in a multi-tree setting. + nb_tr : int, optional + the maximal number of iterations for the subgradient descent + step : float, optional + the step size of the descent + tol : float, optional + Convergence tolerance. The descent stops if the L2 norm of the difference + between two consecutive iterations is smaller than tol. + init_measure : array_like, shape (k), optional + The starting point of the descent, default is None + + Returns + ------- + cur_mes : array_like, shape (k) + The computed fixed-support barycenter supported on the k leaves. + + References + ---------- + "Fixed Support Tree-Sliced Wasserstein Barycenter" + """ + + nx = get_backend(length_list, measures) + + assert ( + tree_list.shape[0] == length_list.shape[0] == measures.shape[0] + and tree_list.shape[1] == length_list.shape[1] + ), "dimension error in the input" + + nb_leafs = measures.shape[2] + + if init_measure is None: + cur_mes = nx.ones(nb_leafs) / nb_leafs + else: + cur_mes = init_measure + + prepared_trees = pre_process_trees(tree_list, length_list, measures) + + nb_tree = len(prepared_trees) + + for itr in range(nb_itr): + old_mes = cur_mes.copy() + g_total = nx.zeros(nb_leafs) + + for tree_data in prepared_trees: + B = tree_data["B"] + B_mes_sorted = tree_data["B_mes_sorted"] + nb_nodes = tree_data["nb_nodes"] + nb_mes = tree_data["nb_mes"] + + cur_B = B.dot(cur_mes) + + g_tree = get_gradient(cur_B, B_mes_sorted, B, nb_mes, nb_nodes) + g_total += g_tree + + g_mean = g_total / nb_tree + + g_mean /= np.linalg.norm(g_mean, ord=2) + + cur_mes -= step * g_mean + + cur_mes = proj_simplex(cur_mes) + + if np.linalg.norm(cur_mes - old_mes) < tol: + break + + return cur_mes diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 17d459a43..08e5ff81e 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -3,6 +3,7 @@ # Author: Remi Flamary # Paul Krzakala # Sonia Mazelet +# Thibaut Germain # # License: MIT License @@ -15,50 +16,85 @@ loss_linear_batch, bregman_projection_batch, bregman_log_projection_batch, + proximal_bregman_log_plan_batch, ) from ot import solve +from contextlib import nullcontext import pytest from ot.backend import torch -@pytest.mark.parametrize("solver", ["sinkhorn", "log_sinkhorn"]) +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) +@pytest.mark.parametrize("method", ["auto", "proximal", "sinkhorn", "log_sinkhorn"]) @pytest.mark.parametrize("reg_type", ["kl", "entropy"]) -def test_solve_batch(solver, reg_type): +def test_solve_batch_vs_solve(reg, method, reg_type): """Check that solve_batch gives the same results as solve for each instance in the batch.""" - batchsize = 4 - n = 16 - rng = np.random.RandomState(0) - - M = rng.rand(batchsize, n, n) - - reg = 0.1 - max_iter = 10000 - tol = 1e-5 - res = solve_batch( - M, - a=None, - b=None, + should_fail = method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0) + + ctx = pytest.raises(Exception) if should_fail else nullcontext() + + with ctx: + tol = 1e-4 + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + + base_plan = np.zeros((batchsize, n, d)) + base_value = np.zeros(batchsize) + for i in range(batchsize): + C_i = C[i] + res_i = solve(C_i, reg=reg, tol=tol, reg_type=reg_type) + base_plan[i] = res_i.plan + base_value[i] = res_i.value_linear + + res = solve_batch( + C, + max_iter=10000, + tol=tol, + grad="detach", + reg=reg, + method=method, + reg_type=reg_type, + inner_reg=1e-3, + ) + plan = res.plan + value = res.value_linear + np.testing.assert_allclose(plan, base_plan, atol=tol * 10) + np.testing.assert_allclose(value, base_value, atol=tol * 10) + + +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) +@pytest.mark.parametrize("inner_iter", [1, 5, 10]) +def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): + tol = 1e-4 + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + res = proximal_bregman_log_plan_batch( + nx.from_numpy(C), reg=reg, - max_iter=max_iter, + inner_reg=1e-3, + max_iter=10000, tol=tol, - solver=solver, - reg_type=reg_type, + inner_iter=inner_iter, grad="detach", ) - plan_batch = res.plan - values_batch = res.value_linear - + plan = nx.to_numpy(res["T"]) for i in range(batchsize): - M_i = M[i] + C_i = C[i] res_i = solve( - M_i, a=None, b=None, reg=reg, max_iter=max_iter, tol=tol, reg_type=reg_type + C_i, + reg=reg, + tol=tol, ) plan_i = res_i.plan - value_i = res_i.value_linear - np.testing.assert_allclose(plan_i, plan_batch[i], atol=1e-05) - np.testing.assert_allclose(value_i, values_batch[i], atol=1e-4) + np.testing.assert_allclose(plan_i, plan[i], atol=tol * 10) def test_bregman_batch(): @@ -78,9 +114,10 @@ def test_bregman_batch(): @pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean", "minkowski", "kl"]) -def test_metrics(metric): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_sample_solve_batch_vs_solve_batch(metric, method): """Check that all functions run without error.""" - + tol = 1e-5 batchsize = 2 n = 4 d = 2 @@ -93,11 +130,10 @@ def test_metrics(metric): is_positive = M >= 0 np.testing.assert_equal(is_positive.all(), True) - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, metric=metric) + res = solve_sample_batch( + X, X, reg=0.1, max_iter=10, tol=tol, metric=metric, method=method + ) # Compute loss loss = res.value_linear # loss given by solver @@ -105,34 +141,35 @@ def test_metrics(metric): loss3 = loss_linear_samples_batch( X, X, res.plan, metric=metric ) # recompute loss from plan and samples - np.testing.assert_allclose(loss, loss2, atol=1e-5) - np.testing.assert_allclose(loss, loss3, atol=1e-5) + np.testing.assert_allclose(loss, loss2, atol=tol * 10) + np.testing.assert_allclose(loss, loss3, atol=tol * 10) @pytest.mark.skipif(not torch, reason="torch not installed") @pytest.mark.parametrize("grad", ["detach", "envelope", "autodiff", "last_step"]) -def test_gradients_torch(grad): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_gradients_torch(grad, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 d = 2 - for solver in ["sinkhorn", "log_sinkhorn"]: - X = torch.randn((batchsize, n, d), requires_grad=True) - M = dist_batch(X, X) - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, solver=solver) - loss = res.value_linear.sum() - loss_plan = res.plan.sum() - if grad == "detach": - assert loss.grad == None - elif grad == "envelope": - loss.backward() - assert X.grad is not None - elif grad in ["autodiff", "last_step"]: - loss_plan.backward() - assert X.grad is not None - - -def test_backend(nx): + X = torch.randn((batchsize, n, d), requires_grad=True) + M = dist_batch(X, X) + res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, method=method) + loss = res.value_linear.sum() + loss_plan = res.plan.sum() + if grad == "detach": + assert loss.grad == None + elif grad == "envelope": + loss.backward() + assert X.grad is not None + elif grad in ["autodiff", "last_step"]: + loss_plan.backward() + assert X.grad is not None + + +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_backend(nx, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 @@ -140,30 +177,5 @@ def test_backend(nx): X = np.random.randn(batchsize, n, d) X = nx.from_numpy(X) M = dist_batch(X, X) - solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5) - - -def test_metric_default_parameters(): - """Check that all functions with default parameters run without error.""" - - batchsize = 2 - n = 4 - d = 2 - rng = np.random.RandomState(0) - X = rng.rand(batchsize, n, d) - M = dist_batch(X, X) - is_positive = M >= 0 - np.testing.assert_equal(is_positive.all(), True) - - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1) - - # Compute loss - loss_linear_batch(M, res.plan) # recompute loss from plan - loss_linear_samples_batch(X, X, res.plan) # recompute loss from plan and samples - assert np.isfinite(loss_linear_batch(M, res.plan)).all() - assert np.isfinite(loss_linear_samples_batch(X, X, res.plan)).all() + solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, method=method) + solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, method=method) diff --git a/test/test_da.py b/test/test_da.py index bb548e27f..df224a247 100644 --- a/test/test_da.py +++ b/test/test_da.py @@ -166,7 +166,7 @@ def test_sinkhorn_lpl1_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -258,7 +258,7 @@ def test_sinkhorn_l1l2_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -356,7 +356,7 @@ def test_sinkhorn_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -472,7 +472,7 @@ def test_unbalanced_sinkhorn_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -571,7 +571,7 @@ def test_emd_transport_class(nx): n_semisup = nx.sum(otda_semi.cost_) # check that the cost matrix norms are indeed different - assert np.allclose( + assert not np.allclose( n_unsup, n_semisup, atol=1e-7 ), "semisupervised mode is not working" @@ -586,6 +586,39 @@ def test_emd_transport_class(nx): ) +@pytest.skip_backend("tf") +def test_semisupervised_cost_correction(nx): + # gh-664: the label-aware cost correction must push apart labeled source + # and target samples with *different* labels, and must leave pairs + # involving an unlabeled (-1) sample untouched. + rng = np.random.RandomState(0) + Xs = rng.randn(6, 2) + Xt = rng.randn(6, 2) + ys = np.array([0, 0, 0, 1, 1, 1]) + yt = np.array([0, 1, 0, 1, 0, 1]) + Xs, ys, Xt, yt = nx.from_numpy(Xs, ys, Xt, yt) + + # all labels known: different-label pairs get the (scaled) limit_max, + # same-label pairs keep their original, smaller cost + otda = ot.da.EMDTransport(limit_max=10) + otda.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt) + cost = nx.to_numpy(otda.cost_) + limit = float(nx.to_numpy(otda.limit_max)) + ys_np, yt_np = nx.to_numpy(ys), nx.to_numpy(yt) + mismatch = ys_np[:, None] != yt_np[None, :] + assert np.all(cost[mismatch] >= limit - 1e-9), "different labels not penalized" + assert np.all(cost[~mismatch] < limit), "same labels wrongly penalized" + + # semi-supervised: unlabeled (-1) targets must never be forbidden + yt_semi = nx.from_numpy(np.array([0, -1, 0, -1, 0, -1])) + otda2 = ot.da.EMDTransport(limit_max=10) + otda2.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt_semi) + cost2 = nx.to_numpy(otda2.cost_) + limit2 = float(nx.to_numpy(otda2.limit_max)) + unlabeled = nx.to_numpy(yt_semi) == -1 + assert np.all(cost2[:, unlabeled] < limit2), "unlabeled targets wrongly forbidden" + + @pytest.skip_backend("jax") @pytest.skip_backend("tf") @pytest.mark.parametrize("kernel", ["linear", "gaussian"])