diff --git a/README.md b/README.md index 4ce4082ab..950a509e9 100644 --- a/README.md +++ b/README.md @@ -455,18 +455,20 @@ Artificial Intelligence. \[82] Bonet, C., Nadjahi, K., Séjourné, T., Fatras, K., & Courty, N. (2024). [Slicing Unbalanced Optimal Transport](https://openreview.net/forum?id=AjJTg5M0r8). Transactions on Machine Learning Research. -[83] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920). +\[83] Germain, T., Flamary, R., Kostic, V. R., & Lounici, K. (2025). [A Spectral-Grassmann Wasserstein Metric for Operator Representations of Dynamical Systems](https://arxiv.org/abs/2509.24920). -[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15. +\[84] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15. -[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385. +\[85] Mahey, G., Chapel, L., Gasso, G., Bonet, C., & Courty, N. (2023). [Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics](https://proceedings.neurips.cc/paper_files/paper/2023/hash/6f1346bac8b02f76a631400e2799b24b-Abstract-Conference.html). Advances in Neural Information Processing Systems, 36, 35350–35385. -[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661. +\[86] Tanguy, E., Chapel, L., Delon, J. (2025). [Sliced Transport Plans](https://arxiv.org/abs/2508.01243) arXiv preprint 2506.03661. -[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations. +\[87] Liu, X., Diaz Martin, R., Bai Y., Shahbazi A., Thorpe M., Aldroubi A., Kolouri, S. (2024). [Expected Sliced Transport Plans](https://openreview.net/forum?id=P7O1Vt1BdU). International Conference on Learning Representations. -[88] Bouveyron, C. & Corneli, M. (2026). [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88. +\[88] Bouveyron, C. & Corneli, M. (2026). [Scaling optimal transport to high-dimensional Gaussian distributions with application to domain adaptation](https://hal.science/hal-04930868v4/file/Article-OT-HDGauss-v4.pdf). Statistics and Computing 36.2 (2026): 88. -[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. +\[89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. -[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. +\[90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. + +\[91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & Courty, N. (2021). [Minibatch optimal transport distances; analysis and applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. \ No newline at end of file diff --git a/RELEASES.md b/RELEASES.md index eea94981d..dad6fa306 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -38,7 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Update the geomloss wrapper to the new version and API (PR #826) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) -- Update sgot cost function and example (PR #830) +- Implement debiased OT solvers in `ot.solve_sample`. #### Closed issues @@ -61,7 +61,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Fix documentation build on master with submodules (PR #818) - Fix failing test for unbalanced solver with generic regularization (PR #824) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - +- Update sgot cost function and example (PR #830) ## 0.9.6.post1 diff --git a/examples/plot_debias_sink_div.py b/examples/plot_debias_sink_div.py new file mode 100644 index 000000000..d6f44bafd --- /dev/null +++ b/examples/plot_debias_sink_div.py @@ -0,0 +1,254 @@ +# -*- coding: utf-8 -*- +""" +====================================== +Sinkhorn Divergence and Debiased OT solvers +====================================== + +This example shows how to use the debiased OT solvers in `ot.solve_sample` to +compute Sinkhorn divergences and debiased Minibatch solutions. The debiased OT solvers +can be used with balanced and unbalanced OT problems, and with different +regularization types (entropic, L2, group lasso). +""" + +# Author: Remi Flamary +# +# License: MIT License +# sphinx_gallery_thumbnail_number = 3 + +# %% + +import numpy as np +import matplotlib.pylab as pl +import ot +import ot.plot +from ot.datasets import make_1D_gauss as gauss + +############################################################################## +# Generate data +# ------------- + + +# %% +def sample_ball(n, radius=1.0, center=(0.0, 0.0)): + np.random.seed(0) + theta = 2 * np.pi * np.random.rand(n) + r = radius * np.sqrt(np.random.rand(n)) + + x = r * np.cos(theta) + center[0] + y = r * np.sin(theta) + center[1] + + return np.stack((x, y), axis=1) + + +def sample_two_balls(n, radius=1.0, sep=1): + assert n % 2 == 0, "n must be even" + centers = ((-sep, -sep), (sep, sep)) + n_half = n // 2 + X1 = sample_ball(n_half, radius, centers[0]) + X2 = sample_ball(n_half, radius, centers[1]) + + perm = np.random.permutation(n_half * 2) + X = np.vstack((X1, X2)) + X = X[perm] + return X + + +n = 50 + +x1 = sample_ball(n, radius=1.0, center=(0, 0)) +x2 = sample_two_balls(n, radius=1.0, sep=1.5) + +pl.figure(1, figsize=(5, 5)) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", alpha=0.7) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", alpha=0.7) +pl.legend() +pl.title("Two distributions") +ax = pl.axis() + + +############################################################################## +# Compute Sinkhorn divergence and visualize plans +# ----------------------------------------------- +# The Sinkhorn divergence is computed by setting the `debias` parameter to +# `True` in the `ot.solve_sample` function. The resulting value is the Sinkhorn +# divergence. The Sinkhorn divergences is computed as: +# +# .. math:: +# S_\epsilon(\mu, \nu) = OT_\epsilon(\mu, \nu) - \frac{1}{2} OT_\epsilon(\mu, \mu) - \frac{1}{2} OT_\epsilon(\nu, \nu) +# +# The entropic OT plans for each of those terms can be accessed in the `log` +# attribute of the result, and can be visualized using the +# `ot.plot.plot2D_samples_mat` function. + +res = ot.solve_sample(x1, x2, reg=0.1, debias=True) + +print("Sinkhorn divergence: ", res.value) + +plan_11 = res.log["res_aa"].plan +plan_12 = res.log["res_ab"].plan +plan_22 = res.log["res_bb"].plan + +# +pl.figure(2, figsize=(15, 5)) + +pl.subplot(1, 3, 1) +ot.plot.plot2D_samples_mat(x1, x1, plan_11, thr=0.05) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) +pl.axis(ax) +pl.title("Plan between source and source") +pl.subplot(1, 3, 2) +ot.plot.plot2D_samples_mat(x1, x2, plan_12, thr=0.05) +pl.scatter(x1[:, 0], x1[:, 1], label="Source distribution", zorder=2) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", zorder=2) +pl.axis(ax) +pl.title("Plan between source and target") +pl.subplot(1, 3, 3) +ot.plot.plot2D_samples_mat(x2, x2, plan_22, thr=0.05) +pl.scatter(x2[:, 0], x2[:, 1], label="Target distribution", color="C1", zorder=2) +pl.axis(ax) +pl.title("Plan between target and target") + + +############################################################################## +# Debiased Minibatch OT +# --------------------------------- +# +# Doing OT on minibatches leads to a similar bias than using entropic +# regularization since the average OT plan is densified due to the stochasticity +# of the minibatch sampling. On a given minibatch, the debiased loss can be +# computed by setting the `debias` parameter to `'split'`that split the data +# points in each distributions in two and computes the debias OT loss as: +# +# .. math:: +# \tilde{OT}_m(\mu, \nu) = \frac{1}{2}(\hat{OT}_m(\mu_1, \nu_1) + \hat{OT}_m(\mu_2, \nu_2) - \hat{OT}_m(\nu_1, \nu_2) - \hat{OT}_m(\mu_1, \nu_2)) +# + +# %% solve OT minibtach and visualize the plans + +res = ot.solve_sample(x1, x2, debias="split") + +print("Debiased minibatch OT loss: ", res.value) + +# recover the plans for each of the four terms in the debiased loss +plan_11 = res.log["res_aa"].plan +plan_12 = res.log["res_ab1"].plan +plan_21 = res.log["res_ab2"].plan +plan_22 = res.log["res_bb"].plan +sel_a1 = res.log["sel_a1"] +sel_a2 = res.log["sel_a2"] +sel_b1 = res.log["sel_b1"] +sel_b2 = res.log["sel_b2"] + +nb1 = plan_11.shape[0] +nb2 = plan_22.shape[0] + +pl.figure(4, figsize=(15, 3)) + +pl.subplot(1, 4, 1) +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="$\mu_1$", zorder=2) +pl.scatter( + x1[sel_a2, 0], x1[sel_a2, 1], label=r"$\mu_2$", zorder=2, color="C0", alpha=0.5 +) +pl.scatter(x2[sel_b1, 0], x2[sel_b1, 1], label=r"$\nu_1$", zorder=2, color="C1") +pl.scatter( + x2[sel_b2, 0], x2[sel_b2, 1], label=r"$\nu_2$", zorder=2, color="C1", alpha=0.5 +) +pl.title("Minibatch split") +pl.axis(ax) +pl.legend() + + +pl.subplot(1, 4, 2) +ot.plot.plot2D_samples_mat(x1[sel_a1], x1[sel_a2], plan_11, thr=0.05) +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], zorder=2) +pl.scatter( + x1[sel_a2, 0], + x1[sel_a2, 1], + zorder=2, + color="C0", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between source and source") +pl.subplot(1, 4, 3) +ot.plot.plot2D_samples_mat(x1[sel_a1], x2[sel_b1], plan_12, thr=0.05) +ot.plot.plot2D_samples_mat(x1[sel_a2], x2[sel_b2], plan_21, thr=0.05, alpha=0.5) + +pl.scatter(x1[sel_a1, 0], x1[sel_a1, 1], label="Source distribution", zorder=2) +pl.scatter( + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" +) +pl.scatter( + x1[sel_a2, 0], + x1[sel_a2, 1], + label="Source distribution", + zorder=2, + color="C0", + alpha=0.5, +) +pl.scatter( + x2[sel_b2, 0], + x2[sel_b2, 1], + label="Target distribution", + zorder=2, + color="C1", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between source and target") + +pl.subplot(1, 4, 4) +ot.plot.plot2D_samples_mat(x2[sel_b1], x2[sel_b2], plan_22, thr=0.05) +pl.scatter( + x2[sel_b1, 0], x2[sel_b1, 1], label="Target distribution", zorder=2, color="C1" +) +pl.scatter( + x2[sel_b2, 0], + x2[sel_b2, 1], + label="Target distribution", + zorder=2, + color="C1", + alpha=0.5, +) +pl.axis(ax) +pl.title("Plan between target and target") + + +############################################################################## +# Comparison of the divergences +# ------------------------------------------------- + +# %% move a distribution and compute Sinkhorn divergence and Sinkhorn distance +reg = 0.1 + +sep_list = np.linspace(0, 1.0, 10) +sink_list = [] +sink_div_list = [] +ot_mb_list = [] +ot_mb_sink_list = [] +for sep in sep_list: + x2sep = sample_two_balls(n, radius=1.0, sep=sep) + sink_list.append( + ot.solve_sample( + x1, + x2sep, + reg=reg, + ).value + ) + sink_div_list.append(ot.solve_sample(x1, x2sep, reg=reg, debias=True).value) + ot_mb_list.append(ot.solve_sample(x1, x2sep, debias="split").value) + + ot_mb_sink_list.append(ot.solve_sample(x1, x2sep, reg=1, debias="split").value) + +pl.figure(3) +pl.plot(sep_list, sink_list, label="Sinkhorn loss", color="C0") +pl.plot(sep_list, sink_div_list, label="Sinkhorn divergence", color="C1") +pl.plot(sep_list, ot_mb_list, label="Debiased MB OT", color="C2") +pl.plot(sep_list, ot_mb_sink_list, label="Debiased MB Sinkhorn", color="C3") +pl.xlabel("Separation between distributions") +pl.ylabel("Loss/Divergence") +pl.title("Comparison of biased VS debiased OT losses") +pl.grid() +pl.legend() + +# %% diff --git a/examples/plot_quickstart_guide.py b/examples/plot_quickstart_guide.py index d5cfac8e6..6ba3d410a 100644 --- a/examples/plot_quickstart_guide.py +++ b/examples/plot_quickstart_guide.py @@ -467,8 +467,8 @@ def df(G): # loss_fgw = ot.gromov.fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1) # loss_efgw = ot.gromov.entropic_fused_gromov_wasserstein2(C1, C2, M, a, b, alpha=0.1, epsilon=reg) # -# Large scale OT -# -------------- +# Large scale OT and approximations +# --------------------------------- # # We discuss here strategies to solve large scale OT problems using approximations # of the exact OT problem. @@ -583,11 +583,67 @@ def df(G): print(f"Exact OT loss = {loss:1.3f}") print(f"Bures-Wasserstein distance = {bw_value:1.3f}") +# %% +# One can also use the HD gaussian assumption (low rank covariance + diagonal) +# that has better properties in high dimension. The rank of the covariance +# matrices can be controlled with the :code:`rank` parameter. + +hdbw_value = ot.solve_sample(x1, x2, a, b, method="gaussian_hd", rank=1).value + +print(f"Bures-Wasserstein distance = {bw_value:1.3f}") +print(f"High Dimensional Bures-Wasserstein distance = {hdbw_value}") + # %% # .. note:: # The Gaussian Wasserstein problem can be solved with the classic API using the # :func:`ot.gaussian.empirical_bures_wasserstein_distance` function. # +# Sliced Wasserstein +# ~~~~~~~~~~~~~~~~~~ +# +# The Sliced Wasserstein distance is a Wasserstein distance between +# empirical distributions that is computed by projecting the samples on random +# directions and averaging the Wasserstein distances between the projected +# distributions. It can be used as an approximation of the Wasserstein distance +# between empirical distributions. + +sw_value = ot.solve_sample(x1, x2, a, b, method="sliced", n_projections=10).value +max_sw_value = ot.solve_sample( + x1, x2, a, b, method="max_sliced", n_projections=10 +).value + +print(f"Exact OT loss = {loss:1.3f}") +print(f"Sliced Wasserstein distance = {sw_value:1.3f}") +print(f"Max Sliced Wasserstein distance = {max_sw_value:1.3f}") + +# %% +# Binary Space Partitioning (BSP) OT +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# One can also use the BSP OT approximation that is based on a recursive +# partitioning of the space and computes the Wasserstein distance between the +# empirical distributions by solving small OT problems between the samples in each +# partition. The number of partitions can be controlled with the +# `n_projections` parameter. + +# BSP can only find bijections so require same number of points +x1_bsp = np.concatenate([x1, x1], axis=0) + +sol_bsp = ot.solve_sample(x1_bsp, x2, method="bsp", n_projections=10) +bsp_value = sol_bsp.value +bsp_sparse_plan = sol_bsp.sparse_plan + +# sphinx_gallery_start_ignore + +pl.figure(1, (3, 3)) +plot2D_samples_mat(x1_bsp, x2, bsp_sparse_plan) +pl.plot(x1_bsp[:, 0], x1_bsp[:, 1], "ob", label="Source samples", **style) +pl.plot(x2[:, 0], x2[:, 1], "or", label="Target samples", **style) +pl.title("BSP OT plan loss={:.3f}".format(bsp_value)) +pl.show() + +# sphinx_gallery_end_ignore +# %% # Comparing all OT plans # ---------------------- # diff --git a/ot/backend.py b/ot/backend.py index 08ab22bdd..7749c948d 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -1867,9 +1867,7 @@ def randperm(self, size, type_as=None): if not isinstance(size, int): raise ValueError("size must be an integer") if type_as is not None: - return jnp.asarray( - jax.random.permutation(subkey, size), dtype=type_as.dtype - ) + return jnp.asarray(jax.random.permutation(subkey, size)) else: return jax.random.permutation(subkey, size) @@ -2444,7 +2442,6 @@ def randperm(self, size, type_as=None): ) return torch.randperm( n=size, - dtype=type_as.dtype, generator=generator, device=type_as.device, ) @@ -2909,7 +2906,7 @@ def randperm(self, size, type_as=None): return self.rng_.permutation(size) else: with cp.cuda.Device(type_as.device): - return cp.asarray(self.rng_.permutation(size), dtype=type_as.dtype) + return cp.asarray(self.rng_.permutation(size)) def coo_matrix(self, data, rows, cols, shape=None, type_as=None): data = self.from_numpy(data) @@ -3360,7 +3357,7 @@ def randperm(self, size, type_as=None): ) else: return tf.random.experimental.stateless_shuffle( - tf.range(size, dtype=type_as.dtype), seed=local_seed + tf.range(size), seed=local_seed ) def _convert_to_index_for_coo(self, tensor): diff --git a/ot/gaussian.py b/ot/gaussian.py index 59d368afa..103d102d8 100644 --- a/ot/gaussian.py +++ b/ot/gaussian.py @@ -171,15 +171,6 @@ def bures_wasserstein_mapping_hd(ms, mt, Us, Ut, ls, lt, sigma2_s, sigma2_t, log """ nx = get_backend(ms, mt, Us, Ut, ls, lt, sigma2_s, sigma2_t) - print(ms.dtype) - print(mt.dtype) - print(Us.dtype) - print(Ut.dtype) - print(ls.dtype) - print(lt.dtype) - print(sigma2_s.dtype) - print(sigma2_t.dtype) - p = Us.shape[0] # source @@ -429,8 +420,8 @@ def empirical_bures_wasserstein_mapping_hd( Ut = Qt[:, -dt:] lt = a_t - sgm2_t - sgm2_s = nx.unsqueeze(sgm2_s, 0) - sgm2_t = nx.unsqueeze(sgm2_t, 0) + sgm2_s = nx.maximum(nx.unsqueeze(sgm2_s, 0), reg) + sgm2_t = nx.maximum(nx.unsqueeze(sgm2_t, 0), reg) if log: A, b, log = bures_wasserstein_mapping_hd( @@ -889,20 +880,20 @@ def empirical_bures_wasserstein_distance_hd( eigs = nx.eigh(Cs) a_s = eigs[0][-ds:] - sgm2_s = (nx.trace(Cs) - nx.sum(a_s)) / (p - ds) + sgm2_s = (nx.trace(Cs) - nx.sum(a_s)) / (p - ds + reg) Qs = eigs[1] Us = Qs[:, -ds:] ls = a_s - sgm2_s eigt = nx.eigh(Ct) a_t = eigt[0][-dt:] - sgm2_t = (nx.trace(Ct) - nx.sum(a_t)) / (p - dt) + sgm2_t = (nx.trace(Ct) - nx.sum(a_t)) / (p - dt + reg) Qt = eigt[1] Ut = Qt[:, -dt:] lt = a_t - sgm2_t - sgm2_s = nx.unsqueeze(sgm2_s, 0) - sgm2_t = nx.unsqueeze(sgm2_t, 0) + sgm2_s = nx.maximum(nx.unsqueeze(sgm2_s, 0), reg) + sgm2_t = nx.maximum(nx.unsqueeze(sgm2_t, 0), reg) if log: W, log = bures_wasserstein_distance_hd( diff --git a/ot/solvers/_bary.py b/ot/solvers/_bary.py index 7ec7469d4..988a5dc37 100644 --- a/ot/solvers/_bary.py +++ b/ot/solvers/_bary.py @@ -12,7 +12,7 @@ from ..lp import free_support_barycenter_generic_costs from ..backend import get_backend -from ._linear import solve, solve_sample, lst_method_lazy +from ._linear import solve, solve_sample, lst_method_solve_sample def _bary_sample_bcd( @@ -434,7 +434,7 @@ def solve_bary_sample( """ - if method is not None and method.lower() in lst_method_lazy: + if method is not None and method.lower() in lst_method_solve_sample: raise NotImplementedError( f"method {method} operating on lazy tensors is not implemented yet" ) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index a550df315..ad27c6e91 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -8,7 +8,7 @@ # # License: MIT License -from ..utils import OTResult, dist +from ..utils import OTResult, dist, split_sample_ratio, apply_scaler from ..lp import emd2, emd2_lazy, wasserstein_1d from ..backend import get_backend from ..unbalanced import mm_unbalanced, sinkhorn_knopp_unbalanced, lbfgsb_unbalanced @@ -18,17 +18,30 @@ empirical_sinkhorn_nystroem2, old_geomloss, ) +from ..sliced import sliced_wasserstein_distance, max_sliced_wasserstein_distance +from ..bsp import compute_bspot_bijection from ..smooth import smooth_ot_dual -from ..gaussian import empirical_bures_wasserstein_distance +from ..gaussian import ( + empirical_bures_wasserstein_distance, + empirical_bures_wasserstein_distance_hd, +) from ..factored import factored_optimal_transport from ..lowrank import lowrank_sinkhorn from ..optim import cg from warnings import warn +lst_valid_methods_solve = [ + "sinkhorn", # sinkhorn for unbalanced + "sinkhorn_log", # sinkhorn for unbalanced + "mm", # unbalanced + "lbfgsb", # unbalanced +] + -lst_method_lazy = [ +lst_method_solve_sample = [ "1d", "gaussian", + "gaussian_hd", "lowrank", "nystroem", "factored", @@ -37,8 +50,13 @@ "geomloss_tensorized", "geomloss_online", "geomloss_multiscale", + "sliced", + "max_sliced", + "bsp", ] +all_valid_methods_solve_sample = lst_valid_methods_solve + lst_method_solve_sample + def solve( M, @@ -294,6 +312,13 @@ def solve( if reg is None: reg = 0 + if method is not None and method not in lst_valid_methods_solve: + raise ValueError( + "Unknown method={} for solve, must be in {}".format( + method, lst_valid_methods_solve + ) + ) + # default values for solutions potentials = None value = None @@ -463,8 +488,15 @@ def solve( if reg_type.lower() == "entropy": value = value_linear + reg * nx.sum(plan * nx.log(plan + 1e-16)) else: + # stabilize kl of 0 mass + if nx.any(a == 0) or nx.any(b == 0): + a = a + 1e-15 * nx.min(a) + b = b + 1e-15 * nx.min(b) + print( + "Warning: a or b has 0 mass, adding small mass to avoid NaN in KL divergence." + ) value = value_linear + reg * nx.kl_div( - plan, a[:, None] * b[None, :] + plan, a[:, None] * b[None, :] + 1e-15 ) if grad == "envelope": # set the gradient at convergence @@ -619,6 +651,10 @@ def solve_sample( verbose=False, grad="autodiff", random_state=None, + debias=False, + n_projections=50, + projections=None, + scaler=None, ): r"""Solve the discrete optimal transport problem using the samples in the source and target domains. @@ -655,17 +691,17 @@ def solve_sample( c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. - If :math:`\texttt{reg_type}=`'entropy', then :math:`\mathbf{c} = 1_{dim_a} 1_{dim_b}^T`. + If `reg_type='entropy'`, then :math:`\mathbf{c} = 1_{dim_a} 1_{dim_b}^T`. reg_type : str, optional Type of regularization :math:`R` either "KL", "L2", "entropy", by default "KL" unbalanced : float or indexable object of length 1 or 2 Marginal relaxation term. If it is a scalar or an indexable object of length 1, then the same relaxation is applied to both marginal relaxations. - The balanced OT can be recovered using :math:`unbalanced=float("inf")`. + The balanced OT can be recovered using `unbalanced=float("inf")`. For semi-relaxed case, use either - :math:`unbalanced=(float("inf"), scalar)` or - :math:`unbalanced=(scalar, float("inf"))`. + `unbalanced=(float("inf"), scalar)` or + `unbalanced=(scalar, float("inf"))`. If unbalanced is an array, it must have the same backend as input arrays `(a, b, M)`. unbalanced_type : str, optional @@ -679,7 +715,21 @@ def solve_sample( method : str, optional Method for solving the problem, this can be used to select the solver for unbalanced problems (see :any:`ot.solve`), or to select a specific - large scale solver. + large scale solver. Available methods are: + + - "1d" for 1D OT solver done in parallel for each dimension. + - "gaussian" for Gaussian OT solver that estimates mean and + covariance and solve the closed form solution) + - "gaussian_hd" for high-dimensional Gaussian OT solver with rank given by `rank` parameter . + - "lowrank" for low-rank Sinkhorn solver (see :any:`ot.lowrank`) + - "nystroem" for Nystroem Sinkhorn solver + - "factored" for factored OT solver + - "geomloss" for GeomLoss Sinkhorn solver + - "sliced" for sliced wasserstein distance (see :any:`ot.sliced`) + - "max_sliced" for max sliced wasserstein distance (see + :any:`ot.sliced`) + - "bsp" for BSP-OT solver (only for distribution with same number of + samples and uniform weights) n_threads : int, optional Number of OMP threads for exact OT solver, by default 1 max_iter : int, optional @@ -706,6 +756,20 @@ def solve_sample( detached. This is useful for memory saving when only the value is needed. random_state : int, optional The random state for sampling the components in each distribution for method='nystroem'. + debias : bool, optional + Whether to use the debiased version of the OT problem, by default False + if True, the value returned is the Sinkhorn divergence [23] but the plan is + still the Sinkhorn plan. The results for all pairs of problems and + provided in the log dictionary. if debiased='split', then X_a and X_b + are split into two halves and the debiased value is computed using the + two halves as proposed in minibatch OT [91]. + n_projections : int, optional + Number of projections for sliced OT, by default 50 + projections : array-like, shape (n_projections, dim), optional + Projections for sliced OT, by default None (random projections are used) + scaler : callable, optional + Function to scale the input data, following :any:`ot.utils.apply_scaler`. + by default None (no scaling). Returns ------- @@ -812,6 +876,22 @@ def solve_sample( res = ot.solve_sample(xa, xb, a, b, reg=1.0, reg_type='L2') + - **Debiased OT and Sinkhorn divergence** (when ``debias=True``): + + One can compute the Sinkhorn divergence between two distributions using the + following code: + + .. code-block:: python + + div = ot.solve_sample(xa, xb, a, b, reg=1.0, debias=True).value + + Debiasing for all existing solvers is implemented and in particular one can + recover the minibatch exact OT debiasing [91] using the following code: + + .. code-block:: python + + div = ot.solve_sample(xa, xb, a, b, debias='split').value + - **Unbalanced OT [41]** (when ``unbalanced!=None``): .. math:: @@ -879,7 +959,7 @@ def solve_sample( Corresponds to a low rank approximation of entropic OT (for a squared Euclidean cost) that runs in linear time. - - **Gaussian Bures-Wasserstein [2]** (when ``method='gaussian'``): + - **Gaussian Bures-Wasserstein** (when ``method='gaussian'``): This method computes the Gaussian Bures-Wasserstein distance between two Gaussian distributions estimated from the empirical distributions @@ -901,6 +981,14 @@ def solve_sample( # recover the squared Gaussian Bures-Wasserstein distance BW_dist = res.value + The Gaussian-HD variant of this method can be used for high-dimensional data with the following code: + + .. code-block:: python + + res = ot.solve_sample(xa, xb, method='gaussian_hd', rank=10) + + where the rank parameter controls the rank of the covariance matrices used in the computation of the Bures-Wasserstein distance. + - **Wasserstein 1d [1]** (when ``method='1D'``): This method computes the Wasserstein distance between two 1d distributions @@ -914,6 +1002,38 @@ def solve_sample( # recover the squared Wasserstein distances W_dists = res.value + - **Sliced Wasserstein distance** (when ``method='sliced'``): + + This method computes the sliced Wasserstein distance between two + distributions. The sliced Wasserstein distance + is defined as the average of the Wasserstein distances between the projected + distributions on random directions. + + .. code-block:: python + + swd = ot.solve_sample(xa, xb, method='sliced', n_projections=50).value + + - **Max Sliced Wasserstein [34]** (when ``method='max_sliced'``): + + This method computes the max sliced Wasserstein distance between two + distributions. + + .. code-block:: python + + mswd = ot.solve_sample(xa, xb, method='max_sliced', n_projections=50).value + + - **Binary Space Partitioning OT (BSP)** (when ``method='bsp'``): + + This method computes the BSP-OT distance between two distributions and + alignments. The number of partitioning directions can be set with the `n_projections` parameter. + + .. code-block:: python + + res = ot.solve_sample(xa, xb, method='bsp', n_projections=50) + + bsp_ot_distance = res.value + bsp_plan = res.plan + .. _references-solve-sample: References @@ -936,6 +1056,10 @@ def solve_sample( Optimal Transport. Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics (AISTATS). + .. [23] Aude, G., Peyré, G., Cuturi, M., Learning Generative Models with + Sinkhorn Divergences, Proceedings of the Twenty-First International + Conference on Artificial Intelligence and Statistics, (AISTATS) 21, 2018 + .. [34] Feydy, J., Séjourné, T., Vialard, F. X., Amari, S. I., Trouvé, A., & Peyré, G. (2019, April). Interpolating between optimal transport and MMD using Sinkhorn divergences. In The 22nd International Conference @@ -957,10 +1081,129 @@ def solve_sample( .. [80] Altschuler, J., Bach, F., Rudi, A., Niles-Weed, J. (2019). Massively scalable Sinkhorn distances via the Nyström method. NeurIPS. + .. [91] Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribonval, R., & + Courty, N. (2021). [Minibatch optimal transport distances; analysis and + applications](https://arxiv.org/pdf/2101.01792). arXiv preprint arXiv:2101.01792. """ - if method is not None and method.lower() in lst_method_lazy: + if method is not None and method.lower() not in all_valid_methods_solve_sample: + raise ValueError( + "Unknown method={} for solve_sample. Available methods are: {}".format( + method, all_valid_methods_solve_sample + ) + ) + + if scaler is not None: + X_a, Xb = apply_scaler(X_a, X_b, scaler=scaler) + + if debias: + dict_params = dict( + metric=metric, + reg=reg, + c=c, + reg_type=reg_type, + unbalanced=unbalanced, + unbalanced_type=unbalanced_type, + lazy=lazy, + batch_size=batch_size, + method=method, + n_threads=n_threads, + max_iter=max_iter, + plan_init=plan_init, + rank=rank, + scaling=scaling, + potentials_init=potentials_init, + X_init=X_init, + tol=tol, + verbose=verbose, + grad=grad, + random_state=random_state, + debias=False, + ) + + nx = get_backend(X_a, X_b, a, b) + + if nx.__name__ == "tf": + raise NotImplementedError( + "Debiasing is not implemented for TensorFlow backend." + ) + + if isinstance(debias, str) and debias == "split": + # split the samples into two halves with each half the mass + X_a1, X_a2, a1, a2, sel_a1, sel_a2 = split_sample_ratio(X_a, a, nx=nx) + X_b1, X_b2, b1, b2, sel_b1, sel_b2 = split_sample_ratio(X_b, b, nx=nx) + + # compute the four OT problems + resaa = solve_sample(X_a1, X_a2, a=a1, b=a2, **dict_params) + resbb = solve_sample(X_b1, X_b2, a=b1, b=b2, **dict_params) + resab1 = solve_sample(X_a1, X_b1, a=a1, b=b1, **dict_params) + resab2 = solve_sample(X_a2, X_b2, a=a2, b=b2, **dict_params) + + # compute debiased values + value = (resab1.value + resab2.value) - (resaa.value + resbb.value) + value_linear = value + + log = { + "res_aa": resaa, + "res_bb": resbb, + "res_ab1": resab1, + "res_ab2": resab2, + "sel_a1": sel_a1, + "sel_a2": sel_a2, + "sel_b1": sel_b1, + "sel_b2": sel_b2, + "a1": a1, + "a2": a2, + "b1": b1, + "b2": b2, + } + + res = OTResult( + value=value, + value_linear=value_linear, + plan=None, # no plan for debiased version + potentials=None, # no potentials for debiased version + sparse_plan=None, # no sparse plan for debiased version + lazy_plan=None, # no lazy plan for debiased version + status="Debiased", + log=log, + backend=nx, + ) + + else: + # standard debiasing à la sinkhorn divergence + + resaa = solve_sample(X_a, X_a, a=a, b=a, **dict_params) + + resbb = solve_sample(X_b, X_b, a=b, b=b, **dict_params) + + resab = solve_sample(X_a, X_b, a=a, b=b, **dict_params) + + # compute debiased values + value = resab.value - 0.5 * (resaa.value + resbb.value) + value_linear = resab.value_linear - 0.5 * ( + resaa.value_linear + resbb.value_linear + ) + + log = {"res_aa": resaa, "res_bb": resbb, "res_ab": resab} + + res = OTResult( + value=value, + value_linear=value_linear, + plan=resab.plan, + potentials=resab.potentials, + sparse_plan=resab.sparse_plan, + lazy_plan=resab.lazy_plan, + status=resab.status, + log=log, + backend=nx, + ) + + # return debiased result + return res + + if method is not None and method.lower() in all_valid_methods_solve_sample: lazy0 = lazy lazy = True @@ -989,14 +1232,7 @@ def solve_sample( return res - elif ( - lazy - and method is None - and (reg is None or reg == 0) - and unbalanced is None - and X_a is not None - and X_b is not None - ): + elif lazy and method is None and (reg is None or reg == 0) and unbalanced is None: # Use lazy EMD solver with coordinates (no regularization, balanced) nx = get_backend(X_a, X_b, a, b) value_linear, log = emd2_lazy( @@ -1032,6 +1268,7 @@ def solve_sample( value_linear = None plan = None lazy_plan = None + sparse_plan = None status = None log = None @@ -1050,6 +1287,89 @@ def solve_sample( value = wasserstein_1d(X_a, X_b, a, b, p=p) value_linear = value + elif method == "sliced": # Sliced Wasserstein distance + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + value, log = sliced_wasserstein_distance( + X_a, + X_b, + a=a, + b=b, + p=p, + n_projections=n_projections, + projections=projections, + seed=random_state, + log=True, + ) + value_linear = value + + elif method == "max_sliced": # Max Sliced Wasserstein distance + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + value, log = max_sliced_wasserstein_distance( + X_a, + X_b, + a=a, + b=b, + p=p, + n_projections=n_projections, + projections=projections, + seed=random_state, + log=True, + ) + value_linear = value + + elif method == "bsp": # BSP-OT solver + if metric == "sqeuclidean": + p = 2 + elif metric in ["euclidean", "cityblock"]: + p = 1 + else: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + if (X_a.shape[0] != X_b.shape[0]) or a is not None or b is not None: + raise ValueError( + "BSP-OT solver requires the same number of samples in both distributions and uniform weights (provided as None)." + ) + + if random_state is None: + random_state = 0 + + n = X_a.shape[0] + + value, perm, perms = compute_bspot_bijection( + X_a, X_b, p=p, seed=random_state, n_plans=n_projections + ) + value_linear = value + sparse_plan = nx.coo_matrix( + nx.ones(n, type_as=X_a) / n, + nx.arange(n), + perm, + shape=(n, n), + type_as=X_a, + ) + + log = { + "perm": perm, + "perms": perms, + } + elif method == "gaussian": # Gaussian Bures-Wasserstein if metric.lower() not in ["sqeuclidean"]: raise ( @@ -1059,8 +1379,39 @@ def solve_sample( if reg is None: reg = 1e-6 + if a is not None and len(a.shape) == 1: + a = a.reshape(-1, 1) + if b is not None and len(b.shape) == 1: + b = b.reshape(-1, 1) + value, log = empirical_bures_wasserstein_distance( - X_a, X_b, reg=reg, log=True + X_a, X_b, reg=reg, ws=a, wt=b, log=True + ) + value = value**2 # return the value (squared bures distance) + value_linear = value # return the value + + elif method == "gaussian_hd": # Gaussian BW for high-dimensional data + if metric.lower() not in ["sqeuclidean"]: + raise ( + NotImplementedError('Not implemented metric="{}"'.format(metric)) + ) + + if reg is None: + reg = 1e-5 + + if a is not None and len(a.shape) == 1: + a = a.reshape(-1, 1) + if b is not None and len(b.shape) == 1: + b = b.reshape(-1, 1) + + if rank > X_a.shape[1] or rank > X_b.shape[1]: + warn( + "Rank is larger than the number of features. Using full rank for Gaussian Bures-Wasserstein distance." + ) + rank = min(X_a.shape[1], X_b.shape[1]) + + value, log = empirical_bures_wasserstein_distance_hd( + X_a, X_b, d_intrinsic=rank, ws=a, wt=b, reg=reg, log=True ) value = value**2 # return the value (squared bures distance) value_linear = value # return the value @@ -1257,6 +1608,7 @@ def solve_sample( value=value, lazy_plan=lazy_plan, value_linear=value_linear, + sparse_plan=sparse_plan, plan=plan, status=status, backend=nx, diff --git a/ot/utils.py b/ot/utils.py index f79d69809..fcd823415 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -2047,3 +2047,104 @@ def fun_numpy(x): return nx.to_numpy(fun(nx.from_numpy(x))) return fun_numpy + + +def split_sample_ratio( + X_a, a=None, ratio=0.5, random_split=False, random_state=None, nx=None +): + """Split distribution according to a ratio of weights (using point ordering). + + + Parameters + ---------- + X_a : array-like, shape (n_samples_a, dim) + samples in the source domain + a : array-like, shape (dim_a,), optional + Samples weights in the source domain (default is uniform) + nx : backend, optional + Backend for array operations, by default None (auto-detect) + ratio : float, optional + Ratio of the split, by default 0.5 + random_split : bool, optional + Whether to split randomly, by default False + random_state : int, optional + Random state for reproducibility, by default None + + Returns + ------- + + X_a1 : array-like, shape (n_samples_a1, dim) + First half of the samples in the source domain + X_a2 : array-like, shape (n_samples_a2, dim) + Second half of the samples in the source domain + a1 : array-like, shape (dim_a1,) + First half of the weights in the source domain + a2 : array-like, shape (dim_a2,) + Second half of the weights in the source domain + sel_a1 : slice, ndarray-like + Slice or indexes for the first half of the samples in the source domain + sel_a2 : slice, ndarray-like + Slice or indexes for the second half of the samples in the source domain + """ + + if ratio < 0 or ratio > 1: + raise ValueError("ratio should be in [0, 1]") + + if nx is None: + nx = get_backend(X_a, a) + + n_a = X_a.shape[0] + + if a is None: + a = nx.ones(n_a, type_as=X_a) / n_a + + if random_split: + if random_state is not None: + nx.seed(random_state) + perm = nx.randperm(n_a, type_as=X_a) + X_a = X_a[perm] + a = a[perm] + + # find the split indices + acs = nx.cumsum(a) + thr_a = ratio * nx.sum(a) + idx_a = nx.searchsorted(acs, thr_a, side="right") + + # split the samples and weights + X_a1, X_a2 = X_a[: idx_a + 1], X_a[idx_a:] + + # compute weights for each half and adjust the last/first weight to sum to 0.5 + a01, a02 = a[:idx_a], a[idx_a + 1 :] + v_idx = a[idx_a] + a1_1 = nx.maximum(v_idx * nx.detach((thr_a - nx.sum(a01)) / v_idx), 0) + a2_0 = nx.maximum(v_idx * nx.detach((nx.sum(a) - thr_a - nx.sum(a02)) / v_idx), 0) + # concat mass or remove samples if mass is 0 + if a1_1 > 0: + a1 = nx.concatenate([a01, nx.reshape(a1_1, (1,))]) + if random_split: + sel_a1 = perm[: idx_a + 1] + else: + sel_a1 = slice(0, idx_a + 1) + else: + X_a1 = X_a[:idx_a] + if random_split: + sel_a1 = perm[:idx_a] + else: + sel_a1 = slice(0, idx_a) + a1 = a01 + + if a2_0 > 0: + a2 = nx.concatenate([nx.reshape(a2_0, (1,)), a02]) + if random_split: + sel_a2 = perm[idx_a:] + else: + sel_a2 = slice(idx_a, n_a) + else: + X_a2 = X_a[idx_a + 1 :] + if random_split: + sel_a2 = perm[idx_a + 1 :] + else: + sel_a2 = slice(idx_a + 1, n_a) + a2 = a02 + + return X_a1, X_a2, a1, a2, sel_a1, sel_a2 diff --git a/test/test_solvers.py b/test/test_solvers.py index 505cdb3af..6cc8137cc 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -13,7 +13,8 @@ import ot from ot.bregman import geomloss from ot.backend import torch -from ot.solvers._linear import lst_method_lazy +from ot.solvers._linear import lst_method_solve_sample +from ot.utils import DataScaler lst_reg = [None, 0.1] @@ -32,9 +33,15 @@ {"method": "1d", "metric": "euclidean"}, {"method": "gaussian"}, {"method": "gaussian", "reg": 1}, + {"method": "gaussian_hd", "rank": 1}, + {"method": "gaussian_hd", "rank": 3}, {"method": "factored", "rank": 2}, {"method": "lowrank", "rank": 2, "max_iter": 5}, {"method": "nystroem", "rank": 2}, + {"method": "sliced", "n_projections": 10}, + {"method": "sliced", "n_projections": 10, "metric": "euclidean"}, + {"method": "max_sliced", "n_projections": 10}, + {"method": "max_sliced", "n_projections": 10, "metric": "euclidean"}, ] lst_parameters_solve_sample_NotImplemented = [ @@ -43,6 +50,10 @@ "method": "gaussian", "metric": "euclidean", }, # fail gaussian on metric not euclidean + { + "method": "gaussian_hd", + "metric": "euclidean", + }, # fail gaussian on metric not euclidean { "method": "factored", "metric": "euclidean", @@ -61,10 +72,18 @@ "reg": 1, "unbalanced": 1, }, # fail lazy for unbalanced and regularized + { + "method": "sliced", + "metric": "wrong_metric", + }, # fail sliced on metric not euclidean + { + "method": "max_sliced", + "metric": "wrong_metric", + }, # fail sliced on metric not euclidean ] lst_parameters_solve_bary_sample_NotImplemented = [ - {"method": method} for method in lst_method_lazy + {"method": method} for method in lst_method_solve_sample ] + [ {"lazy": True}, # fail lazy {"metric": "cosine"}, # fail on invalid metric @@ -735,33 +754,36 @@ def test_solve_sample_geomloss(nx, metric): sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=False, method="geomloss") assert_allclose_sol(sol0, sol) - sol1 = ot.solve_sample( - xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_tensorized" - ) - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) - - sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_online") - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) - - sol1 = ot.solve_sample( - xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_multiscale" - ) - np.testing.assert_allclose( - nx.to_numpy(sol1.plan), - nx.to_numpy(sol.plan), - rtol=1e-5, - atol=1e-5, - ) + # commented because geomloss_tensorized and geomloss_online are not + # implemented in geomloss 0.2.0 yet + + # sol1 = ot.solve_sample( + # xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_tensorized" + # ) + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) + + # sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_online") + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) + + # sol1 = ot.solve_sample( + # xb, yb, ab, bb, reg=1, lazy=True, method="geomloss_multiscale" + # ) + # np.testing.assert_allclose( + # nx.to_numpy(sol1.plan), + # nx.to_numpy(sol.plan), + # rtol=1e-5, + # atol=1e-5, + # ) sol1 = ot.solve_sample(xb, yb, ab, bb, reg=1, lazy=True, method="geomloss") np.testing.assert_allclose( @@ -793,10 +815,134 @@ def test_solve_sample_methods(nx, method_params): assert_allclose_sol(sol, solb) sol2 = ot.solve_sample(x, x, **method_params) - if method_params["method"] not in ["factored", "lowrank", "nystroem"]: + if method_params["method"] not in [ + "factored", + "lowrank", + "nystroem", + "gaussian_hd", + ]: np.testing.assert_allclose(sol2.value, 0, atol=1e-10) +def test_zero_mass_solvers(): + # test that solvers handle zero mass distributions correctly + n_samples_s = 10 + n_samples_t = 9 + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + # Set one of the distributions to zero mass + a[0] = 0.0 + b[0] = 0.0 + + a /= a.sum() + b /= b.sum() + + res = ot.solve_sample(x, y, a, b) + res_reg = ot.solve_sample(x, y, a, b, reg=1.0) + + assert np.isnan(res.value) == False + assert np.isnan(res_reg.value) == False + + +@pytest.skip_backend("tf", reason="Not implemented for tf backend") +@pytest.mark.parametrize("debias", [True, False, "split"]) +@pytest.mark.parametrize("reg", [None, 10]) +def test_solve_sample_debias(nx, debias, reg): + n_samples_s = 10 + n_samples_t = 9 + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + xb, yb, ab, bb = nx.from_numpy(x, y, a, b) + + sol = ot.solve_sample(x, y, reg=reg, debias=debias) + solb = ot.solve_sample(xb, yb, ab, bb, reg=reg, debias=debias) + + # check some attributes (no need ) + assert_allclose_sol(sol, solb) + + +def test_solve_sample_bsp(nx): + n_samples_s = 10 + n_samples_t = 10 # same number of samples for source and target + n_features = 2 + rng = np.random.RandomState(42) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + + xb, yb = nx.from_numpy(x, y) + + sol = ot.solve_sample(x, y, method="bsp") + solb = ot.solve_sample(xb, yb, method="bsp") + + # check some attributes (no need ) + assert_allclose_sol(sol, solb) + + with pytest.raises(ValueError): + # bsp method requires same number of samples for source and target + ot.solve_sample(x, y[:5], method="bsp") + + with pytest.raises(NotImplementedError): + # bsp method requires same number of samples for source and target + ot.solve_sample(x, y, method="bsp", metric="wrong_metric") + + +def test_solvers_bad_method(): + n_samples_s = 20 + n_samples_t = 7 + n_features = 2 + rng = np.random.RandomState(0) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + + C = ot.dist(x, y) + + with pytest.raises(ValueError): + ot.solve_sample(x, y, method="invalid_method") + + with pytest.raises(ValueError): + ot.solve(C, method="invalid_method") + + +@pytest.mark.parametrize("norm", ["standard", "minmax", "l2"]) +def test_solve_sample_scaler(nx, norm): + n_samples_s = 20 + n_samples_t = 7 + n_features = 2 + rng = np.random.RandomState(0) + + x = rng.randn(n_samples_s, n_features) + y = rng.randn(n_samples_t, n_features) + a = ot.utils.unif(n_samples_s) + b = ot.utils.unif(n_samples_t) + + xb, yb, ab, bb = nx.from_numpy(x, y, a, b) + + scaler0 = DataScaler(norm=norm) + scaler0.fit(x) + + scaler1 = DataScaler(norm=norm) + scaler1.fit(xb) + + sol0 = ot.solve_sample(x, y, a, b, scaler=scaler0) + sol1 = ot.solve_sample(xb, yb, ab, bb, scaler=scaler1) + + assert_allclose_sol(sol0, sol1) + + @pytest.mark.parametrize("method_params", lst_parameters_solve_sample_NotImplemented) def test_solve_sample_NotImplemented(nx, method_params): n_samples_s = 20 diff --git a/test/test_utils.py b/test/test_utils.py index 05885a1c2..b2d6b4508 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -957,3 +957,39 @@ def test_datascaler_backend_mismatch_raises(): scaler._nx = object() with pytest.raises(ValueError, match="Backend mismatch"): scaler.transform(X) + + +@pytest.skip_backend("tf") +@pytest.mark.parametrize("ratio", [0.3, 0.5, 0.51, 0.7]) +@pytest.mark.parametrize("n", [10, 50, 100, 101]) +@pytest.mark.parametrize("random", [True, False]) +def test_split_sample_ratio(nx, ratio, n, random): + rng = np.random.RandomState(0) + X = rng.normal(0, 1, (100, 3)) + a = rng.uniform(size=(100,)) + a = a / np.sum(a) + X = nx.from_numpy(X) + a = nx.from_numpy(a) + + seed = 42 + + # test uniform weights + X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( + X, ratio=ratio, random_split=random, random_state=seed, nx=nx + ) + + np.testing.assert_allclose(nx.to_numpy(a1).sum(), ratio, atol=1e-5) + np.testing.assert_allclose(nx.to_numpy(a2).sum(), 1 - ratio, atol=1e-5) + + assert X1.shape[0] == a1.shape[0] + assert X2.shape[0] == a2.shape[0] + + # test non uniform weights + X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio( + X, ratio=ratio, a=a, random_split=random, random_state=seed + ) + + with pytest.raises(ValueError): + ot.utils.split_sample_ratio(X, ratio=1.5, a=a, random_state=seed, nx=nx) + + #