From bb33a5e03c51f9ef235632601d367d094068e041 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 24 Oct 2025 15:07:05 +0200 Subject: [PATCH 01/15] add fugw batch loss to ot.batch._quadratic --- ot/batch/_quadratic.py | 152 ++++++++++++++++++++++ test/batch/test_solve_unbalanced_batch.py | 0 2 files changed, 152 insertions(+) create mode 100644 test/batch/test_solve_unbalanced_batch.py diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index 0da4b8962..e549c7e72 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -152,6 +152,90 @@ def h2(C2): return compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=symmetric) +def div_to_product_batch( + T, a, b, T1=None, T2=None, divergence="kl", mass=True, nx=None +): + r"""Fast computation of the Bregman divergence between a batch of arbitrary measures and a product measures. + Only support for Kullback-Leibler and half-squared L2 divergences. + + - For half-squared L2 divergence: + + .. math:: + \frac{1}{2} || \pi - a \otimes b ||^2 + = \frac{1}{2} \Big[ \sum_{i, j} \pi_{ij}^2 + (\sum_i a_i^2) ( \sum_j b_j^2) - 2 \sum_{i, j} a_i \pi_{ij} b_j \Big] + + - For Kullback-Leibler divergence: + + .. math:: + KL(\pi | a \otimes b) + = \langle \pi, \log \pi \rangle - \langle \pi_1, \log a \rangle + - \langle \pi_2, \log b \rangle - m(\pi) + m(a) m(b) + + where : + + - :math:`\pi` is the (`dim_a`, `dim_b`) transport plan + - :math:`\pi_1` and :math:`\pi_2` are the marginal distributions + - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions + - :math:`m` denotes the mass of the measure + + Parameters + ---------- + pi : array-like (B, n, m) + Transport plan for each problem in the batch + a : array-like (B,n) + Unnormalized histogram of dimension `n` for each problem in the batch + b : array-like (B,m) + Unnormalized histogram of dimension `m` for each problem in the batch + T1 : array-like (B, n), optional (default = None) + Marginal distribution with respect to the first dimension of the transport plan for each problem in the batch + Only used in case of Kullback-Leibler divergence. + T2 : array-like (B, m), optional (default = None) + Marginal distribution with respect to the second dimension of the transport plan for each problem in the batch + Only used in case of Kullback-Leibler divergence. + divergence : string, default = "kl" + Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) + mass : bool, optional. Default is False. + Only used in case of Kullback-Leibler divergence. + If False, calculate the relative entropy. + If True, calculate the Kullback-Leibler divergence. + nx : backend, optional + If let to its default value None, a backend test will be conducted. + + Returns + ------- + Bregman divergence between an arbitrary measure and a product measure for each problem in the batch. + """ + + arr = [T, a, b, T1, T2] + + if nx is None: + nx = get_backend(*arr, T1, T2) + + if divergence == "kl": + if T1 is None: + T1 = nx.sum(T, 2) + if T2 is None: + T2 = nx.sum(T, 1) + + if divergence == "kl": + res = ( + nx.sum((T * nx.log(T + 1.0 * (T == 0))), (1, 2)) + - nx.sum(T1 * nx.log(a), 1) + - nx.sum(T2 * nx.log(b), 1) + ) + if mass: + res = res - nx.sum(T1, 1) + nx.sum(a, 1) * nx.sum(b, 1) + + elif divergence == "l2": + res = ( + nx.sum(T**2, (1, 2)) + + nx.sum(a**2, 1) * nx.sum(b**2, 1) + - 2 * nx.sum((a * (T @ b[:, :, None]).squeeze(-1)), 1) + ) / 2 + + return res + + def loss_quadratic_batch(L, T, recompute_const=False, symmetric=True, nx=None): r""" Computes the gromov-wasserstein cost given a cost tensor and transport plan. Batched version. @@ -266,6 +350,74 @@ def loss_quadratic_samples_batch( ) +def loss_fugw_batch( + L, M, T, alpha=0.5, reg_marginals=1, symmetric=True, divergence="kl", nx=None +): + r""" + Computes the fused unbalanced gromov-wasserstein cost given a cost tensor (Gromov term), a cost matrix between features across domains (linear term) and a transport plan. Batched version. + + Parameters + ---------- + L : dict + Cost tensor as returned by `tensor_batch`. + M : array-like, shape (B, n, m) + Cost matrix between features across domains. + T : array-like, shape (B, n, m) + Transport plan. + alpha : float or array-like( B,) optional + Weight the quadratic term (alpha*Gromov) and the linear term + ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. If alpha + a scalar it is used for all problems in the batch. + reg_marginals : float or array-like( B,) optional + Marginal relaxation terms. If rho is + a scalar it is used for all problems in the batch. + symmetric : bool, optional + Whether to use symmetric version. Default is True. + divergence : string, default = "kl" + Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) + nx : module, optional + Backend to use. Default is None. + + Examples + -------- + >>> import numpy as np + >>> from ot.batch import tensor_batch, loss_quadratic_batch + >>> # Create batch of cost matrices + >>> C1 = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices + >>> C2 = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices + >>> a = np.ones((3, 5)) / 5 # Uniform source distributions + >>> b = np.ones((3, 4)) / 4 # Uniform target distributions + >>> L = tensor_batch(a, b, C1, C2, loss='sqeuclidean') + >>> # Use the uniform transport plan for testing + >>> T = np.ones((3, 5, 4)) / (5 * 4) + >>> loss = loss_quadratic_batch(L, T, recompute_const=True) + >>> loss.shape + (3,) + + See Also + -------- + ot.batch.tensor_batch : From computing the cost tensor L. + ot.batch.solve_gromov_batch : For finding the optimal transport plan T. + """ + if nx is None: + nx = get_backend(T) + + Q = loss_quadratic_batch(L, T, recompute_const=True, symmetric=symmetric, nx=nx) + + L = loss_linear_batch(M, T, nx=nx) + + unbalanced = div_to_product_batch( + T, + a=nx.sum(T, axis=2), + b=nx.sum(T, axis=1), + divergence=divergence, + mass=True, + nx=nx, + ) + + return (1 - alpha) * L + alpha * Q + reg_marginals * unbalanced + + def solve_gromov_batch( C1, C2, diff --git a/test/batch/test_solve_unbalanced_batch.py b/test/batch/test_solve_unbalanced_batch.py new file mode 100644 index 000000000..e69de29bb From e58591e5e5a8404aea09eaa3e234afb64a6eee03 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 22 May 2026 16:37:19 +0200 Subject: [PATCH 02/15] add tests and fix functions in batch module --- ot/batch/_linear.py | 2 +- ot/batch/_quadratic.py | 183 ++++++++++++++++++---- test/batch/test_solve_batch.py | 2 +- test/batch/test_solve_gromov_batch.py | 77 ++++++++- test/batch/test_solve_unbalanced_batch.py | 0 5 files changed, 228 insertions(+), 36 deletions(-) delete mode 100644 test/batch/test_solve_unbalanced_batch.py diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index a63fcb404..1a9ec1955 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -147,7 +147,7 @@ def loss_linear_batch(M, T, nx=None): return nx.sum(M * T, axis=(1, 2)) -def loss_linear_samples_batch(X, Y, T, metric="l2"): +def loss_linear_samples_batch(X, Y, T, metric="sqeuclidean"): r"""Computes the linear optimal transport loss given samples and transport plan. This is the equivalent of calling `dist_batch` and then `loss_linear_batch`. diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index e549c7e72..982ef7cfd 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -10,8 +10,9 @@ from ..utils import OTResult from ot.backend import get_backend -from ot.batch._linear import loss_linear_batch +from ot.batch._linear import loss_linear_batch, loss_linear_samples_batch from ot.batch._utils import bmv, bop, bregman_log_projection_batch +from ot.utils import list_to_array def tensor_batch( @@ -289,7 +290,7 @@ def loss_quadratic_samples_batch( C2, T, loss="sqeuclidean", - symmetric=None, + symmetric=True, nx=None, logits=None, recompute_const=False, @@ -351,7 +352,17 @@ def loss_quadratic_samples_batch( def loss_fugw_batch( - L, M, T, alpha=0.5, reg_marginals=1, symmetric=True, divergence="kl", nx=None + a, + b, + L, + M, + T, + alpha=0.5, + reg_marginals=1, + symmetric=True, + divergence="kl", + recompute_const=True, + nx=None, ): r""" Computes the fused unbalanced gromov-wasserstein cost given a cost tensor (Gromov term), a cost matrix between features across domains (linear term) and a transport plan. Batched version. @@ -364,58 +375,172 @@ def loss_fugw_batch( Cost matrix between features across domains. T : array-like, shape (B, n, m) Transport plan. - alpha : float or array-like( B,) optional + alpha : float, array-like or list (B,) optional Weight the quadratic term (alpha*Gromov) and the linear term ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. If alpha a scalar it is used for all problems in the batch. - reg_marginals : float or array-like( B,) optional + reg_marginals : float array-like or list(B,) optional Marginal relaxation terms. If rho is a scalar it is used for all problems in the batch. symmetric : bool, optional Whether to use symmetric version. Default is True. divergence : string, default = "kl" Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) + recompute_const : bool, optional + Whether to recompute the constant term. Default is True. This should be set to True if T does not satisfy the marginal constraints. nx : module, optional Backend to use. Default is None. + """ + if nx is None: + nx = get_backend(T) - Examples - -------- - >>> import numpy as np - >>> from ot.batch import tensor_batch, loss_quadratic_batch - >>> # Create batch of cost matrices - >>> C1 = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices - >>> C2 = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices - >>> a = np.ones((3, 5)) / 5 # Uniform source distributions - >>> b = np.ones((3, 4)) / 4 # Uniform target distributions - >>> L = tensor_batch(a, b, C1, C2, loss='sqeuclidean') - >>> # Use the uniform transport plan for testing - >>> T = np.ones((3, 5, 4)) / (5 * 4) - >>> loss = loss_quadratic_batch(L, T, recompute_const=True) - >>> loss.shape - (3,) + B = T.shape[0] - See Also - -------- - ot.batch.tensor_batch : From computing the cost tensor L. - ot.batch.solve_gromov_batch : For finding the optimal transport plan T. + if isinstance(alpha, list): + alpha = list_to_array(alpha, nx=nx) + + if isinstance(reg_marginals, list): + reg_marginals = list_to_array(reg_marginals, nx=nx) + + if hasattr(alpha, "ndim") and alpha.ndim > 0: + if alpha.ndim != 1 or alpha.shape[0] != B: + raise ValueError( + f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" + ) + + if hasattr(reg_marginals, "ndim") and reg_marginals.ndim > 0: + if reg_marginals.ndim != 1 or reg_marginals.shape[0] != B: + raise ValueError( + f"If reg_marginals is not a scalar, it must have shape ({B},), got {reg_marginals.shape}" + ) + + quadratic = loss_quadratic_batch( + L, T, recompute_const=recompute_const, symmetric=symmetric, nx=nx + ) + + linear = loss_linear_batch(M, T, nx=nx) + + unbalanced = div_to_product_batch( + T, + a, + b, + divergence=divergence, + mass=True, + nx=nx, + ) + + return (1 - alpha) * linear + alpha * quadratic + reg_marginals * unbalanced + + +def loss_fugw_samples_batch( + a, + b, + C1, + C2, + X, + Y, + T, + alpha=0.5, + reg_marginals=1, + symmetric=True, + divergence="kl", + recompute_const=True, + metric_linear="sqeuclidean", + metric_quadratic="sqeuclidean", + logits=None, + nx=None, +): + r""" + Computes the fused unbalanced gromov-wasserstein cost given a cost tensor (quadratic term), a cost matrix between features across domains (linear term) and a transport plan. Batched version. + + Parameters + ---------- + a : array-like, shape (B, n) + Source distributions. + b : array-like, shape (B, m) + Target distributions. + C1 : array-like, shape (B, n, n) or (B, n, n, d) + Source cost matrices for the quadratic term. + C2 : array-like, shape (B, m, m) or (B, n, n, d) + Target cost matrices for the quadratic term. + X : array-like, shape (B, n, d) + Samples from source distribution for the linear term + Y : array-like, shape (B, m, d) + Samples from target distribution for the linear term + T : array-like, shape (B, n, m) + Transport plan. + alpha : float or array-like or list(B,) optional + Weight the quadratic term (alpha*Gromov) and the linear term + ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. If alpha + a scalar it is used for all problems in the batch. + reg_marginals : float or array-like or list(B,) optional + Marginal relaxation terms. If rho is + a scalar it is used for all problems in the batch. + symmetric : bool, optional + Whether to use symmetric version. Default is True. + divergence : string, default = "kl" + Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) + recompute_const : bool, optional + Whether to recompute the constant term. Default is True. This should be set to True if T does not satisfy the marginal constraints. + metric_linear : str, optional + Metric for the linear term, 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' + metric_quadratic : str, optional + Metric to use for the quadratic term. Supported values: 'sqeuclidean', 'kl'. + Default is 'sqeuclidean'. + logits : bool, optional + For KL divergence, whether inputs are logits (unnormalized log probabilities). + If True, inputs are treated as logits. Default is None. + nx : module, optional + Backend to use. Default is None. """ if nx is None: nx = get_backend(T) - Q = loss_quadratic_batch(L, T, recompute_const=True, symmetric=symmetric, nx=nx) + B = T.shape[0] + + if isinstance(alpha, list): + alpha = list_to_array(alpha, nx=nx) + + if isinstance(reg_marginals, list): + reg_marginals = list_to_array(reg_marginals, nx=nx) + + if hasattr(alpha, "ndim") and alpha.ndim > 0: + if alpha.ndim != 1 or alpha.shape[0] != B: + raise ValueError( + f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" + ) + + if hasattr(reg_marginals, "ndim") and reg_marginals.ndim > 0: + if reg_marginals.ndim != 1 or reg_marginals.shape[0] != B: + raise ValueError( + f"If reg_marginals is not a scalar, it must have shape ({B},), got {reg_marginals.shape}" + ) + + quadratic = loss_quadratic_samples_batch( + a, + b, + C1, + C2, + T, + loss=metric_quadratic, + symmetric=symmetric, + nx=nx, + logits=logits, + recompute_const=recompute_const, + ) - L = loss_linear_batch(M, T, nx=nx) + linear = loss_linear_samples_batch(X, Y, T, metric=metric_linear) unbalanced = div_to_product_batch( T, - a=nx.sum(T, axis=2), - b=nx.sum(T, axis=1), + a, + b, divergence=divergence, mass=True, nx=nx, ) - return (1 - alpha) * L + alpha * Q + reg_marginals * unbalanced + return (1 - alpha) * linear + alpha * quadratic + reg_marginals * unbalanced def solve_gromov_batch( diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 45a7e69fe..fc9a74492 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -1,4 +1,4 @@ -"""Tests for module bregman on OT with bregman projections""" +"""Tests for module batch""" # Author: Remi Flamary # Kilian Fatras diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index e0029689b..1518b41d8 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -1,19 +1,24 @@ -"""Tests for module bregman on OT with bregman projections""" +"""Tests for module batch""" # Author: Remi Flamary -# Kilian Fatras -# Quang Huy Tran -# Eduardo Fernandes Montesuma +# Sonia Mazelet + # # License: MIT License import numpy as np -from ot.batch import solve_gromov_batch, loss_quadratic_samples_batch +from ot.batch import ( + solve_gromov_batch, + loss_quadratic_batch, + loss_linear_batch, + loss_quadratic_samples_batch, +) from ot import solve_gromov from ot.batch._linear import dist_batch import pytest from itertools import product from ot.backend import torch +from ot.batch._quadratic import tensor_batch, loss_fugw_batch, loss_fugw_samples_batch def test_solve_gromov_batch(): @@ -133,3 +138,65 @@ def test_backend(nx): C = np.random.randn(batchsize, n, n, d) C = nx.from_numpy(C) solve_gromov_batch(C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False) + + +def test_fugw_loss(): + """Check that loss_fugw_batch and loss_fugw_samples_batch run without error.""" + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + X = rng.rand(batchsize, n, d) + Y = rng.rand(batchsize, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + reg_marginals = 0 + T = rng.rand(batchsize, n, n) + L = tensor_batch(a=a, b=a, C1=C1, C2=C2, loss="sqeuclidean") + alpha = rng.rand() + reg_marginals = rng.rand() + + loss_fugw = loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw_sample = loss_fugw_samples_batch( + a, a, C1, C2, X, Y, T, alpha=alpha, reg_marginals=reg_marginals + ) + assert np.isfinite(loss_fugw).all() + assert np.isfinite(loss_fugw_sample).all() + + alpha = rng.rand(batchsize) + reg_marginals = rng.rand(batchsize) + loss_fugw = loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw_sample = loss_fugw_samples_batch( + a, a, C1, C2, X, Y, T, alpha=alpha, reg_marginals=reg_marginals + ) + assert np.isfinite(loss_fugw).all() + assert np.isfinite(loss_fugw_sample).all() + + +def test_valid_fugw_loss_endpoints(): + """Check that loss_fugw_batch gives the same results as solve_gromov_batch and solve_linear_batch for alpha=0 and alpha=1.""" + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + reg_marginals = 0 + T = rng.rand(batchsize, n, n) + L = tensor_batch(a=a, b=a, C1=C1, C2=C2, loss="sqeuclidean") + + loss_fugw = loss_fugw_batch( + a, a, L, M, T, alpha=0.0, divergence="l2", reg_marginals=reg_marginals + ) + loss_linear = loss_linear_batch(M, T) + np.testing.assert_allclose(loss_fugw, loss_linear, atol=1e-5) + + loss_fugw = loss_fugw_batch( + a, a, L, M, T, alpha=1.0, divergence="l2", reg_marginals=reg_marginals + ) + loss_gromov = loss_quadratic_batch(L, T, recompute_const=True) + np.testing.assert_allclose(loss_fugw, loss_gromov, atol=1e-5) diff --git a/test/batch/test_solve_unbalanced_batch.py b/test/batch/test_solve_unbalanced_batch.py deleted file mode 100644 index e69de29bb..000000000 From 29e1378a8ca80a5f6b371ca055f5314460ad2211 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 22 May 2026 17:03:40 +0200 Subject: [PATCH 03/15] update RELEASES.md --- RELEASES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/RELEASES.md b/RELEASES.md index 0f8918cac..6ff2b2acf 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -14,6 +14,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397) - Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765) - Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765) +- Add batch FUGW loss to `ot.batch` (PR #775) #### Closed issues From 5f2822c7db362aa43ba189245448a72712f20fc9 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 27 May 2026 09:59:16 +0200 Subject: [PATCH 04/15] increase test coverage --- test/batch/test_solve_gromov_batch.py | 29 +++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index 1518b41d8..2b8088129 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -167,12 +167,41 @@ def test_fugw_loss(): alpha = rng.rand(batchsize) reg_marginals = rng.rand(batchsize) + alpha_list = [alpha[i] for i in range(batchsize)] + reg_marginals_list = [reg_marginals[i] for i in range(batchsize)] + loss_fugw = loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) loss_fugw_sample = loss_fugw_samples_batch( a, a, C1, C2, X, Y, T, alpha=alpha, reg_marginals=reg_marginals ) + loss_fugw_list = loss_fugw_batch( + a, a, L, M, T, alpha=alpha_list, reg_marginals=reg_marginals_list + ) + loss_fugw_sample_list = loss_fugw_samples_batch( + a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + ) + assert np.isfinite(loss_fugw).all() assert np.isfinite(loss_fugw_sample).all() + assert np.isfinite(loss_fugw_list).all() + assert np.isfinite(loss_fugw_sample_list).all() + + # check that invalid alpha shape raise an error + alpha = rng.rand(batchsize + 1) + with pytest.raises(ValueError): + loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw_samples_batch( + a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + ) + + # check that invalid rho shape raise an error + alpha = rng.rand(batchsize) + reg_marginals = rng.rand(batchsize + 1) + with pytest.raises(ValueError): + loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw_samples_batch( + a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + ) def test_valid_fugw_loss_endpoints(): From 3bf84be9c56f0f4d847049ce61285fb1a61bb681 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 27 May 2026 13:44:09 +0200 Subject: [PATCH 05/15] add an example and an additional test --- examples/backends/plot_gradient_descent.py | 228 +++++++++++++++++++++ ot/batch/_quadratic.py | 4 + test/batch/test_solve_batch.py | 25 +++ test/batch/test_solve_gromov_batch.py | 104 +++++++++- 4 files changed, 350 insertions(+), 11 deletions(-) create mode 100644 examples/backends/plot_gradient_descent.py diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py new file mode 100644 index 000000000..7ff944792 --- /dev/null +++ b/examples/backends/plot_gradient_descent.py @@ -0,0 +1,228 @@ +# -*- coding: utf-8 -*- +r""" +=============================================================================== +Solve Fused Unbalanced Gromov Wasserstein with gradient descent +=============================================================================== + +Since the FUGW loss is differentiable, it can be minimized with gradient descent. +We show how to do this with the `loss_fugw_batch` function and compare the results with +the dedicated FUGW solver `fused_unbalanced_gromov_wasserstein`. +""" + +# Author: Rémi Flamary +# Sonia Mazelet +# +# License: MIT License + +# sphinx_gallery_thumbnail_number = 2 + +import numpy as np +import matplotlib.pylab as pl +import torch + +import ot +from ot.batch._quadratic import loss_fugw_batch, tensor_batch +from ot.gromov import fused_unbalanced_gromov_wasserstein +from sklearn.manifold import MDS + + +# %% +# Generation of source and target graphs +# ---------------- + +rng = np.random.RandomState(42) + + +def get_sbm(n, nc, ratio, P): + nbpc = np.round(n * ratio).astype(int) + n = np.sum(nbpc) + C = np.zeros((n, n)) + for c1 in range(nc): + for c2 in range(c1 + 1): + if c1 == c2: + for i in range(np.sum(nbpc[:c1]), np.sum(nbpc[: c1 + 1])): + for j in range(np.sum(nbpc[:c2]), i): + if rng.rand() <= P[c1, c2]: + C[i, j] = 1 + else: + for i in range(np.sum(nbpc[:c1]), np.sum(nbpc[: c1 + 1])): + for j in range(np.sum(nbpc[:c2]), np.sum(nbpc[: c2 + 1])): + if rng.rand() <= P[c1, c2]: + C[i, j] = 1 + + return C + C.T + + +def get_position_colors(x): + xmin = x.min(axis=0, keepdims=True) + xmax = x.max(axis=0, keepdims=True) + xnorm = (x - xmin) / np.maximum(xmax - xmin, 1e-15) + colors_x = pl.cm.viridis(xnorm[:, 0])[:, :3] + colors_y = pl.cm.viridis(xnorm[:, 1])[:, :3] + return np.sqrt(colors_x * colors_y) + + +def plot_graph(x, C, color="C0", s=100): + for j in range(C.shape[0]): + for i in range(j): + if C[i, j] > 0: + pl.plot([x[i, 0], x[j, 0]], [x[i, 1], x[j, 1]], alpha=0.2, color="k") + pl.scatter(x[:, 0], x[:, 1], c=color, s=s, zorder=10, edgecolors="k") + + +n1 = 30 +n2 = 20 +nc1 = 3 +nc2 = 2 +ratio1 = np.array([0.33, 0.33, 0.33]) +ratio2 = np.array([0.5, 0.5]) + +P1 = np.array([[0.8, 0.08, 0.0], [0.08, 0.8, 0.08], [0.0, 0.08, 0.8]]) +P2 = np.array(0.6 * np.eye(2) + 0.05 * np.ones((2, 2))) +C1 = get_sbm(n1, nc1, ratio1, P1) +C2 = get_sbm(n2, nc2, ratio2, P2) + +# get 2d position for nodes +x1 = MDS( + metric="precomputed", random_state=0, n_init=1, init="classical_mds" +).fit_transform(1 - C1) +x2 = MDS( + metric="precomputed", random_state=0, n_init=1, init="classical_mds" +).fit_transform(1 - C2) + +colors1 = get_position_colors(x1) +colors2 = get_position_colors(x2) + + +pl.figure(1, (10, 5)) +pl.clf() +pl.subplot(1, 2, 1) +plot_graph(x1, C1, color=colors1) +pl.title("SBM source graph") +pl.axis("off") +pl.subplot(1, 2, 2) +plot_graph(x2, C2, color=colors2) +pl.title("SBM target graph") +_ = pl.axis("off") + + +# %% +# Solve FUGW with gradient descent +# ---------------- + +# Even though `loss_fugw_batch` supports batches of problems, we use a +# batch of size 1 here for clarity. + +a = ot.unif(C1.shape[0]) +b = ot.unif(C2.shape[0]) +M = ot.dist(x1, x2) +M /= M.max() + +a_torch = torch.tensor(a[None, :]) +b_torch = torch.tensor(b[None, :]) +C1_torch = torch.tensor(C1[None, :, :]) +C2_torch = torch.tensor(C2[None, :, :]) +M_torch = torch.tensor(M[None, :, :]) +L = tensor_batch(a_torch, b_torch, C1_torch, C2_torch, loss="sqeuclidean") + +alpha = 0.7 +reg_marginals = 1.0 +lr = 1e-5 +nb_iter_max = 300 + +T_torch = (a_torch[:, :, None] * b_torch[:, None, :]).clone().requires_grad_(True) +loss_iter = [] +mass_iter = [] + +for i in range(nb_iter_max): + loss = loss_fugw_batch( + a_torch, + b_torch, + L, + M_torch, + T_torch, + alpha=alpha, + reg_marginals=reg_marginals, + divergence="kl", + recompute_const=True, + )[0] + + loss_iter.append(float(loss.detach())) + mass_iter.append(float(T_torch.detach().sum())) + loss.backward() + + with torch.no_grad(): + T_torch -= lr * T_torch.grad + T_torch.clamp_(min=1e-12) + T_torch.grad.zero_() + +T_gd = T_torch.detach().cpu().numpy()[0] + +pl.figure(2, (10, 4)) +pl.clf() +pl.subplot(1, 2, 1) +pl.plot(loss_iter) +pl.grid() +pl.title("FUGW loss along gradient descent") +pl.xlabel("Iterations") +pl.subplot(1, 2, 2) +pl.plot(mass_iter) +pl.grid() +pl.title("Transport mass") +_ = pl.xlabel("Iterations") + + +# %% +# Compare with the dedicated FUGW solver +# ------------------------------------- +# +# The dedicated solver uses a block coordinate descent scheme. We compare the +# coupling it returns with the coupling obtained by direct gradient descent on +# `loss_fugw_batch`. The FUGW loss is non convex so minimizing it directly with gradient descent does not +# necessarily give the same solution as the dedicated solver. By comparing the FUGW costs obtained by both methods, +# we find that the BCD solvers gives a better solution than gradient descent. + +T_bcd, _, log = fused_unbalanced_gromov_wasserstein( + C1, + C2, + wx=a, + wy=b, + reg_marginals=reg_marginals, + divergence="kl", + unbalanced_solver="mm", + alpha=alpha, + M=M, + init_pi=np.outer(a, b), + max_iter=100, + tol=1e-7, + max_iter_ot=200, + tol_ot=1e-7, + log=True, +) + + +print("Final batch FUGW loss (gradient descent):", loss_iter[-1]) +print("FUGW cost reported by the dedicated solver:", log["fugw_cost"]) + + +# %% +# Visualize the learned couplings +# ------------------------------- +# We visualize the couplings obtained by both methods to compare them. +# The BCD solver gives sharper plans but both +# methods find couplings that match the structures of the graphs. + +pl.figure(3, (10, 4)) +pl.clf() +pl.subplot(1, 2, 1) +pl.imshow(T_gd, interpolation="nearest") +pl.title("Coupling from gradient descent") +pl.xlabel("Target nodes") +pl.ylabel("Source nodes") +pl.colorbar() +pl.subplot(1, 2, 2) +pl.imshow(T_bcd, interpolation="nearest") +pl.title("Coupling from BCD solver") +pl.xlabel("Target nodes") +pl.ylabel("Source nodes") +_ = pl.colorbar() diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index 982ef7cfd..7ce0ca2fb 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -369,6 +369,10 @@ def loss_fugw_batch( Parameters ---------- + a : array-like, shape (B, n) + Source distributions. + b : array-like, shape (B, m) + Target distributions. L : dict Cost tensor as returned by `tensor_batch`. M : array-like, shape (B, n, m) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index fc9a74492..566116c0a 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -143,3 +143,28 @@ def test_backend(nx): 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() diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index 2b8088129..a98e8d5ba 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -140,7 +140,12 @@ def test_backend(nx): solve_gromov_batch(C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False) -def test_fugw_loss(): +@pytest.mark.parametrize("divergence", ["kl", "l2"]) +@pytest.mark.parametrize( + "metric_linear", ["sqeuclidean", "euclidean", "minkowski", "kl"] +) +@pytest.mark.parametrize("metric_quadratic", ["sqeuclidean", "kl"]) +def test_fugw_loss(divergence, metric_linear, metric_quadratic): """Check that loss_fugw_batch and loss_fugw_samples_batch run without error.""" batchsize = 2 n = 4 @@ -157,41 +162,105 @@ def test_fugw_loss(): L = tensor_batch(a=a, b=a, C1=C1, C2=C2, loss="sqeuclidean") alpha = rng.rand() reg_marginals = rng.rand() + logits = False if metric_quadratic == "kl" else None - loss_fugw = loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw = loss_fugw_batch( + a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals, divergence=divergence + ) loss_fugw_sample = loss_fugw_samples_batch( - a, a, C1, C2, X, Y, T, alpha=alpha, reg_marginals=reg_marginals + a, + a, + C1, + C2, + X, + Y, + T, + alpha=alpha, + reg_marginals=reg_marginals, + divergence=divergence, + metric_linear=metric_linear, + metric_quadratic=metric_quadratic, + logits=logits, ) assert np.isfinite(loss_fugw).all() assert np.isfinite(loss_fugw_sample).all() + # check that alpha and reg_marginals can be passed as lists or arrays of shape (batchsize,) alpha = rng.rand(batchsize) reg_marginals = rng.rand(batchsize) - alpha_list = [alpha[i] for i in range(batchsize)] - reg_marginals_list = [reg_marginals[i] for i in range(batchsize)] + alpha_list = alpha.tolist() + reg_marginals_list = reg_marginals.tolist() - loss_fugw = loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + loss_fugw = loss_fugw_batch( + a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals, divergence=divergence + ) loss_fugw_sample = loss_fugw_samples_batch( - a, a, C1, C2, X, Y, T, alpha=alpha, reg_marginals=reg_marginals + a, + a, + C1, + C2, + X, + Y, + T, + alpha=alpha, + reg_marginals=reg_marginals, + divergence=divergence, + metric_linear=metric_linear, + metric_quadratic=metric_quadratic, + logits=logits, ) loss_fugw_list = loss_fugw_batch( - a, a, L, M, T, alpha=alpha_list, reg_marginals=reg_marginals_list + a, + a, + L, + M, + T, + alpha=alpha_list, + reg_marginals=reg_marginals_list, + divergence=divergence, ) loss_fugw_sample_list = loss_fugw_samples_batch( - a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + a, + a, + C1, + C2, + X, + Y, + T, + alpha=alpha_list, + reg_marginals=reg_marginals_list, + divergence=divergence, + metric_linear=metric_linear, + metric_quadratic=metric_quadratic, + logits=logits, ) assert np.isfinite(loss_fugw).all() assert np.isfinite(loss_fugw_sample).all() assert np.isfinite(loss_fugw_list).all() assert np.isfinite(loss_fugw_sample_list).all() + np.testing.assert_allclose(loss_fugw, loss_fugw_list) + np.testing.assert_allclose(loss_fugw_sample, loss_fugw_sample_list) # check that invalid alpha shape raise an error alpha = rng.rand(batchsize + 1) with pytest.raises(ValueError): loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + with pytest.raises(ValueError): loss_fugw_samples_batch( - a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + a, + a, + C1, + C2, + X, + Y, + T, + alpha=alpha, + reg_marginals=reg_marginals, + divergence=divergence, + metric_linear=metric_linear, + metric_quadratic=metric_quadratic, + logits=logits, ) # check that invalid rho shape raise an error @@ -199,8 +268,21 @@ def test_fugw_loss(): reg_marginals = rng.rand(batchsize + 1) with pytest.raises(ValueError): loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) + with pytest.raises(ValueError): loss_fugw_samples_batch( - a, a, C1, C2, X, Y, T, alpha=alpha_list, reg_marginals=reg_marginals_list + a, + a, + C1, + C2, + X, + Y, + T, + alpha=alpha, + reg_marginals=reg_marginals, + divergence=divergence, + metric_linear=metric_linear, + metric_quadratic=metric_quadratic, + logits=logits, ) From 8f900d74bcb158fd2d7e1fec12a33c5746a12eff Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 27 May 2026 13:49:39 +0200 Subject: [PATCH 06/15] update RELEASES --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index fcafeaa85..9f504ec7c 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -19,7 +19,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver - 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 batch FUGW loss to `ot.batch` (PR #775) +- Add batch FUGW loss to `ot.batch` and fix issues in some default parameters in the batch module (PR #775) #### Closed issues From 82f361b95d8c763537dcf89aba22ea0af674d1b6 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 27 May 2026 14:13:11 +0200 Subject: [PATCH 07/15] add a test --- test/batch/test_solve_batch.py | 5 ++--- test/batch/test_solve_gromov_batch.py | 24 +++++++++++++++++++++++- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 566116c0a..17d459a43 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -1,9 +1,8 @@ """Tests for module batch""" # Author: Remi Flamary -# Kilian Fatras -# Quang Huy Tran -# Eduardo Fernandes Montesuma +# Paul Krzakala +# Sonia Mazelet # # License: MIT License diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index a98e8d5ba..ddc2b5268 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -1,8 +1,10 @@ """Tests for module batch""" # Author: Remi Flamary +# Paul Krzakala # Sonia Mazelet + # # License: MIT License @@ -18,7 +20,13 @@ import pytest from itertools import product from ot.backend import torch -from ot.batch._quadratic import tensor_batch, loss_fugw_batch, loss_fugw_samples_batch +from ot.batch._quadratic import ( + tensor_batch, + loss_fugw_batch, + loss_fugw_samples_batch, + div_to_product_batch, +) +from ot.gromov._utils import div_to_product def test_solve_gromov_batch(): @@ -311,3 +319,17 @@ def test_valid_fugw_loss_endpoints(): ) loss_gromov = loss_quadratic_batch(L, T, recompute_const=True) np.testing.assert_allclose(loss_fugw, loss_gromov, atol=1e-5) + + +def test_div_to_product(): + batchsize = 2 + n = 4 + batchsize = 1 + rng = np.random.RandomState(0) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + res_batch = div_to_product_batch( + T, a, a, T1=None, T2=None, divergence="kl", mass=True, nx=None + ) + res = div_to_product(T[0], a[0], a[0], divergence="kl", mass=True) + np.testing.assert_allclose(res_batch, res, atol=1e-5) From 72223bc8c861a2bb1457cd4ee31f77277594b60e Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 27 May 2026 14:40:08 +0200 Subject: [PATCH 08/15] fix bug --- examples/backends/plot_gradient_descent.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index 7ff944792..aee970c50 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -83,12 +83,8 @@ def plot_graph(x, C, color="C0", s=100): C2 = get_sbm(n2, nc2, ratio2, P2) # get 2d position for nodes -x1 = MDS( - metric="precomputed", random_state=0, n_init=1, init="classical_mds" -).fit_transform(1 - C1) -x2 = MDS( - metric="precomputed", random_state=0, n_init=1, init="classical_mds" -).fit_transform(1 - C2) +x1 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C1) +x2 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C2) colors1 = get_position_colors(x1) colors2 = get_position_colors(x2) From 214afac12981ea87ee308c9c1abc8ee0de7655dc Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 29 May 2026 11:25:37 +0200 Subject: [PATCH 09/15] change test --- examples/backends/plot_gradient_descent.py | 104 +++++++++++++-------- 1 file changed, 63 insertions(+), 41 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index aee970c50..ddb561ac5 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -1,10 +1,10 @@ # -*- coding: utf-8 -*- r""" =============================================================================== -Solve Fused Unbalanced Gromov Wasserstein with gradient descent +Solve Fused Unbalanced Gromov Wasserstein with Adam =============================================================================== -Since the FUGW loss is differentiable, it can be minimized with gradient descent. +Since the FUGW loss is differentiable, it can be minimized with first-order optimization. We show how to do this with the `loss_fugw_batch` function and compare the results with the dedicated FUGW solver `fused_unbalanced_gromov_wasserstein`. """ @@ -19,7 +19,6 @@ import numpy as np import matplotlib.pylab as pl import torch - import ot from ot.batch._quadratic import loss_fugw_batch, tensor_batch from ot.gromov import fused_unbalanced_gromov_wasserstein @@ -53,15 +52,6 @@ def get_sbm(n, nc, ratio, P): return C + C.T -def get_position_colors(x): - xmin = x.min(axis=0, keepdims=True) - xmax = x.max(axis=0, keepdims=True) - xnorm = (x - xmin) / np.maximum(xmax - xmin, 1e-15) - colors_x = pl.cm.viridis(xnorm[:, 0])[:, :3] - colors_y = pl.cm.viridis(xnorm[:, 1])[:, :3] - return np.sqrt(colors_x * colors_y) - - def plot_graph(x, C, color="C0", s=100): for j in range(C.shape[0]): for i in range(j): @@ -70,6 +60,19 @@ def plot_graph(x, C, color="C0", s=100): pl.scatter(x[:, 0], x[:, 1], c=color, s=s, zorder=10, edgecolors="k") +def get_sbm_labels(n, ratio): + nbpc = np.round(n * ratio).astype(int) + return np.concatenate( + [np.full(count, label, dtype=int) for label, count in enumerate(nbpc)] + ) + + +def get_noisy_one_hot(labels, n_classes, noise_level=0.1): + x = np.eye(n_classes)[labels] + x += noise_level * rng.randn(*x.shape) + return x + + n1 = 30 n2 = 20 nc1 = 3 @@ -81,29 +84,43 @@ def plot_graph(x, C, color="C0", s=100): P2 = np.array(0.6 * np.eye(2) + 0.05 * np.ones((2, 2))) C1 = get_sbm(n1, nc1, ratio1, P1) C2 = get_sbm(n2, nc2, ratio2, P2) - -# get 2d position for nodes -x1 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C1) -x2 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C2) - -colors1 = get_position_colors(x1) -colors2 = get_position_colors(x2) +labels1 = get_sbm_labels(n1, ratio1) +labels2 = get_sbm_labels(n2, ratio2) + +# Use noisy one-hot encodings of the SBM classes as node features. +feature_dim = max(nc1, nc2) +x1 = get_noisy_one_hot(labels1, feature_dim) +x2 = get_noisy_one_hot(labels2, feature_dim) +all_features = np.vstack([x1, x2]) +feature_min = all_features[:, :3].min(axis=0, keepdims=True) +feature_max = all_features[:, :3].max(axis=0, keepdims=True) + +# get 2d positions for visualization +pos1 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C1) +pos2 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C2) + +colors1 = np.clip( + (x1 - feature_min) / np.maximum(feature_max - feature_min, 1e-15), 0.0, 1.0 +) +colors2 = np.clip( + (x2 - feature_min) / np.maximum(feature_max - feature_min, 1e-15), 0.0, 1.0 +) pl.figure(1, (10, 5)) pl.clf() pl.subplot(1, 2, 1) -plot_graph(x1, C1, color=colors1) +plot_graph(pos1, C1, color=colors1) pl.title("SBM source graph") pl.axis("off") pl.subplot(1, 2, 2) -plot_graph(x2, C2, color=colors2) +plot_graph(pos2, C2, color=colors2) pl.title("SBM target graph") _ = pl.axis("off") # %% -# Solve FUGW with gradient descent +# Solve FUGW with Adam # ---------------- # Even though `loss_fugw_batch` supports batches of problems, we use a @@ -121,22 +138,31 @@ def plot_graph(x, C, color="C0", s=100): M_torch = torch.tensor(M[None, :, :]) L = tensor_batch(a_torch, b_torch, C1_torch, C2_torch, loss="sqeuclidean") -alpha = 0.7 +alpha = 0.5 reg_marginals = 1.0 -lr = 1e-5 -nb_iter_max = 300 +lr = 1e-2 +nb_iter_max = 1000 -T_torch = (a_torch[:, :, None] * b_torch[:, None, :]).clone().requires_grad_(True) +T0_torch = torch.tensor( + rng.rand(a_torch.shape[0], a_torch.shape[1], b_torch.shape[1]), + dtype=a_torch.dtype, +) +T0_torch /= T0_torch.sum(dim=(1, 2), keepdim=True) +T_torch = torch.log(torch.expm1(T0_torch)).clone().requires_grad_(True) +optimizer = torch.optim.Adam([T_torch], lr=lr) loss_iter = [] mass_iter = [] for i in range(nb_iter_max): + optimizer.zero_grad() + # Positive transport plan parameterized as log(1 + exp(T)). + plan_torch = torch.nn.functional.softplus(T_torch) loss = loss_fugw_batch( a_torch, b_torch, L, M_torch, - T_torch, + plan_torch, alpha=alpha, reg_marginals=reg_marginals, divergence="kl", @@ -144,22 +170,18 @@ def plot_graph(x, C, color="C0", s=100): )[0] loss_iter.append(float(loss.detach())) - mass_iter.append(float(T_torch.detach().sum())) + mass_iter.append(float(plan_torch.detach().sum())) loss.backward() + optimizer.step() - with torch.no_grad(): - T_torch -= lr * T_torch.grad - T_torch.clamp_(min=1e-12) - T_torch.grad.zero_() - -T_gd = T_torch.detach().cpu().numpy()[0] +T_adam = torch.nn.functional.softplus(T_torch).detach().cpu().numpy()[0] pl.figure(2, (10, 4)) pl.clf() pl.subplot(1, 2, 1) pl.plot(loss_iter) pl.grid() -pl.title("FUGW loss along gradient descent") +pl.title("FUGW loss along iterations") pl.xlabel("Iterations") pl.subplot(1, 2, 2) pl.plot(mass_iter) @@ -173,10 +195,10 @@ def plot_graph(x, C, color="C0", s=100): # ------------------------------------- # # The dedicated solver uses a block coordinate descent scheme. We compare the -# coupling it returns with the coupling obtained by direct gradient descent on -# `loss_fugw_batch`. The FUGW loss is non convex so minimizing it directly with gradient descent does not +# coupling it returns with the coupling obtained by direct Adam minimization on +# `loss_fugw_batch`. The FUGW loss is non convex so minimizing it directly with Adam does not # necessarily give the same solution as the dedicated solver. By comparing the FUGW costs obtained by both methods, -# we find that the BCD solvers gives a better solution than gradient descent. +# we find that the BCD solver gives a better solution than direct minimization on this example. T_bcd, _, log = fused_unbalanced_gromov_wasserstein( C1, @@ -197,7 +219,7 @@ def plot_graph(x, C, color="C0", s=100): ) -print("Final batch FUGW loss (gradient descent):", loss_iter[-1]) +print("Final batch FUGW loss:", loss_iter[-1]) print("FUGW cost reported by the dedicated solver:", log["fugw_cost"]) @@ -211,8 +233,8 @@ def plot_graph(x, C, color="C0", s=100): pl.figure(3, (10, 4)) pl.clf() pl.subplot(1, 2, 1) -pl.imshow(T_gd, interpolation="nearest") -pl.title("Coupling from gradient descent") +pl.imshow(T_adam, interpolation="nearest") +pl.title("Coupling from direct minimization") pl.xlabel("Target nodes") pl.ylabel("Source nodes") pl.colorbar() From be772c122638d138e92573016f4ff7ed2959fc90 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 29 May 2026 11:46:30 +0200 Subject: [PATCH 10/15] fix example --- examples/backends/plot_gradient_descent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index ddb561ac5..f9e01e167 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -96,8 +96,8 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): feature_max = all_features[:, :3].max(axis=0, keepdims=True) # get 2d positions for visualization -pos1 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C1) -pos2 = MDS(metric="precomputed", random_state=0, n_init=1).fit_transform(1 - C2) +pos1 = MDS(dissimilarity="precomputed", random_state=0, n_init=1).fit_transform(1 - C1) +pos2 = MDS(dissimilarity="precomputed", random_state=0, n_init=1).fit_transform(1 - C2) colors1 = np.clip( (x1 - feature_min) / np.maximum(feature_max - feature_min, 1e-15), 0.0, 1.0 From 95fd5985044fe74d275076fe71cc64abb1ac258f Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Wed, 3 Jun 2026 14:25:49 +0200 Subject: [PATCH 11/15] fix example --- examples/backends/plot_gradient_descent.py | 95 ++++++++++++------- ot/batch/_quadratic.py | 102 ++++++++++----------- test/batch/test_solve_gromov_batch.py | 27 ++++-- 3 files changed, 127 insertions(+), 97 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index f9e01e167..efe02cbce 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -14,11 +14,12 @@ # # License: MIT License -# sphinx_gallery_thumbnail_number = 2 +# sphinx_gallery_thumbnail_number = 3 import numpy as np import matplotlib.pylab as pl import torch +from time import perf_counter import ot from ot.batch._quadratic import loss_fugw_batch, tensor_batch from ot.gromov import fused_unbalanced_gromov_wasserstein @@ -73,15 +74,15 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): return x -n1 = 30 -n2 = 20 +n1 = 15 +n2 = 10 nc1 = 3 nc2 = 2 ratio1 = np.array([0.33, 0.33, 0.33]) ratio2 = np.array([0.5, 0.5]) -P1 = np.array([[0.8, 0.08, 0.0], [0.08, 0.8, 0.08], [0.0, 0.08, 0.8]]) -P2 = np.array(0.6 * np.eye(2) + 0.05 * np.ones((2, 2))) +P1 = np.array([[0.8, 0.03, 0.0], [0.08, 0.8, 0.03], [0.0, 0.08, 0.8]]) +P2 = np.array(0.8 * np.eye(2) + 0.01 * np.ones((2, 2))) C1 = get_sbm(n1, nc1, ratio1, P1) C2 = get_sbm(n2, nc2, ratio2, P2) labels1 = get_sbm_labels(n1, ratio1) @@ -138,21 +139,25 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): M_torch = torch.tensor(M[None, :, :]) L = tensor_batch(a_torch, b_torch, C1_torch, C2_torch, loss="sqeuclidean") -alpha = 0.5 -reg_marginals = 1.0 -lr = 1e-2 -nb_iter_max = 1000 +alpha_batch = 0.5 +# `loss_fugw_batch` uses alpha as: alpha * quadratic + (1 - alpha) * linear +# while the dedicated solver uses alpha as the coefficient of the linear term. +alpha_bcd = (1 - alpha_batch) / alpha_batch -T0_torch = torch.tensor( - rng.rand(a_torch.shape[0], a_torch.shape[1], b_torch.shape[1]), - dtype=a_torch.dtype, -) -T0_torch /= T0_torch.sum(dim=(1, 2), keepdim=True) +reg_marginals_batch = 0.5 +reg_marginals_bcd = reg_marginals_batch / alpha_batch +lr = 5e-2 +nb_iter_max = 1500 +tol = 1e-7 + +T0_torch = a_torch[:, :, None] * b_torch[:, None, :] T_torch = torch.log(torch.expm1(T0_torch)).clone().requires_grad_(True) optimizer = torch.optim.Adam([T_torch], lr=lr) loss_iter = [] mass_iter = [] +previous_plan_torch = None +tic = perf_counter() for i in range(nb_iter_max): optimizer.zero_grad() # Positive transport plan parameterized as log(1 + exp(T)). @@ -163,16 +168,22 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): L, M_torch, plan_torch, - alpha=alpha, - reg_marginals=reg_marginals, + alpha=alpha_batch, + reg_marginals=reg_marginals_batch, divergence="kl", recompute_const=True, )[0] loss_iter.append(float(loss.detach())) mass_iter.append(float(plan_torch.detach().sum())) + if previous_plan_torch is not None: + err = float(torch.sum(torch.abs(plan_torch.detach() - previous_plan_torch))) + if err < tol: + break + previous_plan_torch = plan_torch.detach().clone() loss.backward() optimizer.step() +time_adam = perf_counter() - tic T_adam = torch.nn.functional.softplus(T_torch).detach().cpu().numpy()[0] @@ -194,53 +205,71 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): # Compare with the dedicated FUGW solver # ------------------------------------- # -# The dedicated solver uses a block coordinate descent scheme. We compare the -# coupling it returns with the coupling obtained by direct Adam minimization on -# `loss_fugw_batch`. The FUGW loss is non convex so minimizing it directly with Adam does not -# necessarily give the same solution as the dedicated solver. By comparing the FUGW costs obtained by both methods, -# we find that the BCD solver gives a better solution than direct minimization on this example. +# The dedicated solver uses a block coordinate descent (BCD) scheme. We compare +# the coupling it returns with the one obtained by direct Adam minimization of +# `loss_fugw_batch`. + + +def evaluate_batch_fugw_loss(plan): + plan_torch = torch.tensor(plan[None, :, :], dtype=M_torch.dtype) + loss = loss_fugw_batch( + a_torch, + b_torch, + L, + M_torch, + plan_torch, + alpha=alpha_batch, + reg_marginals=reg_marginals_batch, + divergence="kl", + recompute_const=True, + )[0] + return float(loss.detach()) + +tic = perf_counter() T_bcd, _, log = fused_unbalanced_gromov_wasserstein( C1, C2, wx=a, wy=b, - reg_marginals=reg_marginals, + reg_marginals=reg_marginals_bcd, divergence="kl", unbalanced_solver="mm", - alpha=alpha, + alpha=alpha_bcd, M=M, init_pi=np.outer(a, b), - max_iter=100, - tol=1e-7, + max_iter=200, + tol=tol, max_iter_ot=200, tol_ot=1e-7, log=True, ) +time_bcd = perf_counter() - tic - -print("Final batch FUGW loss:", loss_iter[-1]) -print("FUGW cost reported by the dedicated solver:", log["fugw_cost"]) +loss_adam_final = evaluate_batch_fugw_loss(T_adam) +loss_bcd_final = evaluate_batch_fugw_loss(T_bcd) # %% # Visualize the learned couplings # ------------------------------- -# We visualize the couplings obtained by both methods to compare them. -# The BCD solver gives sharper plans but both -# methods find couplings that match the structures of the graphs. +# We visualize the couplings obtained by both methods to compare them. On this example, both methods recover similar couplings, +# but direct minimization reaches a lower `loss_fugw_batch` value at the cost +# of a longer runtime. pl.figure(3, (10, 4)) pl.clf() pl.subplot(1, 2, 1) pl.imshow(T_adam, interpolation="nearest") -pl.title("Coupling from direct minimization") +pl.title( + f"Coupling from direct minimization\nloss={loss_adam_final:.3f}, time={time_adam:.2f}s" +) pl.xlabel("Target nodes") pl.ylabel("Source nodes") pl.colorbar() pl.subplot(1, 2, 2) pl.imshow(T_bcd, interpolation="nearest") -pl.title("Coupling from BCD solver") +pl.title(f"Coupling from BCD solver\nloss={loss_bcd_final:.3f}, time={time_bcd:.2f}s") pl.xlabel("Target nodes") pl.ylabel("Source nodes") _ = pl.colorbar() diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index 7ce0ca2fb..5b398fd80 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -153,85 +153,73 @@ def h2(C2): return compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=symmetric) -def div_to_product_batch( - T, a, b, T1=None, T2=None, divergence="kl", mass=True, nx=None -): - r"""Fast computation of the Bregman divergence between a batch of arbitrary measures and a product measures. +def div_between_product_batch(mu, nu, alpha, beta, divergence, nx=None): + r"""Fast computation of the Bregman divergence between batches of product measures. Only support for Kullback-Leibler and half-squared L2 divergences. - - For half-squared L2 divergence: + For half-squared L2 divergence: .. math:: - \frac{1}{2} || \pi - a \otimes b ||^2 - = \frac{1}{2} \Big[ \sum_{i, j} \pi_{ij}^2 + (\sum_i a_i^2) ( \sum_j b_j^2) - 2 \sum_{i, j} a_i \pi_{ij} b_j \Big] + \frac{1}{2} || \mu \otimes \nu, \alpha \otimes \beta ||^2 + = \frac{1}{2} \Big[ ||\alpha||^2 ||\beta||^2 + ||\mu||^2 ||\nu||^2 - 2 \langle \alpha, \mu \rangle \langle \beta, \nu \rangle \Big] - - For Kullback-Leibler divergence: + For Kullback-Leibler divergence: .. math:: - KL(\pi | a \otimes b) - = \langle \pi, \log \pi \rangle - \langle \pi_1, \log a \rangle - - \langle \pi_2, \log b \rangle - m(\pi) + m(a) m(b) + KL(\mu \otimes \nu, \alpha \otimes \beta) + = m(\mu) * KL(\nu, \beta) + m(\nu) * KL(\mu, \alpha) + (m(\mu) - m(\alpha)) * (m(\nu) - m(\beta)) - where : + where: - - :math:`\pi` is the (`dim_a`, `dim_b`) transport plan - - :math:`\pi_1` and :math:`\pi_2` are the marginal distributions - - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions + - :math:`\mu` and :math:`\alpha` are two measures having the same shape. + - :math:`\nu` and :math:`\beta` are two measures having the same shape. - :math:`m` denotes the mass of the measure Parameters ---------- - pi : array-like (B, n, m) - Transport plan for each problem in the batch - a : array-like (B,n) - Unnormalized histogram of dimension `n` for each problem in the batch - b : array-like (B,m) - Unnormalized histogram of dimension `m` for each problem in the batch - T1 : array-like (B, n), optional (default = None) - Marginal distribution with respect to the first dimension of the transport plan for each problem in the batch - Only used in case of Kullback-Leibler divergence. - T2 : array-like (B, m), optional (default = None) - Marginal distribution with respect to the second dimension of the transport plan for each problem in the batch - Only used in case of Kullback-Leibler divergence. + mu : array-like, shape (B, ...) + First factor of each product measure in the batch. + nu : array-like, shape (B, ...) + Second factor of each product measure in the batch. + alpha : array-like, shape (B, ...) + Reference factor with the same shape as `mu`. + beta : array-like, shape (B, ...) + Reference factor with the same shape as `nu`. divergence : string, default = "kl" Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) - mass : bool, optional. Default is False. - Only used in case of Kullback-Leibler divergence. - If False, calculate the relative entropy. - If True, calculate the Kullback-Leibler divergence. nx : backend, optional If let to its default value None, a backend test will be conducted. Returns - ------- - Bregman divergence between an arbitrary measure and a product measure for each problem in the batch. + ---------- + Bregman divergence between two product measures for each problem in the batch. """ - arr = [T, a, b, T1, T2] - if nx is None: - nx = get_backend(*arr, T1, T2) + nx = get_backend(mu, nu, alpha, beta) - if divergence == "kl": - if T1 is None: - T1 = nx.sum(T, 2) - if T2 is None: - T2 = nx.sum(T, 1) + axis_mu = tuple(range(1, mu.ndim)) if mu.ndim > 1 else 0 + axis_nu = tuple(range(1, nu.ndim)) if nu.ndim > 1 else 0 + axis_alpha = tuple(range(1, alpha.ndim)) if alpha.ndim > 1 else 0 + axis_beta = tuple(range(1, beta.ndim)) if beta.ndim > 1 else 0 if divergence == "kl": + m_mu = nx.sum(mu, axis=axis_mu) + m_nu = nx.sum(nu, axis=axis_nu) + m_alpha = nx.sum(alpha, axis=axis_alpha) + m_beta = nx.sum(beta, axis=axis_beta) + const = (m_mu - m_alpha) * (m_nu - m_beta) res = ( - nx.sum((T * nx.log(T + 1.0 * (T == 0))), (1, 2)) - - nx.sum(T1 * nx.log(a), 1) - - nx.sum(T2 * nx.log(b), 1) + m_nu * nx.kl_div(mu, alpha, mass=True, axis=axis_mu) + + m_mu * nx.kl_div(nu, beta, mass=True, axis=axis_nu) + + const ) - if mass: - res = res - nx.sum(T1, 1) + nx.sum(a, 1) * nx.sum(b, 1) elif divergence == "l2": res = ( - nx.sum(T**2, (1, 2)) - + nx.sum(a**2, 1) * nx.sum(b**2, 1) - - 2 * nx.sum((a * (T @ b[:, :, None]).squeeze(-1)), 1) + nx.sum(alpha**2, axis=axis_alpha) * nx.sum(beta**2, axis=axis_beta) + - 2 * nx.sum(alpha * mu, axis=axis_mu) * nx.sum(beta * nu, axis=axis_nu) + + nx.sum(mu**2, axis=axis_mu) * nx.sum(nu**2, axis=axis_nu) ) / 2 return res @@ -424,12 +412,14 @@ def loss_fugw_batch( linear = loss_linear_batch(M, T, nx=nx) - unbalanced = div_to_product_batch( - T, + T1 = nx.sum(T, 2) + T2 = nx.sum(T, 1) + unbalanced = div_between_product_batch( + T1, + T2, a, b, divergence=divergence, - mass=True, nx=nx, ) @@ -535,12 +525,14 @@ def loss_fugw_samples_batch( linear = loss_linear_samples_batch(X, Y, T, metric=metric_linear) - unbalanced = div_to_product_batch( - T, + T1 = nx.sum(T, 2) + T2 = nx.sum(T, 1) + unbalanced = div_between_product_batch( + T1, + T2, a, b, divergence=divergence, - mass=True, nx=nx, ) diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index ddc2b5268..cabb9219b 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -24,9 +24,9 @@ tensor_batch, loss_fugw_batch, loss_fugw_samples_batch, - div_to_product_batch, + div_between_product_batch, ) -from ot.gromov._utils import div_to_product +from ot.gromov._utils import div_between_product def test_solve_gromov_batch(): @@ -321,15 +321,24 @@ def test_valid_fugw_loss_endpoints(): np.testing.assert_allclose(loss_fugw, loss_gromov, atol=1e-5) -def test_div_to_product(): +@pytest.mark.parametrize("divergence", ["kl", "l2"]) +def test_div_between_product(divergence): batchsize = 2 n = 4 - batchsize = 1 + m = 3 rng = np.random.RandomState(0) - a = np.ones((batchsize, n)) - T = rng.rand(batchsize, n, n) - res_batch = div_to_product_batch( - T, a, a, T1=None, T2=None, divergence="kl", mass=True, nx=None + mu = rng.rand(batchsize, n) + nu = rng.rand(batchsize, m) + alpha = rng.rand(batchsize, n) + beta = rng.rand(batchsize, m) + + res_batch = div_between_product_batch( + mu, nu, alpha, beta, divergence=divergence, nx=None + ) + res = np.array( + [ + div_between_product(mu[i], nu[i], alpha[i], beta[i], divergence) + for i in range(batchsize) + ] ) - res = div_to_product(T[0], a[0], a[0], divergence="kl", mass=True) np.testing.assert_allclose(res_batch, res, atol=1e-5) From dcd7b304b8114af0a2c158dee540bb17ae29d6fd Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:01:25 +0200 Subject: [PATCH 12/15] refactor my functions --- examples/backends/plot_gradient_descent.py | 64 ++--- ot/batch/_quadratic.py | 270 ++++++--------------- test/batch/test_solve_gromov_batch.py | 184 +++++++------- 3 files changed, 217 insertions(+), 301 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index efe02cbce..0db8cfbd6 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -21,7 +21,7 @@ import torch from time import perf_counter import ot -from ot.batch._quadratic import loss_fugw_batch, tensor_batch +from ot.batch._quadratic import loss_quadratic_samples_batch, tensor_batch from ot.gromov import fused_unbalanced_gromov_wasserstein from sklearn.manifold import MDS @@ -145,7 +145,7 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): alpha_bcd = (1 - alpha_batch) / alpha_batch reg_marginals_batch = 0.5 -reg_marginals_bcd = reg_marginals_batch / alpha_batch +reg_marginals_bcd = reg_marginals_batch / (2 * alpha_batch) lr = 5e-2 nb_iter_max = 1500 tol = 1e-7 @@ -162,15 +162,16 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): optimizer.zero_grad() # Positive transport plan parameterized as log(1 + exp(T)). plan_torch = torch.nn.functional.softplus(T_torch) - loss = loss_fugw_batch( + loss = loss_quadratic_samples_batch( a_torch, b_torch, - L, - M_torch, + C1_torch, + C2_torch, plan_torch, + M_torch, alpha=alpha_batch, - reg_marginals=reg_marginals_batch, - divergence="kl", + unbalanced=reg_marginals_batch, + unbalanced_type="kl", recompute_const=True, )[0] @@ -187,19 +188,6 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): T_adam = torch.nn.functional.softplus(T_torch).detach().cpu().numpy()[0] -pl.figure(2, (10, 4)) -pl.clf() -pl.subplot(1, 2, 1) -pl.plot(loss_iter) -pl.grid() -pl.title("FUGW loss along iterations") -pl.xlabel("Iterations") -pl.subplot(1, 2, 2) -pl.plot(mass_iter) -pl.grid() -pl.title("Transport mass") -_ = pl.xlabel("Iterations") - # %% # Compare with the dedicated FUGW solver @@ -212,15 +200,16 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): def evaluate_batch_fugw_loss(plan): plan_torch = torch.tensor(plan[None, :, :], dtype=M_torch.dtype) - loss = loss_fugw_batch( + loss = loss_quadratic_samples_batch( a_torch, b_torch, - L, - M_torch, + C1_torch, + C2_torch, plan_torch, + M_torch, alpha=alpha_batch, - reg_marginals=reg_marginals_batch, - divergence="kl", + unbalanced=reg_marginals_batch, + unbalanced_type="kl", recompute_const=True, )[0] return float(loss.detach()) @@ -248,6 +237,25 @@ def evaluate_batch_fugw_loss(plan): loss_adam_final = evaluate_batch_fugw_loss(T_adam) loss_bcd_final = evaluate_batch_fugw_loss(T_bcd) +print(log["fugw_cost"]) +mass_bcd = T_bcd.sum() + +pl.figure(2, (10, 4)) +pl.clf() +pl.subplot(1, 2, 1) +pl.plot(loss_iter, label="Adam") +pl.axhline(loss_bcd_final, color="C1", linestyle="--", label="BCD solver") +pl.grid() +pl.title("FUGW loss along iterations") +pl.xlabel("Iterations") +pl.legend() +pl.subplot(1, 2, 2) +pl.plot(mass_iter, label="Adam") +pl.axhline(mass_bcd, color="C1", linestyle="--", label="BCD solver") +pl.grid() +pl.title("Transport mass") +pl.xlabel("Iterations") +_ = pl.legend() # %% @@ -257,10 +265,12 @@ def evaluate_batch_fugw_loss(plan): # but direct minimization reaches a lower `loss_fugw_batch` value at the cost # of a longer runtime. +vmin = min(T_adam.min(), T_bcd.min()) +vmax = max(T_adam.max(), T_bcd.max()) pl.figure(3, (10, 4)) pl.clf() pl.subplot(1, 2, 1) -pl.imshow(T_adam, interpolation="nearest") +pl.imshow(T_adam, interpolation="nearest", cmap="Blues", vmin=vmin, vmax=vmax) pl.title( f"Coupling from direct minimization\nloss={loss_adam_final:.3f}, time={time_adam:.2f}s" ) @@ -268,7 +278,7 @@ def evaluate_batch_fugw_loss(plan): pl.ylabel("Source nodes") pl.colorbar() pl.subplot(1, 2, 2) -pl.imshow(T_bcd, interpolation="nearest") +pl.imshow(T_bcd, interpolation="nearest", cmap="Blues", vmin=vmin, vmax=vmax) pl.title(f"Coupling from BCD solver\nloss={loss_bcd_final:.3f}, time={time_bcd:.2f}s") pl.xlabel("Target nodes") pl.ylabel("Source nodes") diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index 5b398fd80..437f2eb60 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -10,7 +10,7 @@ from ..utils import OTResult from ot.backend import get_backend -from ot.batch._linear import loss_linear_batch, loss_linear_samples_batch +from ot.batch._linear import loss_linear_batch from ot.batch._utils import bmv, bop, bregman_log_projection_batch from ot.utils import list_to_array @@ -106,7 +106,9 @@ def tensor_batch( if nx is None: nx = get_backend(C1) - if loss == "sqeuclidean": + loss = loss.lower() + + if loss == "sqeuclidean" or loss == "l2": def f1(C1): if C1.ndim == 4: @@ -277,6 +279,10 @@ def loss_quadratic_samples_batch( C1, C2, T, + M=None, + alpha=None, + unbalanced=None, + unbalanced_type="kl", loss="sqeuclidean", symmetric=True, nx=None, @@ -298,15 +304,33 @@ def loss_quadratic_samples_batch( Target cost matrices. T : array-like, shape (B, n, m) Transport plan. + M : array-like, shape (B, n, m) + Cost matrix between features across domains (default is None). + alpha : float, array-like or list (B,) optional + Weight the quadratic term (alpha*Gromov) and the linear term + ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. Not used for + Gromov problem (when M is not provided). By default ``alpha=None`` + corresponds to ``alpha=1`` for Gromov problem (``M==None``) and + ``alpha=0.5`` for Fused Gromov-Wasserstein problem (``M!=None``). + If alpha is a scalar, it is used for all problems in the batch. + unbalanced : float array-like or list(B,) optional + Unbalanced penalization weight :math:`\lambda_u`. If unbalanced is a scalar, it is used for all problems in the batch. + unbalanced_type : string, optional + Type of unbalanced penalization function, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) loss : str, optional Loss function to use. Supported values: 'sqeuclidean', 'kl'. Default is 'sqeuclidean'. - recompute_const : bool, optional - Whether to recompute the constant term. Default is False. This should be set to True if T does not satisfy the marginal constraints. symmetric : bool, optional Whether to use symmetric version. Default is True. nx : module, optional Backend to use. Default is None. + logits : bool, optional + For KL divergence, whether inputs are logits (unnormalized log probabilities). + If True, inputs are treated as logits. Default is None. + recompute_const : bool, optional + Whether to recompute the constant term. Default is False. This should be set to True if T does not satisfy the marginal constraints. + Will be set to True if unbalanced is not None. + Examples -------- @@ -328,215 +352,81 @@ def loss_quadratic_samples_batch( ot.batch.tensor_batch : From computing the cost tensor L. ot.batch.solve_gromov_batch : For finding the optimal transport plan T. """ + if nx is None: + nx = get_backend(T) + if isinstance(loss, str): L = tensor_batch( a, b, C1, C2, symmetric=symmetric, nx=nx, loss=loss, logits=logits ) else: raise ValueError(f"Unknown loss function: {loss}") - return loss_quadratic_batch( - L, T, recompute_const=recompute_const, symmetric=symmetric, nx=nx - ) - - -def loss_fugw_batch( - a, - b, - L, - M, - T, - alpha=0.5, - reg_marginals=1, - symmetric=True, - divergence="kl", - recompute_const=True, - nx=None, -): - r""" - Computes the fused unbalanced gromov-wasserstein cost given a cost tensor (Gromov term), a cost matrix between features across domains (linear term) and a transport plan. Batched version. - - Parameters - ---------- - a : array-like, shape (B, n) - Source distributions. - b : array-like, shape (B, m) - Target distributions. - L : dict - Cost tensor as returned by `tensor_batch`. - M : array-like, shape (B, n, m) - Cost matrix between features across domains. - T : array-like, shape (B, n, m) - Transport plan. - alpha : float, array-like or list (B,) optional - Weight the quadratic term (alpha*Gromov) and the linear term - ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. If alpha - a scalar it is used for all problems in the batch. - reg_marginals : float array-like or list(B,) optional - Marginal relaxation terms. If rho is - a scalar it is used for all problems in the batch. - symmetric : bool, optional - Whether to use symmetric version. Default is True. - divergence : string, default = "kl" - Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) - recompute_const : bool, optional - Whether to recompute the constant term. Default is True. This should be set to True if T does not satisfy the marginal constraints. - nx : module, optional - Backend to use. Default is None. - """ - if nx is None: - nx = get_backend(T) - - B = T.shape[0] - if isinstance(alpha, list): - alpha = list_to_array(alpha, nx=nx) - - if isinstance(reg_marginals, list): - reg_marginals = list_to_array(reg_marginals, nx=nx) - - if hasattr(alpha, "ndim") and alpha.ndim > 0: - if alpha.ndim != 1 or alpha.shape[0] != B: - raise ValueError( - f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" - ) - - if hasattr(reg_marginals, "ndim") and reg_marginals.ndim > 0: - if reg_marginals.ndim != 1 or reg_marginals.shape[0] != B: - raise ValueError( - f"If reg_marginals is not a scalar, it must have shape ({B},), got {reg_marginals.shape}" - ) + if unbalanced is not None: + recompute_const = True quadratic = loss_quadratic_batch( L, T, recompute_const=recompute_const, symmetric=symmetric, nx=nx ) - linear = loss_linear_batch(M, T, nx=nx) - - T1 = nx.sum(T, 2) - T2 = nx.sum(T, 1) - unbalanced = div_between_product_batch( - T1, - T2, - a, - b, - divergence=divergence, - nx=nx, - ) - - return (1 - alpha) * linear + alpha * quadratic + reg_marginals * unbalanced - - -def loss_fugw_samples_batch( - a, - b, - C1, - C2, - X, - Y, - T, - alpha=0.5, - reg_marginals=1, - symmetric=True, - divergence="kl", - recompute_const=True, - metric_linear="sqeuclidean", - metric_quadratic="sqeuclidean", - logits=None, - nx=None, -): - r""" - Computes the fused unbalanced gromov-wasserstein cost given a cost tensor (quadratic term), a cost matrix between features across domains (linear term) and a transport plan. Batched version. - - Parameters - ---------- - a : array-like, shape (B, n) - Source distributions. - b : array-like, shape (B, m) - Target distributions. - C1 : array-like, shape (B, n, n) or (B, n, n, d) - Source cost matrices for the quadratic term. - C2 : array-like, shape (B, m, m) or (B, n, n, d) - Target cost matrices for the quadratic term. - X : array-like, shape (B, n, d) - Samples from source distribution for the linear term - Y : array-like, shape (B, m, d) - Samples from target distribution for the linear term - T : array-like, shape (B, n, m) - Transport plan. - alpha : float or array-like or list(B,) optional - Weight the quadratic term (alpha*Gromov) and the linear term - ((1-alpha)*Wass) in the Fused Gromov-Wasserstein problem. If alpha - a scalar it is used for all problems in the batch. - reg_marginals : float or array-like or list(B,) optional - Marginal relaxation terms. If rho is - a scalar it is used for all problems in the batch. - symmetric : bool, optional - Whether to use symmetric version. Default is True. - divergence : string, default = "kl" - Bregman divergence, either "kl" (Kullback-Leibler divergence) or "l2" (half-squared L2 divergence) - recompute_const : bool, optional - Whether to recompute the constant term. Default is True. This should be set to True if T does not satisfy the marginal constraints. - metric_linear : str, optional - Metric for the linear term, 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' - metric_quadratic : str, optional - Metric to use for the quadratic term. Supported values: 'sqeuclidean', 'kl'. - Default is 'sqeuclidean'. - logits : bool, optional - For KL divergence, whether inputs are logits (unnormalized log probabilities). - If True, inputs are treated as logits. Default is None. - nx : module, optional - Backend to use. Default is None. - """ - if nx is None: - nx = get_backend(T) + if unbalanced is None and M is None: + return quadratic B = T.shape[0] - if isinstance(alpha, list): - alpha = list_to_array(alpha, nx=nx) - - if isinstance(reg_marginals, list): - reg_marginals = list_to_array(reg_marginals, nx=nx) - - if hasattr(alpha, "ndim") and alpha.ndim > 0: - if alpha.ndim != 1 or alpha.shape[0] != B: + if unbalanced is not None: + if unbalanced_type is None: raise ValueError( - f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" + "unbalanced_type must be specified if unbalanced is not None" ) - if hasattr(reg_marginals, "ndim") and reg_marginals.ndim > 0: - if reg_marginals.ndim != 1 or reg_marginals.shape[0] != B: + unbalanced_type = unbalanced_type.lower() + + if unbalanced_type not in ["kl", "l2"]: raise ValueError( - f"If reg_marginals is not a scalar, it must have shape ({B},), got {reg_marginals.shape}" + f"Unknown unbalanced_type: {unbalanced_type}, expected 'kl' or 'l2'" ) - quadratic = loss_quadratic_samples_batch( - a, - b, - C1, - C2, - T, - loss=metric_quadratic, - symmetric=symmetric, - nx=nx, - logits=logits, - recompute_const=recompute_const, - ) + if isinstance(unbalanced, list): + unbalanced = list_to_array(unbalanced, nx=nx) + + if hasattr(unbalanced, "ndim") and unbalanced.ndim > 0: + if unbalanced.ndim != 1 or unbalanced.shape[0] != B: + raise ValueError( + f"If reg_marginals is not a scalar, it must have shape ({B},), got {unbalanced.shape}" + ) + + T1 = nx.sum(T, 2) + T2 = nx.sum(T, 1) + unbalanced_term = div_between_product_batch( + T1, + T2, + a, + b, + divergence=unbalanced_type, + nx=nx, + ) - linear = loss_linear_samples_batch(X, Y, T, metric=metric_linear) - - T1 = nx.sum(T, 2) - T2 = nx.sum(T, 1) - unbalanced = div_between_product_batch( - T1, - T2, - a, - b, - divergence=divergence, - nx=nx, - ) + if M is not None: + if alpha is None: + alpha = 0.5 + if isinstance(alpha, list): + alpha = list_to_array(alpha, nx=nx) + if hasattr(alpha, "ndim") and alpha.ndim > 0: + if alpha.ndim != 1 or alpha.shape[0] != B: + raise ValueError( + f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" + ) + linear = loss_linear_batch(M, T, nx=nx) - return (1 - alpha) * linear + alpha * quadratic + reg_marginals * unbalanced + if M is not None and unbalanced is not None: + return (1 - alpha) * linear + alpha * quadratic + unbalanced * unbalanced_term + + elif M is not None and unbalanced is None: + return (1 - alpha) * linear + alpha * quadratic + + else: + return quadratic + unbalanced * unbalanced_term def solve_gromov_batch( diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index cabb9219b..b69a94ed2 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -22,8 +22,6 @@ from ot.backend import torch from ot.batch._quadratic import ( tensor_batch, - loss_fugw_batch, - loss_fugw_samples_batch, div_between_product_batch, ) from ot.gromov._utils import div_between_product @@ -148,12 +146,9 @@ def test_backend(nx): solve_gromov_batch(C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False) -@pytest.mark.parametrize("divergence", ["kl", "l2"]) -@pytest.mark.parametrize( - "metric_linear", ["sqeuclidean", "euclidean", "minkowski", "kl"] -) -@pytest.mark.parametrize("metric_quadratic", ["sqeuclidean", "kl"]) -def test_fugw_loss(divergence, metric_linear, metric_quadratic): +@pytest.mark.parametrize("unbalanced_type", ["kl", "l2"]) +@pytest.mark.parametrize("loss", ["sqeuclidean", "kl"]) +def test_fugw_loss(unbalanced_type, loss): """Check that loss_fugw_batch and loss_fugw_samples_batch run without error.""" batchsize = 2 n = 4 @@ -161,37 +156,39 @@ def test_fugw_loss(divergence, metric_linear, metric_quadratic): rng = np.random.RandomState(0) C1 = rng.rand(batchsize, n, n, d) C2 = rng.rand(batchsize, n, n, d) - X = rng.rand(batchsize, n, d) - Y = rng.rand(batchsize, n, d) M = rng.rand(batchsize, n, n) a = np.ones((batchsize, n)) - reg_marginals = 0 T = rng.rand(batchsize, n, n) - L = tensor_batch(a=a, b=a, C1=C1, C2=C2, loss="sqeuclidean") alpha = rng.rand() reg_marginals = rng.rand() - logits = False if metric_quadratic == "kl" else None - loss_fugw = loss_fugw_batch( - a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals, divergence=divergence + loss_fugw = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=(loss == "kl"), ) - loss_fugw_sample = loss_fugw_samples_batch( + loss_fugw_unbalanced_only = loss_quadratic_samples_batch( a, a, C1, C2, - X, - Y, T, + M=None, alpha=alpha, - reg_marginals=reg_marginals, - divergence=divergence, - metric_linear=metric_linear, - metric_quadratic=metric_quadratic, - logits=logits, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) - assert np.isfinite(loss_fugw).all() - assert np.isfinite(loss_fugw_sample).all() + assert np.isfinite(loss_fugw_unbalanced_only).all() # check that alpha and reg_marginals can be passed as lists or arrays of shape (batchsize,) alpha = rng.rand(batchsize) @@ -199,102 +196,92 @@ def test_fugw_loss(divergence, metric_linear, metric_quadratic): alpha_list = alpha.tolist() reg_marginals_list = reg_marginals.tolist() - loss_fugw = loss_fugw_batch( - a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals, divergence=divergence - ) - loss_fugw_sample = loss_fugw_samples_batch( + loss_fugw = loss_quadratic_samples_batch( a, a, C1, C2, - X, - Y, T, - alpha=alpha, - reg_marginals=reg_marginals, - divergence=divergence, - metric_linear=metric_linear, - metric_quadratic=metric_quadratic, - logits=logits, - ) - loss_fugw_list = loss_fugw_batch( - a, - a, - L, M, - T, - alpha=alpha_list, - reg_marginals=reg_marginals_list, - divergence=divergence, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) - loss_fugw_sample_list = loss_fugw_samples_batch( + loss_fugw_list = loss_quadratic_samples_batch( a, a, C1, C2, - X, - Y, T, + M, alpha=alpha_list, - reg_marginals=reg_marginals_list, - divergence=divergence, - metric_linear=metric_linear, - metric_quadratic=metric_quadratic, - logits=logits, + unbalanced=reg_marginals_list, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) assert np.isfinite(loss_fugw).all() - assert np.isfinite(loss_fugw_sample).all() assert np.isfinite(loss_fugw_list).all() - assert np.isfinite(loss_fugw_sample_list).all() np.testing.assert_allclose(loss_fugw, loss_fugw_list) - np.testing.assert_allclose(loss_fugw_sample, loss_fugw_sample_list) + + # test that invalid unbalanced_type raise an error + with pytest.raises(ValueError): + loss_fugw = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type="test", + loss=loss, + logits=False, + ) # check that invalid alpha shape raise an error alpha = rng.rand(batchsize + 1) with pytest.raises(ValueError): - loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) - with pytest.raises(ValueError): - loss_fugw_samples_batch( + loss_quadratic_samples_batch( a, a, C1, C2, - X, - Y, T, + M, alpha=alpha, - reg_marginals=reg_marginals, - divergence=divergence, - metric_linear=metric_linear, - metric_quadratic=metric_quadratic, - logits=logits, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) # check that invalid rho shape raise an error alpha = rng.rand(batchsize) reg_marginals = rng.rand(batchsize + 1) with pytest.raises(ValueError): - loss_fugw_batch(a, a, L, M, T, alpha=alpha, reg_marginals=reg_marginals) - with pytest.raises(ValueError): - loss_fugw_samples_batch( + loss_quadratic_samples_batch( a, a, C1, C2, - X, - Y, T, + M, alpha=alpha, - reg_marginals=reg_marginals, - divergence=divergence, - metric_linear=metric_linear, - metric_quadratic=metric_quadratic, - logits=logits, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) -def test_valid_fugw_loss_endpoints(): +@pytest.mark.parametrize("unbalanced_type", ["kl", "l2"]) +@pytest.mark.parametrize("loss", ["sqeuclidean", "kl"]) +def test_valid_fugw_loss_endpoints(unbalanced_type, loss): """Check that loss_fugw_batch gives the same results as solve_gromov_batch and solve_linear_batch for alpha=0 and alpha=1.""" batchsize = 2 n = 4 @@ -306,18 +293,47 @@ def test_valid_fugw_loss_endpoints(): a = np.ones((batchsize, n)) reg_marginals = 0 T = rng.rand(batchsize, n, n) - L = tensor_batch(a=a, b=a, C1=C1, C2=C2, loss="sqeuclidean") - loss_fugw = loss_fugw_batch( - a, a, L, M, T, alpha=0.0, divergence="l2", reg_marginals=reg_marginals + loss_fugw = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M, + alpha=0.0, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) loss_linear = loss_linear_batch(M, T) np.testing.assert_allclose(loss_fugw, loss_linear, atol=1e-5) - loss_fugw = loss_fugw_batch( - a, a, L, M, T, alpha=1.0, divergence="l2", reg_marginals=reg_marginals + loss_fugw = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M, + alpha=1.0, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, + ) + loss_gromov = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, ) - loss_gromov = loss_quadratic_batch(L, T, recompute_const=True) np.testing.assert_allclose(loss_fugw, loss_gromov, atol=1e-5) From 822a95732a39ea4635df9ad3c79bb872c0c25ccb Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Tue, 9 Jun 2026 14:25:45 +0200 Subject: [PATCH 13/15] increase test coverage --- ot/batch/_quadratic.py | 2 +- test/batch/test_solve_gromov_batch.py | 47 +++++++++++++++++++++++++-- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index 437f2eb60..c4e3c524f 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -355,7 +355,7 @@ def loss_quadratic_samples_batch( if nx is None: nx = get_backend(T) - if isinstance(loss, str): + if isinstance(loss, str) and loss in ["sqeuclidean", "kl", "l2"]: L = tensor_batch( a, b, C1, C2, symmetric=symmetric, nx=nx, loss=loss, logits=logits ) diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index b69a94ed2..c242f83f7 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -173,7 +173,7 @@ def test_fugw_loss(unbalanced_type, loss): unbalanced=reg_marginals, unbalanced_type=unbalanced_type, loss=loss, - logits=(loss == "kl"), + logits=False, ) loss_fugw_unbalanced_only = loss_quadratic_samples_batch( a, @@ -188,7 +188,34 @@ def test_fugw_loss(unbalanced_type, loss): loss=loss, logits=False, ) + loss_fugw_no_alpha = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M=None, + alpha=None, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, + ) + loss_fugw_no_unbalanced = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M=M, + alpha=alpha, + loss=loss, + logits=False, + ) assert np.isfinite(loss_fugw_unbalanced_only).all() + assert np.isfinite(loss_fugw).all() + assert np.isfinite(loss_fugw_no_alpha).all() + assert np.isfinite(loss_fugw_no_unbalanced).all() # check that alpha and reg_marginals can be passed as lists or arrays of shape (batchsize,) alpha = rng.rand(batchsize) @@ -227,7 +254,23 @@ def test_fugw_loss(unbalanced_type, loss): assert np.isfinite(loss_fugw_list).all() np.testing.assert_allclose(loss_fugw, loss_fugw_list) - # test that invalid unbalanced_type raise an error + # check that invalid loss raise an error + with pytest.raises(ValueError): + loss_fugw = loss_quadratic_samples_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss="test", + logits=False, + ) + + # check that invalid loss raise an error with pytest.raises(ValueError): loss_fugw = loss_quadratic_samples_batch( a, From a6cf9999d99a8237c284aa8a4b9bb3f4ef10d476 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 19 Jun 2026 17:23:52 +0200 Subject: [PATCH 14/15] add fugw to solv_gromov --- examples/backends/plot_gradient_descent.py | 43 +--- ot/batch/_quadratic.py | 69 +++++- ot/solvers.py | 105 ++++++++- test/batch/test_solve_gromov_batch.py | 254 +++++++++++++++++---- test/test_solvers.py | 6 +- 5 files changed, 390 insertions(+), 87 deletions(-) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py index 0db8cfbd6..95190f267 100644 --- a/examples/backends/plot_gradient_descent.py +++ b/examples/backends/plot_gradient_descent.py @@ -21,7 +21,7 @@ import torch from time import perf_counter import ot -from ot.batch._quadratic import loss_quadratic_samples_batch, tensor_batch +from ot.batch._quadratic import loss_quadratic_batch, tensor_batch from ot.gromov import fused_unbalanced_gromov_wasserstein from sklearn.manifold import MDS @@ -139,13 +139,8 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): M_torch = torch.tensor(M[None, :, :]) L = tensor_batch(a_torch, b_torch, C1_torch, C2_torch, loss="sqeuclidean") -alpha_batch = 0.5 -# `loss_fugw_batch` uses alpha as: alpha * quadratic + (1 - alpha) * linear -# while the dedicated solver uses alpha as the coefficient of the linear term. -alpha_bcd = (1 - alpha_batch) / alpha_batch - -reg_marginals_batch = 0.5 -reg_marginals_bcd = reg_marginals_batch / (2 * alpha_batch) +alpha = 0.5 +reg_marginals = 0.5 lr = 5e-2 nb_iter_max = 1500 tol = 1e-7 @@ -162,15 +157,15 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): optimizer.zero_grad() # Positive transport plan parameterized as log(1 + exp(T)). plan_torch = torch.nn.functional.softplus(T_torch) - loss = loss_quadratic_samples_batch( + loss = loss_quadratic_batch( a_torch, b_torch, C1_torch, C2_torch, plan_torch, M_torch, - alpha=alpha_batch, - unbalanced=reg_marginals_batch, + alpha=alpha, + unbalanced=reg_marginals, unbalanced_type="kl", recompute_const=True, )[0] @@ -200,15 +195,15 @@ def get_noisy_one_hot(labels, n_classes, noise_level=0.1): def evaluate_batch_fugw_loss(plan): plan_torch = torch.tensor(plan[None, :, :], dtype=M_torch.dtype) - loss = loss_quadratic_samples_batch( + loss = loss_quadratic_batch( a_torch, b_torch, C1_torch, C2_torch, plan_torch, M_torch, - alpha=alpha_batch, - unbalanced=reg_marginals_batch, + alpha=alpha, + unbalanced=reg_marginals, unbalanced_type="kl", recompute_const=True, )[0] @@ -216,28 +211,14 @@ def evaluate_batch_fugw_loss(plan): tic = perf_counter() -T_bcd, _, log = fused_unbalanced_gromov_wasserstein( - C1, - C2, - wx=a, - wy=b, - reg_marginals=reg_marginals_bcd, - divergence="kl", - unbalanced_solver="mm", - alpha=alpha_bcd, - M=M, - init_pi=np.outer(a, b), - max_iter=200, - tol=tol, - max_iter_ot=200, - tol_ot=1e-7, - log=True, +result = ot.solve_gromov( + C1, C2, M, a, b, alpha=alpha, reg=0, unbalanced_type="kl", unbalanced=reg_marginals ) time_bcd = perf_counter() - tic loss_adam_final = evaluate_batch_fugw_loss(T_adam) +T_bcd = result.plan loss_bcd_final = evaluate_batch_fugw_loss(T_bcd) -print(log["fugw_cost"]) mass_bcd = T_bcd.sum() pl.figure(2, (10, 4)) diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index c4e3c524f..c54639bb8 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -12,7 +12,7 @@ from ot.backend import get_backend from ot.batch._linear import loss_linear_batch from ot.batch._utils import bmv, bop, bregman_log_projection_batch -from ot.utils import list_to_array +from ot.utils import deprecated, list_to_array def tensor_batch( @@ -227,7 +227,7 @@ def div_between_product_batch(mu, nu, alpha, beta, divergence, nx=None): return res -def loss_quadratic_batch(L, T, recompute_const=False, symmetric=True, nx=None): +def loss_quadratic_tensor_batch(L, T, recompute_const=False, symmetric=True, nx=None): r""" Computes the gromov-wasserstein cost given a cost tensor and transport plan. Batched version. @@ -273,7 +273,36 @@ def loss_quadratic_batch(L, T, recompute_const=False, symmetric=True, nx=None): return nx.sum(LT * T, axis=(1, 2)) +@deprecated("Use ot.batch.loss_quadratic_batch instead.") def loss_quadratic_samples_batch( + a, + b, + C1, + C2, + T, + loss="sqeuclidean", + symmetric=None, + nx=None, + logits=None, + recompute_const=False, + log=False, +): + return loss_quadratic_batch( + a, + b, + C1, + C2, + T, + loss=loss, + symmetric=symmetric, + nx=nx, + logits=logits, + recompute_const=recompute_const, + log=log, + ) + + +def loss_quadratic_batch( a, b, C1, @@ -288,6 +317,7 @@ def loss_quadratic_samples_batch( nx=None, logits=None, recompute_const=False, + log=False, ): r""" Computes the gromov-wasserstein for samples C1, C2 and transport plan. Batched version. @@ -330,6 +360,9 @@ def loss_quadratic_samples_batch( recompute_const : bool, optional Whether to recompute the constant term. Default is False. This should be set to True if T does not satisfy the marginal constraints. Will be set to True if unbalanced is not None. + log : bool, optional + If True, also returns a dictionary containing the different terms of + the loss. Examples @@ -355,6 +388,14 @@ def loss_quadratic_samples_batch( if nx is None: nx = get_backend(T) + if log: + log_dict = {} + log_dict["value_quadratic"] = None + log_dict["value_linear"] = None + log_dict["value_unbalanced"] = None + else: + log_dict = None + if isinstance(loss, str) and loss in ["sqeuclidean", "kl", "l2"]: L = tensor_batch( a, b, C1, C2, symmetric=symmetric, nx=nx, loss=loss, logits=logits @@ -365,11 +406,16 @@ def loss_quadratic_samples_batch( if unbalanced is not None: recompute_const = True - quadratic = loss_quadratic_batch( + quadratic = loss_quadratic_tensor_batch( L, T, recompute_const=recompute_const, symmetric=symmetric, nx=nx ) + if log: + log_dict["value_quadratic"] = quadratic if unbalanced is None and M is None: + if log: + log_dict["value"] = quadratic + return quadratic, log_dict return quadratic B = T.shape[0] @@ -406,6 +452,8 @@ def loss_quadratic_samples_batch( divergence=unbalanced_type, nx=nx, ) + if log: + log_dict["value_unbalanced"] = unbalanced_term if M is not None: if alpha is None: @@ -418,15 +466,22 @@ def loss_quadratic_samples_batch( f"If alpha is not a scalar, it must have shape ({B},), got {alpha.shape}" ) linear = loss_linear_batch(M, T, nx=nx) + if log: + log_dict["value_linear"] = linear if M is not None and unbalanced is not None: - return (1 - alpha) * linear + alpha * quadratic + unbalanced * unbalanced_term + value = (1 - alpha) * linear + alpha * quadratic + unbalanced * unbalanced_term elif M is not None and unbalanced is None: - return (1 - alpha) * linear + alpha * quadratic + value = (1 - alpha) * linear + alpha * quadratic else: - return quadratic + unbalanced * unbalanced_term + value = quadratic + unbalanced * unbalanced_term + + if log: + log_dict["value"] = value + return value, log_dict + return value def solve_gromov_batch( @@ -641,7 +696,7 @@ def solve_gromov_batch( T = nx.detach(T) value_linear = loss_linear_batch(M, T, nx=nx) - value_quadratic = loss_quadratic_batch( + value_quadratic = loss_quadratic_tensor_batch( L, T, nx=nx, recompute_const=True, symmetric=symmetric ) # Always recompute const for accurate value value = (1 - alpha) * value_linear + alpha * value_quadratic diff --git a/ot/solvers.py b/ot/solvers.py index 88cf5c7ab..b49e426a7 100644 --- a/ot/solvers.py +++ b/ot/solvers.py @@ -31,6 +31,7 @@ partial_fused_gromov_wasserstein2, entropic_partial_gromov_wasserstein2, entropic_partial_fused_gromov_wasserstein2, + fused_unbalanced_gromov_wasserstein, ) from .gaussian import empirical_bures_wasserstein_distance from .factored import factored_optimal_transport @@ -674,7 +675,7 @@ def solve_gromov( function :math:`U`. Corresponds to the total transport mass for partial OT. unbalanced_type : str, optional Type of unbalanced penalization function :math:`U` either "KL", "semirelaxed", - "partial", by default "KL" but note that it is not implemented yet. + "partial", by default "KL". n_threads : int, optional Number of OMP threads for exact OT solver, by default 1 method : str, optional @@ -1085,7 +1086,55 @@ def solve_gromov( # potentials = (log['u'], log['v']) TODO elif unbalanced_type.lower() in ["kl", "l2"]: # unbalanced exact OT - raise (NotImplementedError('Unbalanced_type="{}"'.format(unbalanced_type))) + if alpha == 0: # unbalanced Wasserstein problem + res = solve( + M, + a=a, + b=b, + reg=None, + reg_type=reg_type, + unbalanced=unbalanced, + unbalanced_type=unbalanced_type, + method=method, + max_iter=max_iter, + plan_init=plan_init, + tol=tol, + verbose=verbose, + ) + + plan = res.plan + potentials = res.potentials + value_linear = res.value_linear + value = res.value + value_quad = 0 + status = res.status + + else: + if max_iter is None: + max_iter = 100 + if tol is None: + tol = 1e-7 + + # in this function alpha weights the linear and quadratic terms : alpha * quadratic + (1 - alpha) * linear + # while fused_unbalanced_gromov_wasserstein uses alpha as the coefficient of the linear term. + alpha_fugw = (1 - alpha) / alpha + reg_fugw = unbalanced / (2 * alpha) + plan, _, log = fused_unbalanced_gromov_wasserstein( + Ca, + Cb, + a, + b, + reg_marginals=reg_fugw, + divergence=unbalanced_type.lower(), + alpha=alpha_fugw, + M=M, + max_iter=max_iter, + tol=tol, + log=True, + epsilon=0, + ) + value_linear = log["linear_cost"] * alpha + value = log["fugw_cost"] * alpha else: raise ( @@ -1328,6 +1377,58 @@ def solve_gromov( # potentials = (log['u'], log['v']) TODO value = value_noreg + reg * nx.sum(plan * nx.log(plan + 1e-16)) + elif unbalanced_type.lower() in ["kl", "l2"]: + if alpha == 0: # regularized unbalanced Wasserstein problem + res = solve( + M, + a=a, + b=b, + reg=reg, + reg_type=reg_type, + unbalanced=unbalanced, + unbalanced_type=unbalanced_type, + method=method, + max_iter=max_iter, + plan_init=plan_init, + tol=tol, + verbose=verbose, + ) + + plan = res.plan + potentials = res.potentials + value_linear = res.value_linear + value = res.value + value_quad = 0 + status = res.status + + else: + if max_iter is None: + max_iter = 100 + if tol is None: + tol = 1e-7 + + # in this function alpha weights the linear and quadratic terms : alpha * quadratic + (1 - alpha) * linear + # while fused_unbalanced_gromov_wasserstein uses alpha as the coefficient of the linear term. + alpha_fugw = (1 - alpha) / alpha + reg_fugw = unbalanced / (2 * alpha) + epsilon_fugw = reg / (2 * alpha) + plan, _, log = fused_unbalanced_gromov_wasserstein( + Ca, + Cb, + a, + b, + reg_marginals=reg_fugw, + divergence=unbalanced_type.lower(), + alpha=alpha_fugw, + M=M, + max_iter=max_iter, + tol=tol, + log=True, + epsilon=epsilon_fugw, + ) + value_linear = log["linear_cost"] * alpha + value = log["fugw_cost"] * alpha + else: # unbalanced AND regularized OT raise ( NotImplementedError( diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index c242f83f7..5f6956994 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -13,7 +13,6 @@ solve_gromov_batch, loss_quadratic_batch, loss_linear_batch, - loss_quadratic_samples_batch, ) from ot import solve_gromov from ot.batch._linear import dist_batch @@ -23,6 +22,7 @@ from ot.batch._quadratic import ( tensor_batch, div_between_product_batch, + loss_quadratic_samples_batch, ) from ot.gromov._utils import div_between_product @@ -107,7 +107,7 @@ def test_all(loss, logits): res = solve_gromov_batch(C1=C, C2=C, a=a, b=a, loss=loss, logits=logits) loss1 = res.value_quad - loss2 = loss_quadratic_samples_batch( + loss2 = loss_quadratic_batch( a=a, b=a, C1=C, C2=C, T=res.plan, loss=loss, logits=logits ) np.testing.assert_allclose(loss1, loss2, atol=1e-5) @@ -146,10 +146,12 @@ def test_backend(nx): solve_gromov_batch(C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False) -@pytest.mark.parametrize("unbalanced_type", ["kl", "l2"]) -@pytest.mark.parametrize("loss", ["sqeuclidean", "kl"]) -def test_fugw_loss(unbalanced_type, loss): - """Check that loss_fugw_batch and loss_fugw_samples_batch run without error.""" +@pytest.mark.parametrize( + "loss, logits, unbalanced_type", + product(["sqeuclidean", "kl"], [True, False], ["kl", "l2"]), +) +def test_fugw_loss(unbalanced_type, loss, logits): + """Check that loss_quadratic_batch runs without error.""" batchsize = 2 n = 4 d = 2 @@ -162,7 +164,7 @@ def test_fugw_loss(unbalanced_type, loss): alpha = rng.rand() reg_marginals = rng.rand() - loss_fugw = loss_quadratic_samples_batch( + loss_fugw = loss_quadratic_batch( a, a, C1, @@ -173,9 +175,11 @@ def test_fugw_loss(unbalanced_type, loss): unbalanced=reg_marginals, unbalanced_type=unbalanced_type, loss=loss, - logits=False, + logits=logits, ) - loss_fugw_unbalanced_only = loss_quadratic_samples_batch( + + # unbalanced quadratic + loss_fugw_unbalanced_only = loss_quadratic_batch( a, a, C1, @@ -186,9 +190,11 @@ def test_fugw_loss(unbalanced_type, loss): unbalanced=reg_marginals, unbalanced_type=unbalanced_type, loss=loss, - logits=False, + logits=logits, ) - loss_fugw_no_alpha = loss_quadratic_samples_batch( + + # alpha is None + loss_fugw_no_alpha = loss_quadratic_batch( a, a, C1, @@ -199,9 +205,11 @@ def test_fugw_loss(unbalanced_type, loss): unbalanced=reg_marginals, unbalanced_type=unbalanced_type, loss=loss, - logits=False, + logits=logits, ) - loss_fugw_no_unbalanced = loss_quadratic_samples_batch( + + # balanced + loss_fugw_no_unbalanced = loss_quadratic_batch( a, a, C1, @@ -210,20 +218,83 @@ def test_fugw_loss(unbalanced_type, loss): M=M, alpha=alpha, loss=loss, - logits=False, + logits=logits, ) - assert np.isfinite(loss_fugw_unbalanced_only).all() assert np.isfinite(loss_fugw).all() + assert np.isfinite(loss_fugw_unbalanced_only).all() assert np.isfinite(loss_fugw_no_alpha).all() assert np.isfinite(loss_fugw_no_unbalanced).all() + +def test_fugw_backend(nx): + """Check that loss_quadratic_batch runs without error for all backends.""" + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1_np = rng.rand(batchsize, n, n, d) + C1 = nx.from_numpy(C1_np) + C2_np = rng.rand(batchsize, n, n, d) + C2 = nx.from_numpy(C2_np) + M_np = rng.rand(batchsize, n, n) + M = nx.from_numpy(M_np) + a_np = np.ones((batchsize, n)) + a = nx.from_numpy(a_np) + T_np = rng.rand(batchsize, n, n) + T = nx.from_numpy(T_np) + alpha_np = rng.rand() + alpha = nx.from_numpy(np.array(alpha_np)) + unbalanced_np = rng.rand() + unbalanced = nx.from_numpy(np.array(unbalanced_np)) + + loss_fugw_np = loss_quadratic_batch( + a_np, + a_np, + C1_np, + C2_np, + T_np, + M_np, + alpha=alpha_np, + unbalanced=unbalanced_np, + unbalanced_type="kl", + loss="sqeuclidean", + logits=False, + ) + loss_fugw = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="sqeuclidean", + logits=False, + ) + + assert np.allclose(loss_fugw_np, loss_fugw, atol=1e-5) + + +def test_fugw_paramters_arrays(): + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + alpha = rng.rand() # check that alpha and reg_marginals can be passed as lists or arrays of shape (batchsize,) alpha = rng.rand(batchsize) - reg_marginals = rng.rand(batchsize) + unbalanced = rng.rand(batchsize) alpha_list = alpha.tolist() - reg_marginals_list = reg_marginals.tolist() + unbalanced_list = unbalanced.tolist() - loss_fugw = loss_quadratic_samples_batch( + loss_fugw = loss_quadratic_batch( a, a, C1, @@ -231,12 +302,12 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha, - unbalanced=reg_marginals, - unbalanced_type=unbalanced_type, - loss=loss, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", logits=False, ) - loss_fugw_list = loss_quadratic_samples_batch( + loss_fugw_list = loss_quadratic_batch( a, a, C1, @@ -244,9 +315,9 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha_list, - unbalanced=reg_marginals_list, - unbalanced_type=unbalanced_type, - loss=loss, + unbalanced=unbalanced_list, + unbalanced_type="kl", + loss="l2", logits=False, ) @@ -254,9 +325,23 @@ def test_fugw_loss(unbalanced_type, loss): assert np.isfinite(loss_fugw_list).all() np.testing.assert_allclose(loss_fugw, loss_fugw_list) + +def test_fugw_invalid_loss_values(): + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + alpha = rng.rand() + unbalanced = rng.rand() + # check that invalid loss raise an error with pytest.raises(ValueError): - loss_fugw = loss_quadratic_samples_batch( + loss_quadratic_batch( a, a, C1, @@ -264,15 +349,15 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha, - unbalanced=reg_marginals, - unbalanced_type=unbalanced_type, + unbalanced=unbalanced, + unbalanced_type="kl", loss="test", logits=False, ) - # check that invalid loss raise an error + # check that invalid unbalanced_type raise an error with pytest.raises(ValueError): - loss_fugw = loss_quadratic_samples_batch( + loss_quadratic_batch( a, a, C1, @@ -280,16 +365,30 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha, - unbalanced=reg_marginals, + unbalanced=unbalanced, unbalanced_type="test", - loss=loss, + loss="l2", logits=False, ) + +def test_fugw_invalid_shapes(): + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + alpha = rng.rand() + unbalanced = rng.rand() + # check that invalid alpha shape raise an error alpha = rng.rand(batchsize + 1) with pytest.raises(ValueError): - loss_quadratic_samples_batch( + loss_quadratic_batch( a, a, C1, @@ -297,17 +396,17 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha, - unbalanced=reg_marginals, - unbalanced_type=unbalanced_type, - loss=loss, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", logits=False, ) # check that invalid rho shape raise an error alpha = rng.rand(batchsize) - reg_marginals = rng.rand(batchsize + 1) + unbalanced = rng.rand(batchsize + 1) with pytest.raises(ValueError): - loss_quadratic_samples_batch( + loss_quadratic_batch( a, a, C1, @@ -315,9 +414,9 @@ def test_fugw_loss(unbalanced_type, loss): T, M, alpha=alpha, - unbalanced=reg_marginals, - unbalanced_type=unbalanced_type, - loss=loss, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", logits=False, ) @@ -337,7 +436,7 @@ def test_valid_fugw_loss_endpoints(unbalanced_type, loss): reg_marginals = 0 T = rng.rand(batchsize, n, n) - loss_fugw = loss_quadratic_samples_batch( + loss_fugw = loss_quadratic_batch( a, a, C1, @@ -353,7 +452,7 @@ def test_valid_fugw_loss_endpoints(unbalanced_type, loss): loss_linear = loss_linear_batch(M, T) np.testing.assert_allclose(loss_fugw, loss_linear, atol=1e-5) - loss_fugw = loss_quadratic_samples_batch( + loss_fugw = loss_quadratic_batch( a, a, C1, @@ -366,7 +465,7 @@ def test_valid_fugw_loss_endpoints(unbalanced_type, loss): loss=loss, logits=False, ) - loss_gromov = loss_quadratic_samples_batch( + loss_gromov = loss_quadratic_batch( a, a, C1, @@ -401,3 +500,70 @@ def test_div_between_product(divergence): ] ) np.testing.assert_allclose(res_batch, res, atol=1e-5) + + +def test_loss_quadratic_samples_batch_deprecated(): + rng = np.random.RandomState(0) + batchsize = 2 + n = 4 + d = 2 + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + + with pytest.warns(DeprecationWarning, match="loss_quadratic_batch"): + loss_quadratic_samples_batch(a, a, C1, C2, T, loss="sqeuclidean") + + +def test_loss_quadratic_batch_log_balanced(): + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + + value, log = loss_quadratic_batch(a, a, C1, C2, T, loss="sqeuclidean", log=True) + expected = loss_quadratic_batch(a, a, C1, C2, T, loss="sqeuclidean") + + np.testing.assert_allclose(value, expected) + np.testing.assert_allclose(log["value"], expected) + np.testing.assert_allclose(log["value_quadratic"], expected) + assert log["value_linear"] is None + assert log["value_unbalanced"] is None + + +def test_loss_quadratic_batch_log_fugw(): + batchsize = 2 + n = 4 + d = 2 + rng = np.random.RandomState(0) + C1 = rng.rand(batchsize, n, n, d) + C2 = rng.rand(batchsize, n, n, d) + M = rng.rand(batchsize, n, n) + a = np.ones((batchsize, n)) + T = rng.rand(batchsize, n, n) + alpha = rng.rand() + unbalanced = rng.rand() + + value, log = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M=M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="sqeuclidean", + log=True, + ) + + expected = (1 - alpha) * log["value_linear"] + alpha * log["value_quadratic"] + expected = expected + unbalanced * log["value_unbalanced"] + np.testing.assert_allclose(value, expected) + np.testing.assert_allclose(log["value"], expected) diff --git a/test/test_solvers.py b/test/test_solvers.py index 802aca631..b2b4d5d70 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -14,15 +14,15 @@ from ot.backend import torch -lst_reg = [None, 1] +lst_reg = [None, 0.1] lst_reg_type = ["KL", "entropy", "L2", "tuple"] lst_unbalanced = [None, 0.9] lst_unbalanced_type = ["KL", "L2", "TV"] lst_reg_type_gromov = ["entropy"] lst_gw_losses = ["L2", "KL"] -lst_unbalanced_type_gromov = ["KL", "semirelaxed", "partial"] -lst_unbalanced_gromov = [None, 0.9] +lst_unbalanced_type_gromov = ["KL", "semirelaxed", "partial", "L2"] +lst_unbalanced_gromov = [None, 1.0] lst_alpha = [0, 0.4, 0.9, 1] lst_method_params_solve_sample = [ From ea9c2072f80f3b283c689767c664ba2acbfa3c75 Mon Sep 17 00:00:00 2001 From: SoniaMazelet <121769948+SoniaMaz8@users.noreply.github.com> Date: Fri, 19 Jun 2026 17:51:23 +0200 Subject: [PATCH 15/15] change C1 to Ca --- ot/batch/_quadratic.py | 194 +++++++++++++------------- test/batch/test_solve_gromov_batch.py | 12 +- 2 files changed, 103 insertions(+), 103 deletions(-) diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index c54639bb8..ccb8a8792 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -16,7 +16,7 @@ def tensor_batch( - a, b, C1, C2, symmetric=True, nx=None, loss="sqeuclidean", logits=None + a, b, Ca, Cb, symmetric=True, nx=None, loss="sqeuclidean", logits=None ): r""" Compute the Gromov-Wasserstein cost tensor for a batch of problems. @@ -39,9 +39,9 @@ def tensor_batch( Source distributions for each problem in the batch. b : array-like, shape (B, m) Target distributions for each problem in the batch. - C1 : array-like, shape (B, n, n) or (B, n, n, d) + Ca : array-like, shape (B, n, n) or (B, n, n, d) Source cost matrices for each problem. Can be a 3D array for scalar costs or a 4D array for vector-valued costs (edge features). - C2 : array-like, shape (B, m, m) or (B, n, n, d) + Cb : array-like, shape (B, m, m) or (B, n, n, d) Target cost matrices for each problem. Can be a 3D array for scalar costs or a 4D array for vector-valued costs (edge features). symmetric : bool, optional Whether the cost matrices are symmetric. Default is True. @@ -60,10 +60,10 @@ def tensor_batch( Dictionary containing: - constC : array-like, shape (B, n, m) Constant term in the tensor product. - - hC1 : array-like, shape (B, n, n, d) or (B, n, n) - - hC2 : array-like, shape (B, m, m, d) or (B, m, m) - - fC1 : array-like, shape (B, n, n) - - fC2 : array-like, shape (B, m, m) + - hCa : array-like, shape (B, n, n, d) or (B, n, n) + - hCb : array-like, shape (B, m, m, d) or (B, m, m) + - fCa : array-like, shape (B, n, n) + - fCb : array-like, shape (B, m, m) Supported loss functions: @@ -79,7 +79,7 @@ def tensor_batch( .. math:: \ell(a, b) = \sum_i a_i \log\left(\frac{a_i}{b_i}\right) - If ``logits=True``, the entries of C1 are treated as logits (unnormalized log probabilities) + If ``logits=True``, the entries of Ca are treated as logits (unnormalized log probabilities) and the loss becomes: .. math:: @@ -90,11 +90,11 @@ def tensor_batch( >>> import numpy as np >>> from ot.batch import tensor_batch >>> # Create batch of cost matrices - >>> C1 = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices - >>> C2 = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices + >>> Ca = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices + >>> Cb = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices >>> a = np.ones((3, 5)) / 5 # Uniform source distributions >>> b = np.ones((3, 4)) / 4 # Uniform target distributions - >>> L = tensor_batch(a, b, C1, C2, loss='sqeuclidean') + >>> L = tensor_batch(a, b, Ca, Cb, loss='sqeuclidean') References ---------- @@ -104,33 +104,33 @@ def tensor_batch( """ if nx is None: - nx = get_backend(C1) + nx = get_backend(Ca) loss = loss.lower() if loss == "sqeuclidean" or loss == "l2": - def f1(C1): - if C1.ndim == 4: - return nx.sum(C1**2, axis=-1) + def f1(Ca): + if Ca.ndim == 4: + return nx.sum(Ca**2, axis=-1) else: - return C1**2 + return Ca**2 - def f2(C2): - if C2.ndim == 4: - return nx.sum(C2**2, axis=-1) + def f2(Cb): + if Cb.ndim == 4: + return nx.sum(Cb**2, axis=-1) else: - return C2**2 + return Cb**2 - def h1(C1): - if C1.ndim == 3: - C1 = nx.unsqueeze(C1, -1) - return 2 * C1 + def h1(Ca): + if Ca.ndim == 3: + Ca = nx.unsqueeze(Ca, -1) + return 2 * Ca - def h2(C2): - if C2.ndim == 3: - C2 = nx.unsqueeze(C2, -1) - return C2 + def h2(Cb): + if Cb.ndim == 3: + Cb = nx.unsqueeze(Cb, -1) + return Cb elif loss == "kl": assert logits in [ @@ -138,21 +138,21 @@ def h2(C2): False, ], "logits must be either True or False for KL loss" - def f1(C1): - return nx.zeros((C1.shape[0], C1.shape[1], C1.shape[2]), type_as=C1) + def f1(Ca): + return nx.zeros((Ca.shape[0], Ca.shape[1], Ca.shape[2]), type_as=Ca) - def f2(C2): - assert C2.ndim == 4, "C2 must be a bxnxnxd tensor" - fC2 = C2 * nx.log(C2 + 1e-15) # Avoid log(0) - return nx.sum(fC2, axis=-1) + def f2(Cb): + assert Cb.ndim == 4, "Cb must be a bxnxnxd tensor" + fCb = Cb * nx.log(Cb + 1e-15) # Avoid log(0) + return nx.sum(fCb, axis=-1) - def h1(C1): - return C1 if logits else nx.log(C1 + 1e-15) + def h1(Ca): + return Ca if logits else nx.log(Ca + 1e-15) - def h2(C2): - return C2 + def h2(Cb): + return Cb - return compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=symmetric) + return compute_tensor_batch(f1, f2, h1, h2, a, b, Ca, Cb, symmetric=symmetric) def div_between_product_batch(mu, nu, alpha, beta, divergence, nx=None): @@ -249,14 +249,14 @@ def loss_quadratic_tensor_batch(L, T, recompute_const=False, symmetric=True, nx= >>> import numpy as np >>> from ot.batch import tensor_batch, loss_quadratic_batch >>> # Create batch of cost matrices - >>> C1 = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices - >>> C2 = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices + >>> Ca = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices + >>> Cb = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices >>> a = np.ones((3, 5)) / 5 # Uniform source distributions >>> b = np.ones((3, 4)) / 4 # Uniform target distributions - >>> L = tensor_batch(a, b, C1, C2, loss='sqeuclidean') + >>> L = tensor_batch(a, b, Ca, Cb, loss='sqeuclidean') >>> # Use the uniform transport plan for testing >>> T = np.ones((3, 5, 4)) / (5 * 4) - >>> loss = loss_quadratic_batch(L, T, recompute_const=True) + >>> loss = loss_quadratic_tensor_batch(L, T, recompute_const=True) >>> loss.shape (3,) @@ -277,8 +277,8 @@ def loss_quadratic_tensor_batch(L, T, recompute_const=False, symmetric=True, nx= def loss_quadratic_samples_batch( a, b, - C1, - C2, + Ca, + Cb, T, loss="sqeuclidean", symmetric=None, @@ -290,8 +290,8 @@ def loss_quadratic_samples_batch( return loss_quadratic_batch( a, b, - C1, - C2, + Ca, + Cb, T, loss=loss, symmetric=symmetric, @@ -305,8 +305,8 @@ def loss_quadratic_samples_batch( def loss_quadratic_batch( a, b, - C1, - C2, + Ca, + Cb, T, M=None, alpha=None, @@ -320,7 +320,7 @@ def loss_quadratic_batch( log=False, ): r""" - Computes the gromov-wasserstein for samples C1, C2 and transport plan. Batched version. + Computes the gromov-wasserstein for samples Ca, Cb and transport plan. Batched version. Parameters ---------- @@ -328,9 +328,9 @@ def loss_quadratic_batch( Source distributions. b : array-like, shape (B, m) Target distributions. - C1 : array-like, shape (B, n, n) or (B, n, n, d) + Ca : array-like, shape (B, n, n) or (B, n, n, d) Source cost matrices. - C2 : array-like, shape (B, m, m) or (B, n, n, d) + Cb : array-like, shape (B, m, m) or (B, n, n, d) Target cost matrices. T : array-like, shape (B, n, m) Transport plan. @@ -370,13 +370,13 @@ def loss_quadratic_batch( >>> import numpy as np >>> from ot.batch import loss_quadratic_samples_batch >>> # Create batch of cost matrices - >>> C1 = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices - >>> C2 = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices + >>> Ca = np.random.rand(3, 5, 5) # 3 problems, 5x5 source matrices + >>> Cb = np.random.rand(3, 4, 4) # 3 problems, 4x4 target matrices >>> a = np.ones((3, 5)) / 5 # Uniform source distributions >>> b = np.ones((3, 4)) / 4 # Uniform target distributions >>> # Use the uniform transport plan for testing >>> T = np.ones((3, 5, 4)) / (5 * 4) - >>> loss = loss_quadratic_samples_batch(a, b, C1, C2, T, recompute_const=True) + >>> loss = loss_quadratic_samples_batch(a, b, Ca, Cb, T, recompute_const=True) >>> loss.shape (3,) @@ -398,7 +398,7 @@ def loss_quadratic_batch( if isinstance(loss, str) and loss in ["sqeuclidean", "kl", "l2"]: L = tensor_batch( - a, b, C1, C2, symmetric=symmetric, nx=nx, loss=loss, logits=logits + a, b, Ca, Cb, symmetric=symmetric, nx=nx, loss=loss, logits=logits ) else: raise ValueError(f"Unknown loss function: {loss}") @@ -485,8 +485,8 @@ def loss_quadratic_batch( def solve_gromov_batch( - C1, - C2, + Ca, + Cb, reg=1e-2, a=None, b=None, @@ -546,9 +546,9 @@ def solve_gromov_batch( Parameters ---------- - C1 : array-like, shape (B, n, n, d) or (B, n, n) + Ca : array-like, shape (B, n, n, d) or (B, n, n) Samples affinity matrices from source distribution - C2 : array-like, shape (B, n, n, d) or (B, n, n) + Cb : array-like, shape (B, n, n, d) or (B, n, n) Samples affinity matrices from target distribution a : array-like, shape (B, n), optional Marginal distribution of the source samples. If None, uniform distribution is used. @@ -557,9 +557,9 @@ def solve_gromov_batch( loss : str, optional Type of loss function, can be 'sqeuclidean' or 'kl' or a QuadraticMetric instance. symmetric : bool, optional - Either C1 and C2 are to be assumed symmetric or not. + Either Ca and Cb are to be assumed symmetric or not. If let to its default None value, a symmetry test will be conducted. - Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric). + Else if set to True (resp. False), Ca and Cb will be assumed symmetric (resp. asymmetric). M : array-like, shape (dim_a, dim_b), optional Linear cost matrix for Fused Gromov-Wasserstein (default is None). alpha : float, optional @@ -618,24 +618,24 @@ def solve_gromov_batch( # -------------- Setup -------------- # - nx = get_backend(a, b, M, C1, C2, T_init) - B, n, m = (C1.shape[0], C1.shape[1], C2.shape[1]) + nx = get_backend(a, b, M, Ca, Cb, T_init) + B, n, m = (Ca.shape[0], Ca.shape[1], Cb.shape[1]) if a is None: - a = nx.ones((B, n), type_as=C1) / n + a = nx.ones((B, n), type_as=Ca) / n if b is None: - b = nx.ones((B, m), type_as=C2) / m + b = nx.ones((B, m), type_as=Cb) / m if symmetric is None: - symmetric = nx.allclose(C1, transpose(C1, nx=nx), atol=1e-10) and nx.allclose( - C2, transpose(C2, nx=nx), atol=1e-10 + symmetric = nx.allclose(Ca, transpose(Ca, nx=nx), atol=1e-10) and nx.allclose( + Cb, transpose(Cb, nx=nx), atol=1e-10 ) # -------------- Get cost_tensor (quadratic part) -------------- # if isinstance(loss, str): L = tensor_batch( - a, b, C1, C2, symmetric=symmetric, nx=nx, loss=loss, logits=logits + a, b, Ca, Cb, symmetric=symmetric, nx=nx, loss=loss, logits=logits ) else: raise ValueError(f"Unknown loss function: {loss}") @@ -643,7 +643,7 @@ def solve_gromov_batch( # -------------- Get cost_matrix (linear part) -------------- # if M is None and alpha is None: - M = nx.zeros((B, n, m), type_as=C1) + M = nx.zeros((B, n, m), type_as=Ca) alpha = 1.0 # Gromov problem elif M is not None and alpha is None: raise ValueError( @@ -719,11 +719,11 @@ def solve_gromov_batch( ### --------------------- Utility functions for quadratic OT --------------------- ### -def compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=True): +def compute_tensor_batch(f1, f2, h1, h2, a, b, Ca, Cb, symmetric=True): """ Gromov-Wasserstein writes as: - GW(T,C1,C2) = sum_ijkl T_ik T_jl l(C1_ij, C2_kl) = < LxT, T > - Where L is a cost tensor L[i,j,k,l] = l(C1_ij, C2_kl). + GW(T,Ca,Cb) = sum_ijkl T_ik T_jl l(Ca_ij, Cb_kl) = < LxT, T > + Where L is a cost tensor L[i,j,k,l] = l(Ca_ij, Cb_kl). For loss function of form l(a,b) = f1(a) + f2(b) - < h1(a), h2(b) > The tensor product LxT can be computed fast using tensor_product [12]. @@ -735,17 +735,17 @@ def compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=True): International Conference on Machine Learning (ICML). 2016. """ - fC1 = f1(C1) - fC2 = f2(C2) + fCa = f1(Ca) + fCb = f2(Cb) if not symmetric: - fC1 = 0.5 * (fC1 + transpose(fC1)) - fC2 = 0.5 * (fC2 + transpose(fC2)) - hC1 = h1(C1) - hC2 = h2(C2) + fCa = 0.5 * (fCa + transpose(fCa)) + fCb = 0.5 * (fCb + transpose(fCb)) + hCa = h1(Ca) + hCb = h2(Cb) - constC = compute_const_from_marginals(fC1, fC2, a, b) + constC = compute_const_from_marginals(fCa, fCb, a, b) - L = {"constC": constC, "hC1": hC1, "hC2": hC2, "fC1": fC1, "fC2": fC2} + L = {"constC": constC, "hCa": hCa, "hCb": hCb, "fCa": fCa, "fCb": fCb} return L @@ -754,8 +754,8 @@ def tensor_product_batch(L, T, nx=None, recompute_const=False, symmetric=True): """ Compute the tensor product LxT for the cost tensor L and transport plan T. The formula is: - LxT = const - hC1 T hC2^T - const = < fC1 a 1^T + 1 (fC2 b)^T + LxT = const - hCa T hCb^T + const = < fCa a 1^T + 1 (fCb b)^T References ---------- @@ -769,21 +769,21 @@ def tensor_product_batch(L, T, nx=None, recompute_const=False, symmetric=True): if recompute_const: const = compute_const_from_marginals( - L["fC1"], L["fC2"], nx.sum(T, axis=2), nx.sum(T, axis=1), nx=nx + L["fCa"], L["fCb"], nx.sum(T, axis=2), nx.sum(T, axis=1), nx=nx ) else: const = L["constC"] - hC1 = L["hC1"] - hC2 = L["hC2"] + hCa = L["hCa"] + hCb = L["hCb"] - dot = nx.einsum("bijd,bjk->bikd", hC1, T) - dot = nx.einsum("bikd,bjkd->bijd", dot, hC2) + dot = nx.einsum("bijd,bjk->bikd", hCa, T) + dot = nx.einsum("bikd,bjkd->bijd", dot, hCb) dot = nx.sum(dot, axis=-1) if not symmetric: - dot_t = nx.einsum("bijd,bjk->bikd", transpose(hC1), T) - dot_t = nx.einsum("bikd,bjkd->bijd", dot_t, transpose(hC2)) + dot_t = nx.einsum("bijd,bjk->bikd", transpose(hCa), T) + dot_t = nx.einsum("bikd,bjkd->bijd", dot_t, transpose(hCb)) dot_t = nx.sum(dot_t, axis=-1) dot = (dot + dot_t) / 2 # Average the two symmetric terms @@ -796,15 +796,15 @@ def transpose(C, nx=None): return nx.transpose(C, (0, 2, 1)) if C.ndim == 3 else nx.transpose(C, (0, 2, 1, 3)) -def compute_const_from_marginals(fC1, fC2, a, b, nx=None): +def compute_const_from_marginals(fCa, fCb, a, b, nx=None): """ - Compute the constant term f1(C1) a 1^T + 1 b^T f2(C2)^T + Compute the constant term f1(Ca) a 1^T + 1 b^T f2(Cb)^T """ if nx is None: - nx = get_backend(fC1, fC2, a, b) - fC1a = bmv(fC1, a, nx=nx) - fC2b = bmv(fC2, b, nx=nx) - constC = fC1a[:, :, None] + fC2b[:, None, :] + nx = get_backend(fCa, fCb, a, b) + fCaa = bmv(fCa, a, nx=nx) + fCbb = bmv(fCb, b, nx=nx) + constC = fCaa[:, :, None] + fCbb[:, None, :] return constC @@ -813,7 +813,7 @@ def detach_cost_tensor(L, nx=None): Detach the cost tensor L to avoid gradients. """ if nx is None: - nx = get_backend(L["constC"], L["hC1"], L["hC2"]) + nx = get_backend(L["constC"], L["hCa"], L["hCb"]) L_detached = {} for key, value in L.items(): L_detached[key] = nx.detach(value) diff --git a/test/batch/test_solve_gromov_batch.py b/test/batch/test_solve_gromov_batch.py index 5f6956994..d59c89d89 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -56,8 +56,8 @@ def test_solve_gromov_batch(): alpha=alpha, reg=reg, M=M, - C1=C1, - C2=C2, + Ca=C1, + Cb=C2, max_iter=max_iter, tol=tol, max_iter_inner=max_iter_inner, @@ -104,11 +104,11 @@ def test_all(loss, logits): C = np.abs(C) + 1e-6 C = C / np.sum(C, axis=-1, keepdims=True) - res = solve_gromov_batch(C1=C, C2=C, a=a, b=a, loss=loss, logits=logits) + res = solve_gromov_batch(Ca=C, Cb=C, a=a, b=a, loss=loss, logits=logits) loss1 = res.value_quad loss2 = loss_quadratic_batch( - a=a, b=a, C1=C, C2=C, T=res.plan, loss=loss, logits=logits + a=a, b=a, Ca=C, Cb=C, T=res.plan, loss=loss, logits=logits ) np.testing.assert_allclose(loss1, loss2, atol=1e-5) @@ -122,7 +122,7 @@ def test_gradients_torch(grad): d = 2 C = torch.randn((batchsize, n, n, d), requires_grad=True) res = solve_gromov_batch( - C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False, grad=grad + Ca=C, Cb=C, a=None, b=None, loss="sqeuclidean", logits=False, grad=grad ) loss = res.value.sum() loss_plan = res.plan.sum() @@ -143,7 +143,7 @@ def test_backend(nx): d = 2 C = np.random.randn(batchsize, n, n, d) C = nx.from_numpy(C) - solve_gromov_batch(C1=C, C2=C, a=None, b=None, loss="sqeuclidean", logits=False) + solve_gromov_batch(Ca=C, Cb=C, a=None, b=None, loss="sqeuclidean", logits=False) @pytest.mark.parametrize(