Spherical Sliced-Wasserstein Embedding on Sphere

Here, we aim at transforming samples into a uniform distribution on the sphere by minimizing SSW:

\[\min_{x} SSW_2(\nu, \frac{1}{n}\sum_{i=1}^n \delta_{x_i})\]

where \(\nu=\mathrm{Unif}(S^{d-1})\).

# Author: Clément Bonet <clement.bonet@univ-ubs.fr>
#
# License: MIT License

# sphinx_gallery_thumbnail_number = 3

import numpy as np
import matplotlib.pyplot as pl
import matplotlib.animation as animation
import torch
import torch.nn.functional as F

import ot

Data generation

torch.manual_seed(1)

N = 500
x0 = torch.rand(N, 3)
x0 = F.normalize(x0, dim=-1)

Plot data

def plot_sphere(ax):
    # Create a sphere using spherical coordinates
    phi = np.linspace(0, 2 * np.pi, 100)
    theta = np.linspace(0, np.pi, 100)
    phi, theta = np.meshgrid(phi, theta)

    # Compute the spherical coordinates
    X = np.sin(theta) * np.cos(phi)
    Y = np.sin(theta) * np.sin(phi)
    Z = np.cos(theta)

    # Plot the wireframe
    ax.plot_wireframe(X, Y, Z, color="gray", alpha=0.3)


# plot the distributions
pl.figure(1)
ax = pl.axes(projection="3d")
plot_sphere(ax)
ax.scatter(x0[:, 0], x0[:, 1], x0[:, 2], label="Data samples", alpha=0.5)
ax.set_title("Data distribution")
ax.legend()