Sliced Unbalanced optimal transport

This example illustrates the behavior of Sliced UOT versus Unbalanced Sliced OT, introduced in [82]. The first one removes outliers on each slice while the second one removes outliers of the original marginals.

[82] Bonet, C., Nadjahi, K., Séjourné, T., Fatras, K., & Courty, N. (2025). Slicing Unbalanced Optimal Transport. Transactions on Machine Learning Research.

# Author: Clément Bonet <clement.bonet.mapp@polytechnique.edu>
#         Nicolas Courty <nicolas.courty@irisa.fr>
#
# License: MIT License

# sphinx_gallery_thumbnail_number = 4

import numpy as np
import matplotlib.pylab as pl
import ot
import torch
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from sklearn.neighbors import KernelDensity

Generate data

np.random.seed(42)

n_samples = 25  # 500
nb_outliers = 10  # 200


mu_s = np.array([0, 0]) - 0.5
cov_s = 0.2**2 * np.array([[1, 0], [0, 1]])

mu_s_outliers = -np.array([2, 0.5])
cov_s_outliers = 0.05**2 * np.array([[1, 0], [0, 1]])

mu_t = np.array([0, 0]) + 1.5
cov_t = 0.2**2 * np.array([[1, 0], [0, 1]])


def generate_dataset(n_samples):
    # Generate source data (with outliers)
    Xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s)
    Xs_outlier = ot.datasets.make_2D_samples_gauss(
        nb_outliers, mu_s_outliers, cov_s_outliers
    )

    Xs = np.vstack((Xs, Xs_outlier))
    Xs_torch = torch.from_numpy(Xs).type(torch.float)

    # Generate target data
    Xt = ot.datasets.make_2D_samples_gauss(n_samples, mu_t, cov_t)
    Xt_torch = torch.from_numpy(Xt).type(torch.float)

    return Xs_torch, Xt_torch


Xs, Xt = generate_dataset(n_samples)

pl.figure(1)
pl.scatter(Xs[:, 0], Xs[:, 1], color="blue", label="Source data")
pl.scatter(Xt[:, 0], Xt[:, 1], color="red", label="Target data")
pl.xlim(-2.4, 2.4)
pl.ylim(-1, 2.2)
pl.legend()
pl.show()