[MRG] Add QSW sampling for sliced Wasserstein

README.md

@@ -484,3 +484,7 @@ Artificial Intelligence.
 
 \[94] Mezzadri, F. (2007). [How to generate random matrices from the classical compact groups](https://www.ams.org/notices/200705/fea-mezzadri-web.pdf). Notices of the American Mathematical Society, 54(5), 592-604.
 
+\[95] Nguyen, K., Bariletto, N., & Ho, N. (2024). [Quasi-Monte Carlo for 3D Sliced Wasserstein](https://arxiv.org/abs/2309.11713). International Conference on Learning Representations (ICLR).
+
+\[96] Rakhmanov, E. A., Saff, E. B., & Zhou, Y. M. (1994). [Minimal Discrete Energy on the Sphere](https://www.math.vanderbilt.edu/~esaff/texts/155.pdf). Mathematical Research Letters, 1(6), 647-662.
+

RELEASES.md

@@ -5,6 +5,9 @@
 #### New features
 
 - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836)
+- Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized
+  spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`,
+  as described in [95] (PR #838)
 
 #### Closed issues
 

examples/sliced-wasserstein/plot_qsw_3d.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+"""
+=========================================================
+Quasi-Monte Carlo Sliced Wasserstein in 3D
+=========================================================
+
+This example illustrates the Quasi-Sliced Wasserstein (QSW) and Randomized
+Quasi-Sliced Wasserstein (RQSW) sampling schemes introduced in [95], and
+compares them to the default uniform (Monte Carlo) sampling of slicing
+directions.
+
+Sliced Wasserstein (SWD) approximates the Wasserstein distance by averaging
+1D Wasserstein distances over projections onto random directions
+:math:`\\theta` drawn uniformly on the sphere. By default these directions
+are sampled purely at random (Monte Carlo), which introduces some variance
+in the estimate for a given number of projections.
+
+QSW replaces the random directions with a deterministic, low-discrepancy
+point set on the sphere (generalized spiral points), which covers the
+sphere more evenly than random sampling and reduces the approximation
+error, especially in 3D. Since QSW is deterministic it cannot directly be
+used as an unbiased estimator in stochastic settings (e.g. gradient-based
+optimization) -- RQSW addresses this by applying a random rotation to the
+same point set, which preserves both its low discrepancy and its
+unbiasedness.
+
+We first visualize the three sampling schemes on the sphere, then measure
+how fast each one converges to the true Sliced Wasserstein distance
+between two point clouds -- known here in closed form, with no
+approximation error left except from the number of projections itself.
+
+.. [95] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for
+    3D Sliced Wasserstein. International Conference on Learning
+    Representations (ICLR).
+.. [96] Rakhmanov, E. A., Saff, E. B., & Zhou, Y. M. (1994). Minimal
+    Discrete Energy on the Sphere. Mathematical Research Letters, 1(6),
+    647-662.
+"""
+
+# Author: Samuel Vangu <samuelvangu0@gmail.com>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 1
+
+import numpy as np
+import matplotlib.pylab as pl
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (registers the 3D projection)
+
+import ot
+from ot.sliced import get_random_projections, get_projections_spiral
+
+##############################################################################
+# Visualize the three sampling schemes on the sphere
+# ----------------------------------------------------
+# We draw a few hundred directions on :math:`S^2` with each scheme:
+#
+# - ``uniform``: directions are Gaussian vectors normalized to unit norm
+#   (standard Monte Carlo sampling of the sphere).
+# - ``spiral_qmc``: deterministic generalized spiral points -- a simple,
+#   closed-form low-discrepancy point set (Rakhmanov, Saff & Zhou, 1994)
+#   [96]. The same call always returns the same points.
+# - ``randomized_spiral_qmc``: the same spiral point set, rotated by a
+#   random (3, 3) rotation matrix (drawn via QR decomposition of a
+#   Gaussian matrix). The rotation makes the estimator unbiased while
+#   keeping the points as evenly spread out as the deterministic spiral
+#   set.
+
+n_projections = 500
+d = 3
+seed = 42
+
+theta_uniform = get_random_projections(d, n_projections, seed=seed)
+theta_qsw = get_projections_spiral(d, n_projections, randomized=False)
+theta_rqsw = get_projections_spiral(d, n_projections, randomized=True, seed=seed)
+
+fig = pl.figure(1, figsize=(15, 5))
+
+schemes = [
+    (theta_uniform, "Uniform (Monte Carlo)"),
+    (theta_qsw, "QSW (deterministic spiral)"),
+    (theta_rqsw, "RQSW (randomly rotated spiral)"),
+]
+
+for i, (theta, title) in enumerate(schemes):
+    ax = fig.add_subplot(1, 3, i + 1, projection="3d")
+    ax.scatter(theta[0], theta[1], theta[2], c=theta[2], cmap="viridis", s=4, alpha=0.8)
+    ax.set_title(title)
+    ax.set_box_aspect([1, 1, 1])
+    ax.view_init(elev=20, azim=45)
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_zticks([])
+
+pl.tight_layout()
+pl.show()
+
+# Notice how the uniform sample leaves visible gaps and clusters, while QSW
+# and RQSW spread the points much more evenly over the sphere -- this is
+# exactly the low-discrepancy property that reduces the error of the Sliced
+# Wasserstein estimate.
+
+##############################################################################
+# Convergence to the true Sliced Wasserstein distance
+# ------------------------------------------------------
+# We now compare how fast each sampling scheme converges to the *true*
+# SWD as the number of projections grows. To get a reference value with
+# **zero** approximation error -- not even from a finite number of
+# samples -- we build ``Xt`` as a pure translation of ``Xs`` by a fixed
+# vector :math:`\delta`: ``Xt = Xs + delta``.
+#
+# For a rigid translation, the classical 1D Wasserstein identity
+# :math:`W_2(\mu, \mu + c) = |c|` holds *exactly*, for any distribution
+# shape and any (even very small) sample size -- no law-of-large-numbers
+# argument, no Gaussian assumption, just an algebraic identity of optimal
+# transport on the line. Projected onto any direction :math:`\theta`, this
+# gives :math:`W_2(\theta_\# \mu, \theta_\# \nu) = |\theta^T \delta|`
+# exactly, and averaging the square over :math:`\theta` uniform on
+# :math:`S^{d-1}` gives the closed-form identity
+#
+# .. math::
+#     \mathcal{SWD}_2(\mu, \nu) = \frac{\|\delta\|}{\sqrt{d}}
+#
+# Because this holds regardless of ``Xs``'s shape or size, the *only*
+# remaining source of error in the experiment below is the number of
+# projections -- exactly the quantity we want to study.
+
+rng = np.random.RandomState(0)
+
+n_samples = 200
+delta = np.array([1.5, 1.0, -0.5])
+Xs = rng.uniform(-2, 2, (n_samples, d))
+Xt = Xs + delta
+
+# Exact reference: no approximation at all, at any cost.
+sw_true = np.linalg.norm(delta) / np.sqrt(d)
+
+n_proj_list = [10, 20, 50, 100, 200, 500]
+n_trials = 8
+
+errors_uniform = np.zeros((n_trials, len(n_proj_list)))
+errors_rqsw = np.zeros((n_trials, len(n_proj_list)))
+errors_qsw = np.zeros(len(n_proj_list))
+
+for j, n_proj in enumerate(n_proj_list):
+    for t in range(n_trials):
+        sw_uniform = ot.sliced_wasserstein_distance(
+            Xs, Xt, n_projections=n_proj, sampling_slices="uniform", seed=t
+        )
+        sw_rqsw = ot.sliced_wasserstein_distance(
+            Xs,
+            Xt,
+            n_projections=n_proj,
+            sampling_slices="randomized_spiral_qmc",
+            seed=t,
+        )
+        errors_uniform[t, j] = np.abs(sw_uniform - sw_true)
+        errors_rqsw[t, j] = np.abs(sw_rqsw - sw_true)
+
+    sw_qsw = ot.sliced_wasserstein_distance(
+        Xs, Xt, n_projections=n_proj, sampling_slices="spiral_qmc"
+    )
+    errors_qsw[j] = np.abs(sw_qsw - sw_true)
+
+mean_err_uniform = errors_uniform.mean(axis=0)
+std_err_uniform = errors_uniform.std(axis=0)
+mean_err_rqsw = errors_rqsw.mean(axis=0)
+std_err_rqsw = errors_rqsw.std(axis=0)
+
+pl.figure(2, figsize=(6, 5))
+pl.plot(n_proj_list, mean_err_uniform, "o-", label="Uniform (MC)")
+pl.fill_between(
+    n_proj_list,
+    mean_err_uniform - std_err_uniform,
+    mean_err_uniform + std_err_uniform,
+    alpha=0.3,
+)
+pl.plot(n_proj_list, mean_err_rqsw, "s-", label="RQSW")
+pl.fill_between(
+    n_proj_list,
+    mean_err_rqsw - std_err_rqsw,
+    mean_err_rqsw + std_err_rqsw,
+    alpha=0.3,
+)
+pl.plot(n_proj_list, errors_qsw, "^-", label="QSW (deterministic)")
+pl.xscale("log")
+pl.yscale("log")
+pl.xlabel("Number of projections")
+pl.ylabel("Absolute error to the true SWD")
+pl.title("Convergence of the Sliced Wasserstein estimate (3D)")
+pl.legend()
+pl.show()
+
+# QSW and RQSW reach a given accuracy with fewer projections than uniform
+# sampling, and RQSW keeps the estimator unbiased -- so it is a drop-in
+# replacement for uniform sampling in stochastic optimization settings
+# (e.g. Sliced Wasserstein gradient flows) where a deterministic QSW
+# estimate would not be appropriate.
+
+# %%

ot/backend.py

@@ -1225,6 +1225,26 @@ def slogdet(self, a):
         """
         raise NotImplementedError()
 
+    def sin(self, a):
+        r"""
+        Trigonometric sine, element-wise.
+
+        This function follows the api from :any:`numpy.sin`
+
+        See: https://numpy.org/doc/stable/reference/generated/numpy.sin.html
+        """
+        raise NotImplementedError()
+
+    def cos(self, a):
+        r"""
+        Trigonometric cosine, element-wise.
+
+        This function follows the api from :any:`numpy.cos`
+
+        See: https://numpy.org/doc/stable/reference/generated/numpy.cos.html
+        """
+        raise NotImplementedError()
+
 
 class NumpyBackend(Backend):
     """
@@ -1381,6 +1401,12 @@ def conj(self, a):
     def arccos(self, a):
         return np.arccos(a)
 
+    def sin(self, a):
+        return np.sin(a)
+
+    def cos(self, a):
+        return np.cos(a)
+
     def repeat(self, a, repeats, axis=None):
         return np.repeat(a, repeats, axis)
 
@@ -1810,6 +1836,12 @@ def conj(self, a):
     def arccos(self, a):
         return jnp.arccos(a)
 
+    def sin(self, a):
+        return jnp.sin(a)
+
+    def cos(self, a):
+        return jnp.cos(a)
+
     def repeat(self, a, repeats, axis=None):
         return jnp.repeat(a, repeats, axis)
 
@@ -2319,6 +2351,12 @@ def conj(self, a):
     def arccos(self, a):
         return torch.acos(a)
 
+    def sin(self, a):
+        return torch.sin(a)
+
+    def cos(self, a):
+        return torch.cos(a)
+
     def repeat(self, a, repeats, axis=None):
         return torch.repeat_interleave(a, repeats, dim=axis)
 
@@ -2831,6 +2869,12 @@ def conj(self, a):
     def arccos(self, a):
         return cp.arccos(a)
 
+    def sin(self, a):
+        return cp.sin(a)
+
+    def cos(self, a):
+        return cp.cos(a)
+
     def repeat(self, a, repeats, axis=None):
         return cp.repeat(a, repeats, axis)
 
@@ -3277,6 +3321,12 @@ def conj(self, a):
     def arccos(self, a):
         return tnp.arccos(a)
 
+    def sin(self, a):
+        return tnp.sin(a)
+
+    def cos(self, a):
+        return tnp.cos(a)
+
     def repeat(self, a, repeats, axis=None):
         return tnp.repeat(a, repeats, axis)
 

ot/sliced/__init__.py

@@ -15,6 +15,7 @@
     get_projections_sphere,
     get_random_rotations,
     projection_sphere_to_circle,
+    get_projections_spiral,
     projection_sphere_to_ball,
 )
 from ._sliced_distances import (
@@ -43,5 +44,6 @@
     "sliced_wasserstein_sphere",
     "sliced_wasserstein_sphere_unif",
     "linear_sliced_wasserstein_sphere",
+    "get_projections_spiral",
     "stereographic_sliced_wasserstein_sphere",
 ]