diff --git a/README.md b/README.md index 950a509e9..87e5b9732 100644 --- a/README.md +++ b/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. \ 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..5fb40e863 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -30,7 +30,7 @@ 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) @@ -38,6 +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) +- 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`. 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/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/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)