|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +""" |
| 3 | +========================================================= |
| 4 | +UnifOrtho Sliced Wasserstein in high dimension |
| 5 | +========================================================= |
| 6 | +
|
| 7 | +This example illustrates the UnifOrtho sampling scheme for Sliced |
| 8 | +Wasserstein directions, introduced in [97] and recommended for large |
| 9 | +dimensions by a recent numerical and theoretical study [98], and compares |
| 10 | +it to the default uniform (Monte Carlo) sampling of slicing directions. |
| 11 | +
|
| 12 | +Sliced Wasserstein (SWD) approximates the Wasserstein distance by averaging |
| 13 | +1D Wasserstein distances over projections onto random directions |
| 14 | +:math:`\\theta` drawn uniformly on the sphere. By default these directions |
| 15 | +are sampled purely at random (Monte Carlo), which introduces some variance |
| 16 | +in the estimate for a given number of projections. |
| 17 | +
|
| 18 | +UnifOrtho takes a different route that works in *any* dimension: |
| 19 | +instead of drawing directions independently, it draws them in blocks of ``dim`` directions, |
| 20 | +each block is a random orthonormal basis (a draw from the Haar measure on the |
| 21 | +special orthogonal group :math:`\\mathrm{SO}(\\mathrm{dim})`). Directions within a block are |
| 22 | +therefore exactly mutually orthogonal, spreading them out much more evenly |
| 23 | +than independent draws would. |
| 24 | +
|
| 25 | +We first visualize this block structure on the ordinary 3D sphere, purely |
| 26 | +for intuition -- dimension 3 is precisely where QSW/RQSW should be |
| 27 | +preferred in practice, not UnifOrtho, as the convergence comparison right |
| 28 | +after the visualization confirms numerically. We then measure convergence |
| 29 | +to the true Sliced Wasserstein distance in a genuinely high dimension, |
| 30 | +where UnifOrtho is the recommended choice. |
| 31 | +
|
| 32 | +.. [97] Rowland, M., Hron, J., Tang, Y., Choromanski, K., Sarlos, T., & |
| 33 | + Weller, A. (2019). Orthogonal Estimation of Wasserstein Distances. |
| 34 | + Proceedings of the 22nd International Conference on Artificial |
| 35 | + Intelligence and Statistics (AISTATS), PMLR 89. |
| 36 | +.. [98] Petrovic, V., Bardenet, R., & Desolneux, A. (2026). Repulsive |
| 37 | + Monte Carlo on the sphere for the sliced Wasserstein distance. |
| 38 | + Transactions on Machine Learning Research. |
| 39 | +.. [99] Sisouk, K., Delon, J., & Tierny, J. (2025). A User's Guide to |
| 40 | + Sampling Strategies for Sliced Optimal Transport. Transactions on |
| 41 | + Machine Learning Research. |
| 42 | +""" |
| 43 | + |
| 44 | +# Author: Samuel Vangu <samuelvangu0@gmail.com> |
| 45 | +# |
| 46 | +# License: MIT License |
| 47 | + |
| 48 | +# sphinx_gallery_thumbnail_number = 1 |
| 49 | + |
| 50 | +import numpy as np |
| 51 | +import matplotlib.pylab as pl |
| 52 | +from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (registers the 3D projection) |
| 53 | + |
| 54 | +import ot |
| 55 | +from ot.sliced import get_random_projections, get_random_orthogonal_directions |
| 56 | + |
| 57 | +############################################################################## |
| 58 | +# Visualize the block structure on the sphere (d=3, for intuition only) |
| 59 | +# ----------------------------------------------------------------------- |
| 60 | +# We draw 15 directions on :math:`S^2` with each scheme: |
| 61 | +# |
| 62 | +# - ``uniform``: directions are Gaussian vectors normalized to unit norm |
| 63 | +# (standard Monte Carlo sampling of the sphere) -- no structure between |
| 64 | +# the points. |
| 65 | +# - ``unif_ortho``: 5 independent blocks of 3 mutually orthogonal |
| 66 | +# directions each. Each block is colored separately below, so that |
| 67 | +# same-colored points are exactly orthogonal to one another -- this is |
| 68 | +# the structure that is not visible in the uniform sample. |
| 69 | +# |
| 70 | +# A translucent sphere is drawn behind the points purely to make the |
| 71 | +# geometry easier to read. |
| 72 | +# |
| 73 | +# Dimension 3 is used here only because it is the one human beings can |
| 74 | +# actually look at. It is *not* the dimension UnifOrtho is recommended |
| 75 | +# for -- see the convergence experiment below. |
| 76 | + |
| 77 | +d = 3 |
| 78 | +n_blocks = 15 |
| 79 | +n_projections = n_blocks * d |
| 80 | +seed = 42 |
| 81 | + |
| 82 | +theta_uniform = get_random_projections(d, n_projections, seed=seed) |
| 83 | +theta_uniortho = get_random_orthogonal_directions(d, n_projections, seed=seed) |
| 84 | + |
| 85 | +# A plain unit sphere surface, drawn behind the scattered points below. |
| 86 | +u_sphere = np.linspace(0, 2 * np.pi, 40) |
| 87 | +v_sphere = np.linspace(0, np.pi, 20) |
| 88 | +sphere_x = np.outer(np.cos(u_sphere), np.sin(v_sphere)) |
| 89 | +sphere_y = np.outer(np.sin(u_sphere), np.sin(v_sphere)) |
| 90 | +sphere_z = np.outer(np.ones_like(u_sphere), np.cos(v_sphere)) |
| 91 | + |
| 92 | +fig = pl.figure(1, figsize=(10, 5)) |
| 93 | + |
| 94 | +ax1 = fig.add_subplot(1, 2, 1, projection="3d") |
| 95 | +ax1.scatter( |
| 96 | + theta_uniform[0], theta_uniform[1], theta_uniform[2], c="gray", s=25, alpha=0.8 |
| 97 | +) |
| 98 | +ax1.set_title("Uniform (Monte Carlo)") |
| 99 | + |
| 100 | +ax2 = fig.add_subplot(1, 2, 2, projection="3d") |
| 101 | +block_ids = np.repeat(np.arange(n_blocks), d) |
| 102 | +ax2.scatter( |
| 103 | + theta_uniortho[0], |
| 104 | + theta_uniortho[1], |
| 105 | + theta_uniortho[2], |
| 106 | + c=block_ids, |
| 107 | + cmap="tab10", |
| 108 | + s=25, |
| 109 | + alpha=0.9, |
| 110 | +) |
| 111 | +ax2.set_title("UnifOrtho (one color per orthogonal block)") |
| 112 | + |
| 113 | +for ax in (ax1, ax2): |
| 114 | + ax.plot_surface( |
| 115 | + sphere_x, sphere_y, sphere_z, color="lightgray", alpha=0.15, linewidth=0 |
| 116 | + ) |
| 117 | + ax.set_box_aspect([1, 1, 1]) |
| 118 | + ax.view_init(elev=20, azim=45) |
| 119 | + ax.set_xticks([]) |
| 120 | + ax.set_yticks([]) |
| 121 | + ax.set_zticks([]) |
| 122 | + |
| 123 | +pl.tight_layout() |
| 124 | +pl.show() |
| 125 | + |
| 126 | +# Every triple of same-colored points on the right is an exact orthogonal |
| 127 | +# basis of R^3 -- three mutually perpendicular directions. The uniform |
| 128 | +# sample on the left has no such guarantee: any two of its points can end |
| 129 | +# up arbitrarily close to each other. |
| 130 | + |
| 131 | +############################################################################## |
| 132 | +# Convergence at d=3: UnifOrtho is not the right tool here |
| 133 | +# ------------------------------------------------------------ |
| 134 | +# Before moving to high dimension, it is worth checking numerically that |
| 135 | +# UnifOrtho is indeed NOT the best choice at d=3 -- randomized spiral |
| 136 | +# points (RQSW, see the companion example "Quasi-Monte Carlo Sliced |
| 137 | +# Wasserstein in 3D") are. We use the same exact closed-form reference as |
| 138 | +# below (a pure translation). |
| 139 | + |
| 140 | +from ot.sliced import get_projections_spiral # noqa: E402 (kept local to this section) |
| 141 | + |
| 142 | +d3 = 3 |
| 143 | +rng3 = np.random.RandomState(1) |
| 144 | +delta3 = rng3.normal(size=d3) * 1.2 |
| 145 | +Xs3 = rng3.uniform(-2, 2, (100, d3)) |
| 146 | +Xt3 = Xs3 + delta3 |
| 147 | +sw_true3 = np.linalg.norm(delta3) / np.sqrt(d3) |
| 148 | + |
| 149 | +n_proj_list3 = [10, 20, 50, 100] |
| 150 | +n_trials3 = 10 |
| 151 | + |
| 152 | +errors3 = {"uniform": [], "unif_ortho": [], "randomized_spiral_qmc": []} |
| 153 | +for n_proj in n_proj_list3: |
| 154 | + row = {k: [] for k in errors3} |
| 155 | + for t in range(n_trials3): |
| 156 | + for method in errors3: |
| 157 | + val = ot.sliced_wasserstein_distance( |
| 158 | + Xs3, Xt3, n_projections=n_proj, sampling_slices=method, seed=t |
| 159 | + ) |
| 160 | + row[method].append(abs(val - sw_true3)) |
| 161 | + for method in errors3: |
| 162 | + errors3[method].append(np.mean(row[method])) |
| 163 | + |
| 164 | +pl.figure(2, figsize=(6, 5)) |
| 165 | +for method, marker, label in [ |
| 166 | + ("uniform", "o-", "Uniform (MC)"), |
| 167 | + ("unif_ortho", "s-", "UnifOrtho"), |
| 168 | + ("randomized_spiral_qmc", "^-", "RQSW"), |
| 169 | +]: |
| 170 | + pl.plot(n_proj_list3, errors3[method], marker, label=label) |
| 171 | +pl.xscale("log") |
| 172 | +pl.yscale("log") |
| 173 | +pl.xlabel("Number of projections") |
| 174 | +pl.ylabel("Absolute error to the true SWD") |
| 175 | +pl.title("Convergence of the Sliced Wasserstein estimate (d=3)") |
| 176 | +pl.legend() |
| 177 | +pl.show() |
| 178 | + |
| 179 | +# RQSW is consistently the most accurate here, UnifOrtho a clear second |
| 180 | +# (still better than plain uniform sampling, but not the recommended |
| 181 | +# choice), and uniform sampling the least accurate -- exactly the |
| 182 | +# low-dimensional ordering the literature [99] describes. |
| 183 | + |
| 184 | +############################################################################## |
| 185 | +# Convergence to the true Sliced Wasserstein distance, in high dimension |
| 186 | +# -------------------------------------------------------------------------- |
| 187 | +# As with the QSW/RQSW example, we build ``Xt`` as a pure translation of |
| 188 | +# ``Xs`` by a fixed vector :math:`\delta`, which makes the true Sliced |
| 189 | +# Wasserstein distance known exactly, with zero approximation error left |
| 190 | +# except from the number of projections: |
| 191 | +# |
| 192 | +# .. math:: |
| 193 | +# \mathcal{SWD}_2(\mu, \nu) = \frac{\|\delta\|}{\sqrt{d}} |
| 194 | +# |
| 195 | +# This time we work in dimension 30 -- the regime UnifOrtho is recommended |
| 196 | +# for . The list of projection counts below is deliberately chosen to never be a multiple of ``d``: |
| 197 | +# when ``n_projections`` is an exact multiple of 30, UnifOrtho draws a |
| 198 | +# whole number of *complete* orthogonal bases, which for this particular |
| 199 | +# translation experiment happens to recover the exact answer up to |
| 200 | +# floating-point precision (an identity, not an approximation) -- an |
| 201 | +# interesting fact in its own right, but not representative of the |
| 202 | +# typical case we want to illustrate here. |
| 203 | + |
| 204 | +d = 30 |
| 205 | +rng = np.random.RandomState(0) |
| 206 | + |
| 207 | +n_samples = 200 |
| 208 | +delta = rng.normal(size=d) * 1.2 |
| 209 | +Xs = rng.uniform(-2, 2, (n_samples, d)) |
| 210 | +Xt = Xs + delta |
| 211 | + |
| 212 | +# Exact reference: no approximation at all, at any cost. |
| 213 | +sw_true = np.linalg.norm(delta) / np.sqrt(d) |
| 214 | + |
| 215 | +n_proj_list = [35, 65, 95, 190, 380, 760] # all >= d, none a multiple of d |
| 216 | +n_trials = 10 |
| 217 | + |
| 218 | +errors_uniform = np.zeros((n_trials, len(n_proj_list))) |
| 219 | +errors_uniortho = np.zeros((n_trials, len(n_proj_list))) |
| 220 | + |
| 221 | +for j, n_proj in enumerate(n_proj_list): |
| 222 | + for t in range(n_trials): |
| 223 | + sw_uniform = ot.sliced_wasserstein_distance( |
| 224 | + Xs, Xt, n_projections=n_proj, sampling_slices="uniform", seed=t |
| 225 | + ) |
| 226 | + sw_uniortho = ot.sliced_wasserstein_distance( |
| 227 | + Xs, Xt, n_projections=n_proj, sampling_slices="unif_ortho", seed=t |
| 228 | + ) |
| 229 | + errors_uniform[t, j] = np.abs(sw_uniform - sw_true) |
| 230 | + errors_uniortho[t, j] = np.abs(sw_uniortho - sw_true) |
| 231 | + |
| 232 | +mean_err_uniform = errors_uniform.mean(axis=0) |
| 233 | +std_err_uniform = errors_uniform.std(axis=0) |
| 234 | +mean_err_uniortho = errors_uniortho.mean(axis=0) |
| 235 | +std_err_uniortho = errors_uniortho.std(axis=0) |
| 236 | + |
| 237 | +pl.figure(3, figsize=(6, 5)) |
| 238 | +pl.plot(n_proj_list, mean_err_uniform, "o-", label="Uniform (MC)") |
| 239 | +pl.fill_between( |
| 240 | + n_proj_list, |
| 241 | + mean_err_uniform - std_err_uniform, |
| 242 | + mean_err_uniform + std_err_uniform, |
| 243 | + alpha=0.3, |
| 244 | +) |
| 245 | +pl.plot(n_proj_list, mean_err_uniortho, "s-", label="UnifOrtho") |
| 246 | +pl.fill_between( |
| 247 | + n_proj_list, |
| 248 | + mean_err_uniortho - std_err_uniortho, |
| 249 | + mean_err_uniortho + std_err_uniortho, |
| 250 | + alpha=0.3, |
| 251 | +) |
| 252 | +pl.xscale("log") |
| 253 | +pl.yscale("log") |
| 254 | +pl.xlabel("Number of projections") |
| 255 | +pl.ylabel("Absolute error to the true SWD") |
| 256 | +pl.title(f"Convergence of the Sliced Wasserstein estimate (d={d})") |
| 257 | +pl.legend() |
| 258 | +pl.show() |
| 259 | + |
| 260 | +# UnifOrtho consistently reaches a given accuracy with markedly fewer |
| 261 | +# projections than uniform sampling in this high-dimensional setting -- |
| 262 | +# the opposite of the low-dimensional case above, where RQSW is the |
| 263 | +# better choice. As a rule of thumb from the literature [98, 99]: prefer |
| 264 | +# RQSW in low dimension, UnifOrtho in high dimension, and either may do |
| 265 | +# in between. Since UnifOrtho remains an unbiased, stochastic estimator |
| 266 | +# (like RQSW), it is also a drop-in replacement for uniform sampling in |
| 267 | +# stochastic optimization settings. |
| 268 | + |
| 269 | +# %% |
0 commit comments