diff --git a/RELEASES.md b/RELEASES.md index 47c1f19a0..f3c17d427 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -27,7 +27,10 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver a callable, or a no-op (PR #808) - Add optional `scaler` parameter to `sliced_wasserstein_distance` and `max_sliced_wasserstein_distance` (PR #808) - Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723) -- Add cost functions between linear operators following [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), implemented in `ot.sgot` (PR #792) +- Add cost functions between linear operators following + [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), + implemented in `ot.sgot` (PR #792) +- Add batch FUGW loss to `ot.batch` and fix issues in some default parameters in the batch module (PR #775) - Build wheels on ubuntu ARM to avoid QEMU emulation (PR #818) - Add new methods to compute the linear transport map and the related 2-Wasserstein distance betweeen high-dimensional (HD) Gaussian distributions as described in [88], implemented in `ot.gaussian.bures_wasserstein_mapping_hd` and `ot.gaussian.bures_wasserstein_distance_hd`, respectively. Two additional methods estimate the same quantities from the source and destination observed data and are implemented in `ot.gaussian.empirical_bures_wasserstein_mapping_hd` and `ot.gaussian.empirical_bures_wasserstein_distance_hd`, respectively (PR #814) diff --git a/examples/backends/plot_gradient_descent.py b/examples/backends/plot_gradient_descent.py new file mode 100644 index 000000000..95190f267 --- /dev/null +++ b/examples/backends/plot_gradient_descent.py @@ -0,0 +1,266 @@ +# -*- coding: utf-8 -*- +r""" +=============================================================================== +Solve Fused Unbalanced Gromov Wasserstein with Adam +=============================================================================== + +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`. +""" + +# Author: Rémi Flamary +# Sonia Mazelet +# +# License: MIT License + +# 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_quadratic_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 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") + + +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 = 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.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) +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(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 +) +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(pos1, C1, color=colors1) +pl.title("SBM source graph") +pl.axis("off") +pl.subplot(1, 2, 2) +plot_graph(pos2, C2, color=colors2) +pl.title("SBM target graph") +_ = pl.axis("off") + + +# %% +# Solve FUGW with Adam +# ---------------- + +# 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.5 +reg_marginals = 0.5 +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)). + plan_torch = torch.nn.functional.softplus(T_torch) + loss = loss_quadratic_batch( + a_torch, + b_torch, + C1_torch, + C2_torch, + plan_torch, + M_torch, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type="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] + + +# %% +# Compare with the dedicated FUGW solver +# ------------------------------------- +# +# 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_quadratic_batch( + a_torch, + b_torch, + C1_torch, + C2_torch, + plan_torch, + M_torch, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type="kl", + recompute_const=True, + )[0] + return float(loss.detach()) + + +tic = perf_counter() +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) +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() + + +# %% +# Visualize the learned couplings +# ------------------------------- +# 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. + +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", cmap="Blues", vmin=vmin, vmax=vmax) +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", 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") +_ = pl.colorbar() 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 0da4b8962..ccb8a8792 100644 --- a/ot/batch/_quadratic.py +++ b/ot/batch/_quadratic.py @@ -12,10 +12,11 @@ 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 deprecated, list_to_array 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. @@ -38,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. @@ -59,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: @@ -78,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:: @@ -89,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 ---------- @@ -103,31 +104,33 @@ def tensor_batch( """ if nx is None: - nx = get_backend(C1) + nx = get_backend(Ca) - if loss == "sqeuclidean": + loss = loss.lower() - def f1(C1): - if C1.ndim == 4: - return nx.sum(C1**2, axis=-1) + if loss == "sqeuclidean" or loss == "l2": + + 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 [ @@ -135,24 +138,96 @@ 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(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(Ca): + return Ca if logits else nx.log(Ca + 1e-15) + + def h2(Cb): + return Cb + + return compute_tensor_batch(f1, f2, h1, h2, a, b, Ca, Cb, symmetric=symmetric) - 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 h1(C1): - return C1 if logits else nx.log(C1 + 1e-15) +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. - def h2(C2): - return C2 + For half-squared L2 divergence: + + .. math:: + \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: + + .. math:: + 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)) - return compute_tensor_batch(f1, f2, h1, h2, a, b, C1, C2, symmetric=symmetric) + where: + - :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 -def loss_quadratic_batch(L, T, recompute_const=False, symmetric=True, nx=None): + Parameters + ---------- + 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) + nx : backend, optional + If let to its default value None, a backend test will be conducted. + + Returns + ---------- + Bregman divergence between two product measures for each problem in the batch. + """ + + if nx is None: + nx = get_backend(mu, nu, alpha, beta) + + 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 = ( + 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 + ) + + elif divergence == "l2": + res = ( + 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 + + +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. @@ -174,14 +249,14 @@ def loss_quadratic_batch(L, T, recompute_const=False, symmetric=True, nx=None): >>> 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,) @@ -198,20 +273,54 @@ 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, + Ca, + Cb, T, loss="sqeuclidean", symmetric=None, nx=None, logits=None, recompute_const=False, + log=False, +): + return loss_quadratic_batch( + a, + b, + Ca, + Cb, + T, + loss=loss, + symmetric=symmetric, + nx=nx, + logits=logits, + recompute_const=recompute_const, + log=log, + ) + + +def loss_quadratic_batch( + a, + b, + Ca, + Cb, + T, + M=None, + alpha=None, + unbalanced=None, + unbalanced_type="kl", + loss="sqeuclidean", + symmetric=True, + nx=None, + logits=None, + recompute_const=False, + 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 ---------- @@ -219,34 +328,55 @@ def loss_quadratic_samples_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. + 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. + log : bool, optional + If True, also returns a dictionary containing the different terms of + the loss. + Examples -------- >>> 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,) @@ -255,20 +385,108 @@ 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 isinstance(loss, str): + 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 + a, b, Ca, Cb, symmetric=symmetric, nx=nx, loss=loss, logits=logits ) else: raise ValueError(f"Unknown loss function: {loss}") - return loss_quadratic_batch( + + if unbalanced is not None: + recompute_const = True + + 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] + + if unbalanced is not None: + if unbalanced_type is None: + raise ValueError( + "unbalanced_type must be specified if unbalanced is not None" + ) + + unbalanced_type = unbalanced_type.lower() + + if unbalanced_type not in ["kl", "l2"]: + raise ValueError( + f"Unknown unbalanced_type: {unbalanced_type}, expected 'kl' or 'l2'" + ) + + 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, + ) + if log: + log_dict["value_unbalanced"] = unbalanced_term + + 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) + if log: + log_dict["value_linear"] = linear + + if M is not None and unbalanced is not None: + value = (1 - alpha) * linear + alpha * quadratic + unbalanced * unbalanced_term + + elif M is not None and unbalanced is None: + value = (1 - alpha) * linear + alpha * quadratic + + else: + value = quadratic + unbalanced * unbalanced_term + + if log: + log_dict["value"] = value + return value, log_dict + return value def solve_gromov_batch( - C1, - C2, + Ca, + Cb, reg=1e-2, a=None, b=None, @@ -328,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. @@ -339,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 @@ -400,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}") @@ -425,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( @@ -478,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 @@ -501,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]. @@ -517,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 @@ -536,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 ---------- @@ -551,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 @@ -578,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 @@ -595,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/ot/solvers.py b/ot/solvers.py index b84d4b1d5..22a831c93 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 @@ -679,7 +680,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 @@ -1090,7 +1091,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 ( @@ -1333,6 +1382,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_batch.py b/test/batch/test_solve_batch.py index 45a7e69fe..17d459a43 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -1,9 +1,8 @@ -"""Tests for module bregman on OT with bregman projections""" +"""Tests for module batch""" # Author: Remi Flamary -# Kilian Fatras -# Quang Huy Tran -# Eduardo Fernandes Montesuma +# Paul Krzakala +# Sonia Mazelet # # License: MIT License @@ -143,3 +142,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 e0029689b..d59c89d89 100644 --- a/test/batch/test_solve_gromov_batch.py +++ b/test/batch/test_solve_gromov_batch.py @@ -1,19 +1,30 @@ -"""Tests for module bregman on OT with bregman projections""" +"""Tests for module batch""" # Author: Remi Flamary -# Kilian Fatras -# Quang Huy Tran -# Eduardo Fernandes Montesuma +# Paul Krzakala +# 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, +) 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, + div_between_product_batch, + loss_quadratic_samples_batch, +) +from ot.gromov._utils import div_between_product def test_solve_gromov_batch(): @@ -45,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, @@ -93,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_samples_batch( - a=a, b=a, C1=C, C2=C, T=res.plan, loss=loss, logits=logits + loss2 = loss_quadratic_batch( + a=a, b=a, Ca=C, Cb=C, T=res.plan, loss=loss, logits=logits ) np.testing.assert_allclose(loss1, loss2, atol=1e-5) @@ -111,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() @@ -132,4 +143,427 @@ 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( + "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 + 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() + reg_marginals = rng.rand() + + loss_fugw = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=logits, + ) + + # unbalanced quadratic + loss_fugw_unbalanced_only = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M=None, + alpha=alpha, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=logits, + ) + + # alpha is None + loss_fugw_no_alpha = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M=None, + alpha=None, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=logits, + ) + + # balanced + loss_fugw_no_unbalanced = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M=M, + alpha=alpha, + loss=loss, + logits=logits, + ) + 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) + unbalanced = rng.rand(batchsize) + alpha_list = alpha.tolist() + unbalanced_list = unbalanced.tolist() + + loss_fugw = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", + logits=False, + ) + loss_fugw_list = loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha_list, + unbalanced=unbalanced_list, + unbalanced_type="kl", + loss="l2", + logits=False, + ) + + assert np.isfinite(loss_fugw).all() + 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_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="test", + logits=False, + ) + + # check that invalid unbalanced_type raise an error + with pytest.raises(ValueError): + loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="test", + 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_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", + logits=False, + ) + + # check that invalid rho shape raise an error + alpha = rng.rand(batchsize) + unbalanced = rng.rand(batchsize + 1) + with pytest.raises(ValueError): + loss_quadratic_batch( + a, + a, + C1, + C2, + T, + M, + alpha=alpha, + unbalanced=unbalanced, + unbalanced_type="kl", + loss="l2", + logits=False, + ) + + +@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 + 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) + + loss_fugw = loss_quadratic_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_quadratic_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_batch( + a, + a, + C1, + C2, + T, + unbalanced=reg_marginals, + unbalanced_type=unbalanced_type, + loss=loss, + logits=False, + ) + np.testing.assert_allclose(loss_fugw, loss_gromov, atol=1e-5) + + +@pytest.mark.parametrize("divergence", ["kl", "l2"]) +def test_div_between_product(divergence): + batchsize = 2 + n = 4 + m = 3 + rng = np.random.RandomState(0) + 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) + ] + ) + 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 = [