[MRG]Add UnifOrtho sampling of slicing directions for Sliced Wasserstein

README.md

@@ -488,3 +488,8 @@ Artificial Intelligence.
 
 \[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.
 
+\[97] Rowland, M., Hron, J., Tang, Y., Choromanski, K., Sarlos, T., & Weller, A. (2019). [Orthogonal Estimation of Wasserstein Distances](https://proceedings.mlr.press/v89/rowland19a.html). Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS), PMLR 89:186-195.
+
+\[98] Petrovic, V., Bardenet, R., & Desolneux, A. (2026). [Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance](https://openreview.net/forum?id=JSiTmB6Ehu). Transactions on Machine Learning Research.
+
+\[99] Sisouk, K., Delon, J., & Tierny, J. (2025). [A User's Guide to Sampling Strategies for Sliced Optimal Transport](https://openreview.net/forum?id=ECBepTWAFG). Transactions on Machine Learning Research.

RELEASES.md

@@ -9,6 +9,9 @@
 - 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)
+- Add UnifOrtho sliced Wasserstein sampling via independent orthogonal
+  blocks, selectable with `sampling_slices` in `sliced_wasserstein_distance`,
+  as described in [97] (PR #853)
 
 #### Closed issues
 

examples/sliced-wasserstein/plot_uniortho_high_dim.py

@@ -0,0 +1,269 @@
+# -*- coding: utf-8 -*-
+"""
+=========================================================
+UnifOrtho Sliced Wasserstein in high dimension
+=========================================================
+
+This example illustrates the UnifOrtho sampling scheme for Sliced
+Wasserstein directions, introduced in [97] and recommended for large
+dimensions by a recent numerical and theoretical study [98], and compares
+it 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.
+
+UnifOrtho takes a different route that works in *any* dimension:
+instead of drawing directions independently, it draws them in blocks of ``dim`` directions,
+each block is a random orthonormal basis (a draw from the Haar measure on the
+special orthogonal group :math:`\\mathrm{SO}(\\mathrm{dim})`). Directions within a block are
+therefore exactly mutually orthogonal, spreading them out much more evenly
+than independent draws would.
+
+We first visualize this block structure on the ordinary 3D sphere, purely
+for intuition -- dimension 3 is precisely where QSW/RQSW should be
+preferred in practice, not UnifOrtho, as the convergence comparison right
+after the visualization confirms numerically. We then measure convergence
+to the true Sliced Wasserstein distance in a genuinely high dimension,
+where UnifOrtho is the recommended choice.
+
+.. [97] Rowland, M., Hron, J., Tang, Y., Choromanski, K., Sarlos, T., &
+    Weller, A. (2019). Orthogonal Estimation of Wasserstein Distances.
+    Proceedings of the 22nd International Conference on Artificial
+    Intelligence and Statistics (AISTATS), PMLR 89.
+.. [98] Petrovic, V., Bardenet, R., & Desolneux, A. (2026). Repulsive
+    Monte Carlo on the sphere for the sliced Wasserstein distance.
+    Transactions on Machine Learning Research.
+.. [99] Sisouk, K., Delon, J., & Tierny, J. (2025). A User's Guide to
+    Sampling Strategies for Sliced Optimal Transport. Transactions on
+    Machine Learning Research.
+"""
+
+# 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_random_orthogonal_directions
+
+##############################################################################
+# Visualize the block structure on the sphere (d=3, for intuition only)
+# -----------------------------------------------------------------------
+# We draw 15 directions on :math:`S^2` with each scheme:
+#
+# - ``uniform``: directions are Gaussian vectors normalized to unit norm
+#   (standard Monte Carlo sampling of the sphere) -- no structure between
+#   the points.
+# - ``unif_ortho``: 5 independent blocks of 3 mutually orthogonal
+#   directions each. Each block is colored separately below, so that
+#   same-colored points are exactly orthogonal to one another -- this is
+#   the structure that is not visible in the uniform sample.
+#
+# A translucent sphere is drawn behind the points purely to make the
+# geometry easier to read.
+#
+# Dimension 3 is used here only because it is the one human beings can
+# actually look at. It is *not* the dimension UnifOrtho is recommended
+# for -- see the convergence experiment below.
+
+d = 3
+n_blocks = 15
+n_projections = n_blocks * d
+seed = 42
+
+theta_uniform = get_random_projections(d, n_projections, seed=seed)
+theta_uniortho = get_random_orthogonal_directions(d, n_projections, seed=seed)
+
+# A plain unit sphere surface, drawn behind the scattered points below.
+u_sphere = np.linspace(0, 2 * np.pi, 40)
+v_sphere = np.linspace(0, np.pi, 20)
+sphere_x = np.outer(np.cos(u_sphere), np.sin(v_sphere))
+sphere_y = np.outer(np.sin(u_sphere), np.sin(v_sphere))
+sphere_z = np.outer(np.ones_like(u_sphere), np.cos(v_sphere))
+
+fig = pl.figure(1, figsize=(10, 5))
+
+ax1 = fig.add_subplot(1, 2, 1, projection="3d")
+ax1.scatter(
+    theta_uniform[0], theta_uniform[1], theta_uniform[2], c="gray", s=25, alpha=0.8
+)
+ax1.set_title("Uniform (Monte Carlo)")
+
+ax2 = fig.add_subplot(1, 2, 2, projection="3d")
+block_ids = np.repeat(np.arange(n_blocks), d)
+ax2.scatter(
+    theta_uniortho[0],
+    theta_uniortho[1],
+    theta_uniortho[2],
+    c=block_ids,
+    cmap="tab10",
+    s=25,
+    alpha=0.9,
+)
+ax2.set_title("UnifOrtho (one color per orthogonal block)")
+
+for ax in (ax1, ax2):
+    ax.plot_surface(
+        sphere_x, sphere_y, sphere_z, color="lightgray", alpha=0.15, linewidth=0
+    )
+    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()
+
+# Every triple of same-colored points on the right is an exact orthogonal
+# basis of R^3 -- three mutually perpendicular directions. The uniform
+# sample on the left has no such guarantee: any two of its points can end
+# up arbitrarily close to each other.
+
+##############################################################################
+# Convergence at d=3: UnifOrtho is not the right tool here
+# ------------------------------------------------------------
+# Before moving to high dimension, it is worth checking numerically that
+# UnifOrtho is indeed NOT the best choice at d=3 -- randomized spiral
+# points (RQSW, see the companion example "Quasi-Monte Carlo Sliced
+# Wasserstein in 3D") are. We use the same exact closed-form reference as
+# below (a pure translation).
+
+from ot.sliced import get_projections_spiral  # noqa: E402 (kept local to this section)
+
+d3 = 3
+rng3 = np.random.RandomState(1)
+delta3 = rng3.normal(size=d3) * 1.2
+Xs3 = rng3.uniform(-2, 2, (100, d3))
+Xt3 = Xs3 + delta3
+sw_true3 = np.linalg.norm(delta3) / np.sqrt(d3)
+
+n_proj_list3 = [10, 20, 50, 100]
+n_trials3 = 10
+
+errors3 = {"uniform": [], "unif_ortho": [], "randomized_spiral_qmc": []}
+for n_proj in n_proj_list3:
+    row = {k: [] for k in errors3}
+    for t in range(n_trials3):
+        for method in errors3:
+            val = ot.sliced_wasserstein_distance(
+                Xs3, Xt3, n_projections=n_proj, sampling_slices=method, seed=t
+            )
+            row[method].append(abs(val - sw_true3))
+    for method in errors3:
+        errors3[method].append(np.mean(row[method]))
+
+pl.figure(2, figsize=(6, 5))
+for method, marker, label in [
+    ("uniform", "o-", "Uniform (MC)"),
+    ("unif_ortho", "s-", "UnifOrtho"),
+    ("randomized_spiral_qmc", "^-", "RQSW"),
+]:
+    pl.plot(n_proj_list3, errors3[method], marker, label=label)
+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 (d=3)")
+pl.legend()
+pl.show()
+
+# RQSW is consistently the most accurate here, UnifOrtho a clear second
+# (still better than plain uniform sampling, but not the recommended
+# choice), and uniform sampling the least accurate -- exactly the
+# low-dimensional ordering the literature [99] describes.
+
+##############################################################################
+# Convergence to the true Sliced Wasserstein distance, in high dimension
+# --------------------------------------------------------------------------
+# As with the QSW/RQSW example, we build ``Xt`` as a pure translation of
+# ``Xs`` by a fixed vector :math:`\delta`, which makes the true Sliced
+# Wasserstein distance known exactly, with zero approximation error left
+# except from the number of projections:
+#
+# .. math::
+#     \mathcal{SWD}_2(\mu, \nu) = \frac{\|\delta\|}{\sqrt{d}}
+#
+# This time we work in dimension 30 -- the regime UnifOrtho is recommended
+# for . The list of projection counts below is deliberately chosen to never be a multiple of ``d``:
+# when ``n_projections`` is an exact multiple of 30, UnifOrtho draws a
+# whole number of *complete* orthogonal bases, which for this particular
+# translation experiment happens to recover the exact answer up to
+# floating-point precision (an identity, not an approximation) -- an
+# interesting fact in its own right, but not representative of the
+# typical case we want to illustrate here.
+
+d = 30
+rng = np.random.RandomState(0)
+
+n_samples = 200
+delta = rng.normal(size=d) * 1.2
+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 = [35, 65, 95, 190, 380, 760]  # all >= d, none a multiple of d
+n_trials = 10
+
+errors_uniform = np.zeros((n_trials, len(n_proj_list)))
+errors_uniortho = np.zeros((n_trials, 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_uniortho = ot.sliced_wasserstein_distance(
+            Xs, Xt, n_projections=n_proj, sampling_slices="unif_ortho", seed=t
+        )
+        errors_uniform[t, j] = np.abs(sw_uniform - sw_true)
+        errors_uniortho[t, j] = np.abs(sw_uniortho - sw_true)
+
+mean_err_uniform = errors_uniform.mean(axis=0)
+std_err_uniform = errors_uniform.std(axis=0)
+mean_err_uniortho = errors_uniortho.mean(axis=0)
+std_err_uniortho = errors_uniortho.std(axis=0)
+
+pl.figure(3, 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_uniortho, "s-", label="UnifOrtho")
+pl.fill_between(
+    n_proj_list,
+    mean_err_uniortho - std_err_uniortho,
+    mean_err_uniortho + std_err_uniortho,
+    alpha=0.3,
+)
+pl.xscale("log")
+pl.yscale("log")
+pl.xlabel("Number of projections")
+pl.ylabel("Absolute error to the true SWD")
+pl.title(f"Convergence of the Sliced Wasserstein estimate (d={d})")
+pl.legend()
+pl.show()
+
+# UnifOrtho consistently reaches a given accuracy with markedly fewer
+# projections than uniform sampling in this high-dimensional setting --
+# the opposite of the low-dimensional case above, where RQSW is the
+# better choice. As a rule of thumb from the literature [98, 99]: prefer
+# RQSW in low dimension, UnifOrtho in high dimension, and either may do
+# in between. Since UnifOrtho remains an unbiased, stochastic estimator
+# (like RQSW), it is also a drop-in replacement for uniform sampling in
+# stochastic optimization settings.
+
+# %%

ot/sliced/__init__.py

@@ -17,6 +17,7 @@
     projection_sphere_to_circle,
     get_projections_spiral,
     projection_sphere_to_ball,
+    get_random_orthogonal_directions,
 )
 from ._sliced_distances import (
     sliced_wasserstein_distance,
@@ -46,4 +47,5 @@
     "linear_sliced_wasserstein_sphere",
     "get_projections_spiral",
     "stereographic_sliced_wasserstein_sphere",
+    "get_random_orthogonal_directions",
 ]

ot/sliced/_sliced_distances.py

@@ -11,7 +11,11 @@
 
 from ..backend import get_backend
 from ..utils import list_to_array, apply_scaler
-from ._utils import get_random_projections, get_projections_spiral
+from ._utils import (
+    get_random_projections,
+    get_projections_spiral,
+    get_random_orthogonal_directions,
+)
 from ..lp import wasserstein_1d
 
 
@@ -40,11 +44,18 @@ def sliced_wasserstein_distance(
     - :math:`\theta_\# \mu` stands for the pushforwards of the projection :math:`X \in \mathbb{R}^d \mapsto \langle \theta, X \rangle`
 
     By default, the projection directions :math:`\theta` are sampled uniformly
-    at random. Setting ``sampling_slices`` to ``"spiral_qmc"`` or ``"randomized_spiral_qmc"`` instead
-    uses Quasi-Monte Carlo point sets on the sphere (generalized spiral
-    points), which can reduce the approximation error for a given
-    ``n_projections`` [95]. These two options are
-    only implemented for ``dim == 3``.
+    at random. Two families of alternatives are available through
+    ``sampling_slices``, each better suited to a different regime:
+
+    - ``"spiral_qmc"`` / ``"randomized_spiral_qmc"`` use a deterministic,
+      low-discrepancy point set on the sphere (generalized spiral points),
+      only defined for ``dim == 3`` [95].
+    - ``"unif_ortho"`` uses independent blocks of mutually orthogonal
+      directions (UnifOrtho), defined for any dimension.
+
+    Recent numerical and theoretical studies [98, 99] recommend
+    ``"randomized_spiral_qmc"`` in low dimensions and
+    ``"unif_ortho"`` for large ``dim``, with no clear winner in between.
 
     Parameters
     ----------
@@ -95,6 +106,12 @@ def sliced_wasserstein_distance(
           point set as ``"spiral_qmc"``, with a random rotation applied, giving an
           unbiased estimator suitable for stochastic optimization. Only
           implemented for ``dim == 3``.
+        - ``"unif_ortho"``: UnifOrtho [97] -- independent blocks of
+          mutually orthogonal directions, each block drawn from the Haar
+          measure on :math:`\mathrm{SO}(\mathrm{dim})`. Defined for any
+          ``dim``, and recommended in particular for large ``dim`` [98, 99].
+          See :any:`get_random_orthogonal_directions` for details, including
+          how ``n_projections`` not being a multiple of ``dim`` is handled.
 
     Returns
     -------
@@ -118,6 +135,9 @@ def sliced_wasserstein_distance(
     .. [31] Bonneel, Nicolas, et al. "Sliced and radon wasserstein barycenters of measures." Journal of Mathematical Imaging and Vision 51.1 (2015): 22-45
     .. [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.
+    .. [97] Rowland, M., Hron, J., Tang, Y., Choromanski, K., Sarlos, T., & Weller, A. (2019). "Orthogonal Estimation of Wasserstein Distances." Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS), PMLR 89.
+    .. [98] Petrovic, V., Bardenet, R., & Desolneux, A. (2026). "Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance." Transactions on Machine Learning Research.
+    .. [99] Sisouk, K., Delon, J., & Tierny, J. (2025). "A User's Guide to Sampling Strategies for Sliced Optimal Transport." Transactions on Machine Learning Research.
     """
 
     X_s, X_t = list_to_array(X_s, X_t)
@@ -160,10 +180,18 @@ def sliced_wasserstein_distance(
                 backend=nx,
                 type_as=X_s,
             )
+        elif method == "unif_ortho":
+            projections = get_random_orthogonal_directions(
+                d,
+                n_projections,
+                seed=seed,
+                backend=nx,
+                type_as=X_s,
+            )
         else:
             raise ValueError(
                 f"Unknown sampling_slices method '{sampling_slices}', "
-                "must be one of 'uniform', 'spiral_qmc', 'randomized_spiral_qmc'"
+                "must be one of 'uniform', 'spiral_qmc', 'randomized_spiral_qmc', 'unif_ortho' "
             )
     else:
         n_projections = projections.shape[1]