From dcfe739935aa779c1b8a1a9fb308a5e53aca4524 Mon Sep 17 00:00:00 2001 From: Samuel Vangu Date: Tue, 25 Aug 2026 19:26:51 +0200 Subject: [PATCH 1/3] Add QSW sampling for sliced Wasserstein --- README.md | 2 + RELEASES.md | 9 + examples/sliced-wasserstein/plot_qsw_3d.py | 175 ++++++++++++++ ot/sliced/__init__.py | 2 + ot/sliced/_sliced_distances.py | 49 +++- ot/sliced/_utils.py | 81 +++++++ test/sliced/test_sliced_distances.py | 264 ++++++++++++++++++++- 7 files changed, 575 insertions(+), 7 deletions(-) create mode 100644 examples/sliced-wasserstein/plot_qsw_3d.py diff --git a/README.md b/README.md index 70092d9ad..e4ea3878a 100644 --- a/README.md +++ b/README.md @@ -480,3 +480,5 @@ Artificial Intelligence. \[92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). [A fast proximal point method for computing exact wasserstein distance.](https://proceedings.mlr.press/v115/xie20b/xie20b.pdf) In Uncertainty in artificial intelligence (pp. 433-453). PMLR. +\[93] 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). + diff --git a/RELEASES.md b/RELEASES.md index f91b42bc0..29fc8e6bd 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -1,4 +1,13 @@ # Releases +## 0.9.8dev +*August 2026* + +#### New features + +- Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized + spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`, + as described in [93] (PR #xxx) + ## 0.9.7.post1 diff --git a/examples/sliced-wasserstein/plot_qsw_3d.py b/examples/sliced-wasserstein/plot_qsw_3d.py new file mode 100644 index 000000000..266288118 --- /dev/null +++ b/examples/sliced-wasserstein/plot_qsw_3d.py @@ -0,0 +1,175 @@ +# -*- 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 [93], 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 a reference Sliced Wasserstein distance +between two 3D distributions. + +.. [93] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for + 3D Sliced Wasserstein. International Conference on Learning + Representations (ICLR). +""" + +# Author: Samuel Vangu +# +# 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 ~1000 directions on :math:`S^2` with each scheme: +# +# - ``uniform``: directions are Gaussian vectors normalized to unit norm +# (standard Monte Carlo sampling of the sphere). +# - ``qsw``: deterministic generalized spiral points -- a simple, +# closed-form low-discrepancy point set (Rakhmanov, Saff & Zhou, 1994). +# The same call always returns the same points. +# - ``rqsw``: 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 QSW set. + +n_projections = 1000 +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 a reference SWD +# value as the number of projections grows, on two 3D Gaussian samples. +# For a fixed number of projections, ``uniform`` and ``rqsw`` are random +# estimators (we average their error over several trials), while ``qsw`` is +# deterministic (a single value per number of projections). + +rng = np.random.RandomState(0) + +n_samples = 200 +Xs = rng.multivariate_normal([0, 0, 0], np.eye(3), n_samples) +Xt = rng.multivariate_normal([1.5, 1.0, -0.5], np.eye(3), n_samples) + +# "Ground truth" SWD, estimated with a very large number of uniform +# projections so that the residual Monte-Carlo error is negligible. +sw_true = ot.sliced_wasserstein_distance(Xs, Xt, n_projections=200000, seed=0) + +n_proj_list = [10, 20, 50, 100, 200, 500, 1000] +n_trials = 20 + +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="rqsw", 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="qsw" + ) + 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 reference 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. + +# %% diff --git a/ot/sliced/__init__.py b/ot/sliced/__init__.py index 709693bb7..bc1033ac9 100644 --- a/ot/sliced/__init__.py +++ b/ot/sliced/__init__.py @@ -14,6 +14,7 @@ get_random_projections, get_projections_sphere, projection_sphere_to_circle, + get_projections_spiral, ) from ._sliced_distances import ( sliced_wasserstein_distance, @@ -38,4 +39,5 @@ "sliced_wasserstein_sphere", "sliced_wasserstein_sphere_unif", "linear_sliced_wasserstein_sphere", + "get_projections_spiral", ] diff --git a/ot/sliced/_sliced_distances.py b/ot/sliced/_sliced_distances.py index 1a6c300b2..c9c859528 100644 --- a/ot/sliced/_sliced_distances.py +++ b/ot/sliced/_sliced_distances.py @@ -11,7 +11,7 @@ from ..backend import get_backend from ..utils import list_to_array, apply_scaler -from ._utils import get_random_projections +from ._utils import get_random_projections, get_projections_spiral from ..lp import wasserstein_1d @@ -26,6 +26,7 @@ def sliced_wasserstein_distance( seed=None, log=False, scaler=None, + sampling_slices="uniform", ): r""" Computes a Monte-Carlo approximation of the p-Sliced Wasserstein distance @@ -38,6 +39,12 @@ 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 ``"qsw"`` or ``"rqsw"`` instead + uses Quasi-Monte Carlo point sets on the sphere (generalized spiral + points), which can reduce the approximation error for a given + ``n_projections`` [93]. These two options are currently + only implemented for ``dim == 3``. Parameters ---------- @@ -54,9 +61,12 @@ def sliced_wasserstein_distance( p: float, optional Power p used for computing the sliced Wasserstein projections: shape (dim, n_projections), optional - Projection matrix (n_projections and seed are not used in this case) + Projection matrix (n_projections, seed and sampling_slices are not + used in this case) seed: int or RandomState or None, optional - Seed used for random number generator + Seed used for random number generator. Ignored if + ``sampling_slices="qsw"`` (the deterministic point set does not + depend on a seed). log: bool, optional if True, sliced_wasserstein_distance returns the projections used and their associated EMD. scaler: None, object with .transform(), or callable, optional @@ -73,6 +83,18 @@ def sliced_wasserstein_distance( See :class:`ot.utils.DataScaler` for a backend-aware scaler that supports joint fitting on multiple distributions. + sampling_slices: str, optional + Method used to sample the projection directions when ``projections`` + is not provided directly. One of: + + - ``"uniform"`` (default): directions sampled uniformly at random on + the sphere (Monte Carlo). + - ``"qsw"``: deterministic Quasi-Sliced Wasserstein directions via + generalized spiral points. Only implemented for ``dim == 3``. + - ``"rqsw"``: Randomized Quasi-Sliced Wasserstein -- the same spiral + point set as ``"qsw"``, with a random rotation applied, giving an + unbiased estimator suitable for stochastic optimization. Only + implemented for ``dim == 3``. Returns ------- @@ -94,6 +116,7 @@ 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 + .. [93] Nguyen, K., Bariletto, N., & Ho, N. (2024). "Quasi-Monte Carlo for 3D Sliced Wasserstein." International Conference on Learning Representations (ICLR). """ X_s, X_t = list_to_array(X_s, X_t) @@ -120,9 +143,23 @@ def sliced_wasserstein_distance( d = X_s.shape[1] if projections is None: - projections = get_random_projections( - d, n_projections, seed, backend=nx, type_as=X_s - ) + if sampling_slices == "uniform": + projections = get_random_projections( + d, n_projections, seed, backend=nx, type_as=X_s + ) + elif sampling_slices == "qsw": + projections = get_projections_spiral( + d, n_projections, randomized=False, backend=nx, type_as=X_s + ) + elif sampling_slices == "rqsw": + projections = get_projections_spiral( + d, n_projections, randomized=True, seed=seed, backend=nx, type_as=X_s + ) + else: + raise ValueError( + f"Unknown sampling_slices method '{sampling_slices}', " + "must be one of 'uniform', 'qsw', 'rqsw'" + ) else: n_projections = projections.shape[1] diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index fc280d52c..24bea73ea 100644 --- a/ot/sliced/_utils.py +++ b/ot/sliced/_utils.py @@ -167,3 +167,84 @@ def projection_sphere_to_circle( ) return Xp_coords, projections + + +def get_projections_spiral( + d, n_projections, randomized=True, seed=None, backend=None, type_as=None +): + r""" + Generates n_projections points on the sphere via generalized + spiral points (Rakhmanov, Saff & Zhou, 1994). + + Only implemented for d=3 (the 2-sphere :math:`S^2`). + + Parameters + ---------- + d : int + dimension of the space. Only d=3 is currently supported. + n_projections : int + number of samples requested + randomized : bool, optional + If True (default), applies a random (d, d) rotation to the + deterministic spiral point set (RQSW), giving an unbiased estimator + suitable for stochastic optimization. If False, returns the plain + deterministic point set (QSW). + seed: int or RandomState, optional + Seed used for the random rotation. Ignored if randomized=False. + backend: + Backend to use for random generation + type_as: type, optional + Type of the returned array + + Returns + ------- + out: ndarray, shape (d, n_projections) + The (optionally rotated) spiral points on the sphere + + Examples + -------- + >>> n_projections = 100 + >>> d = 3 + >>> projs = get_projections_spiral(d, n_projections, randomized=False) + >>> np.allclose(np.sum(np.square(projs), 0), 1.) # doctest: +NORMALIZE_WHITESPACE + True + >>> rprojs = get_projections_spiral(d, n_projections, randomized=True, seed=0) + >>> np.allclose(np.sum(np.square(rprojs), 0), 1.) # doctest: +NORMALIZE_WHITESPACE + True + + """ + if d != 3: + raise ValueError( + f"get_projections_spiral is only implemented for d=3, got d={d}" + ) + + if backend is None: + nx = NumpyBackend() + else: + nx = backend + + # sin/cos are not exposed by the backend abstraction (only arccos, atan2 + # exist), so the deterministic point construction is done in plain NumPy + # and converted to the target backend at the end. + i = np.arange(1, n_projections + 1) + z = 1 - (2 * i - 1) / n_projections + phi1 = np.arccos(z) + phi2 = (1.8 * np.sqrt(n_projections) * phi1) % (2 * np.pi) + theta_np = np.stack( + [np.sin(phi1) * np.cos(phi2), np.sin(phi1) * np.sin(phi2), z], axis=0 + ) # shape (d, n_projections) + theta = nx.from_numpy(theta_np, type_as=type_as) + + if not randomized: + return theta + + if isinstance(seed, np.random.RandomState) and str(nx) == "numpy": + Z = seed.randn(d, d) + else: + if seed is not None: + nx.seed(seed) + Z = nx.randn(d, d, type_as=type_as) + + Q, R = nx.qr(Z) + Q = Q * nx.sign(nx.diag(R))[None, :] + return nx.matmul(Q, theta) diff --git a/test/sliced/test_sliced_distances.py b/test/sliced/test_sliced_distances.py index 8c7998cc4..fe6c1e504 100644 --- a/test/sliced/test_sliced_distances.py +++ b/test/sliced/test_sliced_distances.py @@ -11,7 +11,7 @@ import pytest import ot -from ot.sliced import get_random_projections +from ot.sliced import get_random_projections, get_projections_spiral from ot.backend import tf, torch @@ -403,3 +403,265 @@ def test_sliced_wasserstein_scaler_backend(nx): ) np.testing.assert_allclose(nx.to_numpy(val_b_max), val_np_max, atol=1e-5) + + +def test_get_projections_spiral(): + """Mirrors test_get_random_projections: spiral points must lie on the unit sphere.""" + projections = get_projections_spiral(3, 50, randomized=False) + np.testing.assert_almost_equal(np.sum(projections**2, 0), 1.0) + + +def test_get_projections_spiral_randomized_preserves_sphere(): + """The random rotation applied for RQSW must still yield unit vectors.""" + projections = get_projections_spiral(3, 50, randomized=True, seed=0) + np.testing.assert_almost_equal(np.sum(projections**2, 0), 1.0) + + +def test_get_projections_spiral_wrong_dim_raises(): + """get_projections_spiral is only implemented for d=3.""" + with pytest.raises(ValueError): + get_projections_spiral(5, 50) + + +def test_get_projections_spiral_seed_reproducibility(): + """Same seed must give identical RQSW rotations; different seeds must differ.""" + p1 = get_projections_spiral(3, 50, randomized=True, seed=42) + p2 = get_projections_spiral(3, 50, randomized=True, seed=42) + p3 = get_projections_spiral(3, 50, randomized=True, seed=43) + np.testing.assert_allclose(p1, p2) + assert not np.allclose(p1, p3) + + +def test_get_projections_spiral_deterministic_ignores_seed(): + """QSW (randomized=False) must not depend on seed at all.""" + p1 = get_projections_spiral(3, 50, randomized=False, seed=0) + p2 = get_projections_spiral(3, 50, randomized=False, seed=999) + np.testing.assert_allclose(p1, p2) + + +@pytest.mark.parametrize("sampling_slices", ["qsw", "rqsw"]) +def test_sliced_qsw_same_dist(sampling_slices): + """Same distribution -> SWD approx 0, mirrors test_sliced_same_dist.""" + n = 100 + rng = np.random.RandomState(0) + x = rng.randn(n, 3) + u = ot.utils.unif(n) + + res = ot.sliced_wasserstein_distance( + x, x, u, u, 10, seed=0, sampling_slices=sampling_slices + ) + np.testing.assert_almost_equal(res, 0.0) + + +@pytest.mark.parametrize("sampling_slices", ["qsw", "rqsw"]) +def test_sliced_qsw_different_dists(sampling_slices): + """Different distributions -> SWD > 0, mirrors test_sliced_different_dists.""" + n = 100 + rng = np.random.RandomState(0) + x = rng.randn(n, 3) + y = rng.randn(n, 3) + 2.0 + u = ot.utils.unif(n) + + res = ot.sliced_wasserstein_distance( + x, y, u, u, 10, seed=0, sampling_slices=sampling_slices + ) + assert res > 0.0 + + +def test_sliced_invalid_sampling_slices(): + """An unknown sampling_slices string must raise a clear ValueError.""" + n = 20 + rng = np.random.RandomState(0) + x = rng.randn(n, 3) + with pytest.raises(ValueError, match="sampling_slices"): + ot.sliced_wasserstein_distance(x, x, sampling_slices="not_a_real_method") + + +@pytest.mark.parametrize("sampling_slices", ["qsw", "rqsw"]) +def test_sliced_qsw_wrong_dim_raises(sampling_slices): + """qsw/rqsw are only implemented for dim=3; other dims must raise.""" + n = 20 + rng = np.random.RandomState(0) + x = rng.randn(n, 5) + with pytest.raises(ValueError): + ot.sliced_wasserstein_distance(x, x, sampling_slices=sampling_slices) + + +@pytest.mark.parametrize("sampling_slices", ["qsw", "rqsw"]) +def test_sliced_qsw_backend(nx, sampling_slices): + """qsw/rqsw must work identically across backends, mirrors test_sliced_backend.""" + n = 100 + rng = np.random.RandomState(0) + x = rng.randn(n, 3) + y = rng.randn(2 * n, 3) + + xb, yb = nx.from_numpy(x, y) + + val = ot.sliced_wasserstein_distance( + xb, yb, n_projections=20, seed=0, sampling_slices=sampling_slices + ) + val2 = ot.sliced_wasserstein_distance( + xb, yb, n_projections=20, seed=0, sampling_slices=sampling_slices + ) + + assert nx.to_numpy(val) > 0 + if sampling_slices == "qsw": + # Deterministic: identical seed or not, result must be identical. + assert val == val2 + else: + # RQSW seeded: same seed -> same rotation -> same result. + assert val == val2 + + +def test_qsw_matches_across_backends(): + """QSW (deterministic spiral points) must give EXACTLY the same + projections and the same SW value across every installed backend -- + the point construction itself never touches any backend's RNG (see + get_projections_spiral), so this is a genuine equality, not just + "close enough". + """ + d = 3 + n_projections = 50 + rng = np.random.RandomState(0) + X_s = rng.normal(0, 1, (30, d)) + X_t = rng.normal(1, 1, (30, d)) + + backends = ot.backend.get_backend_list() + assert len(backends) >= 1 + + results = {} + projections_by_backend = {} + for nx in backends: + X_s_b, X_t_b = nx.from_numpy(X_s, X_t) + val, log = ot.sliced_wasserstein_distance( + X_s_b, + X_t_b, + n_projections=n_projections, + sampling_slices="qsw", + log=True, + ) + results[nx.__name__] = nx.to_numpy(val) + projections_by_backend[nx.__name__] = nx.to_numpy(log["projections"]) + + names = list(results.keys()) + reference = results[names[0]] + for name in names[1:]: + np.testing.assert_allclose( + results[name], + reference, + atol=1e-10, + err_msg=f"QSW result on '{name}' does not match '{names[0]}'", + ) + np.testing.assert_allclose( + projections_by_backend[name], + projections_by_backend[names[0]], + atol=1e-10, + err_msg=f"QSW projections on '{name}' do not match '{names[0]}'", + ) + + +def test_rqsw_seed_does_not_match_across_backends(): + """RQSW's random rotation is drawn from each backend's OWN native RNG + (numpy.random.RandomState, torch.Generator, jax's Threefry counter RNG, + tf.random.Generator -- four genuinely different algorithms). The "same" + integer seed therefore does NOT produce the same rotation, or the same + SW value, across backends. + + This is documented, expected behaviour (see the Sobol docstring's + identical caveat), not a bug -- this test exists to confirm that + behaviour explicitly, verified directly: + numpy.RandomState(0).randn(3,3) != torch.Generator().manual_seed(0).randn(3,3). + + Within a SINGLE backend, same-seed reproducibility is already covered + by test_sliced_qsw_backend. + """ + backends = ot.backend.get_backend_list() + if len(backends) < 2: + pytest.skip("Need at least two installed backends to compare") + + d = 3 + n_projections = 50 + rng = np.random.RandomState(0) + X_s = rng.normal(0, 1, (30, d)) + X_t = rng.normal(1, 1, (30, d)) + + results = {} + for nx in backends: + X_s_b, X_t_b = nx.from_numpy(X_s, X_t) + val = ot.sliced_wasserstein_distance( + X_s_b, + X_t_b, + n_projections=n_projections, + seed=0, + sampling_slices="rqsw", + ) + results[nx.__name__] = float(nx.to_numpy(val)) + + names = list(results.keys()) + # At least one pair of backends must differ -- if every backend somehow + # agreed, that would itself be suspicious given the RNGs are unrelated. + values = list(results.values()) + assert len(set(round(v, 8) for v in values)) > 1, ( + f"Expected RQSW to differ across backends with 'the same' seed " + f"(different RNG algorithms), but all backends agreed: {results}" + ) + + +def test_sliced_qsw_beats_uniform_d3(): + """QSW/RQSW should reduce the SW approximation error compared to uniform + random sampling, at equal n_projections, in d=3 -- the paper's central claim. + + Reference SW value computed with a very large number of uniformly-sampled + projections (assumed to have converged close to the true SW distance). + Errors for "uniform" are averaged over several seeds, since a single + draw can be misleading. QSW has no seed to average over (deterministic); RQSW + is also averaged over several seeds. + """ + d = 3 + n_pts = 200 + n_projections = 50 + n_trials = 15 + + rng = np.random.RandomState(0) + X_s = rng.normal(0, 1, (n_pts, d)) + X_t = rng.normal(1, 1, (n_pts, d)) + + # High-precision reference. + reference = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=20000, seed=0, sampling_slices="uniform" + ) + + uniform_errors = [] + rqsw_errors = [] + for seed in range(n_trials): + val_uniform = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=n_projections, seed=seed, sampling_slices="uniform" + ) + val_rqsw = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=n_projections, seed=seed, sampling_slices="rqsw" + ) + uniform_errors.append(abs(val_uniform - reference)) + rqsw_errors.append(abs(val_rqsw - reference)) + + val_qsw = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=n_projections, sampling_slices="qsw" + ) + qsw_error = abs(val_qsw - reference) + + mean_uniform_error = np.mean(uniform_errors) + mean_rqsw_error = np.mean(rqsw_errors) + + print(f"\n[DEBUG] mean uniform error={mean_uniform_error:.6e}") + print( + f"[DEBUG] mean rqsw error={mean_rqsw_error:.6e}, " + f"ratio={mean_uniform_error / mean_rqsw_error:.2f}x" + ) + print( + f"[DEBUG] qsw error={qsw_error:.6e}, " + f"ratio vs mean uniform={mean_uniform_error / qsw_error:.2f}x" + ) + + # Deliberately conservative until measured against the real backend -- + # see test docstring. + assert mean_rqsw_error < mean_uniform_error + assert qsw_error < mean_uniform_error From 57a5a199d982be05a2fad1e9930a365075141efd Mon Sep 17 00:00:00 2001 From: Samuel Vangu Date: Tue, 25 Aug 2026 19:37:17 +0200 Subject: [PATCH 2/3] Add the PR number in the RELEASES.mf file --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 29fc8e6bd..de2a3cc39 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -6,7 +6,7 @@ - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`, - as described in [93] (PR #xxx) + as described in [93] (PR #838) ## 0.9.7.post1 From a055d12b97594be84657fa7b2b139e4dab864843 Mon Sep 17 00:00:00 2001 From: Samuel Vangu Date: Tue, 25 Aug 2026 22:46:16 +0200 Subject: [PATCH 3/3] modifying the structure of the example file --- examples/sliced-wasserstein/plot_qsw_3d.py | 51 ++++++++++++++-------- 1 file changed, 34 insertions(+), 17 deletions(-) diff --git a/examples/sliced-wasserstein/plot_qsw_3d.py b/examples/sliced-wasserstein/plot_qsw_3d.py index 266288118..a1cc9c874 100644 --- a/examples/sliced-wasserstein/plot_qsw_3d.py +++ b/examples/sliced-wasserstein/plot_qsw_3d.py @@ -25,8 +25,9 @@ unbiasedness. We first visualize the three sampling schemes on the sphere, then measure -how fast each one converges to a reference Sliced Wasserstein distance -between two 3D distributions. +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. .. [93] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for 3D Sliced Wasserstein. International Conference on Learning @@ -49,7 +50,7 @@ ############################################################################## # Visualize the three sampling schemes on the sphere # ---------------------------------------------------- -# We draw ~1000 directions on :math:`S^2` with each scheme: +# 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). @@ -61,7 +62,7 @@ # The rotation makes the estimator unbiased while keeping the points # as evenly spread out as the deterministic QSW set. -n_projections = 1000 +n_projections = 500 d = 3 seed = 42 @@ -98,24 +99,40 @@ ############################################################################## # Convergence to the true Sliced Wasserstein distance # ------------------------------------------------------ -# We now compare how fast each sampling scheme converges to a reference SWD -# value as the number of projections grows, on two 3D Gaussian samples. -# For a fixed number of projections, ``uniform`` and ``rqsw`` are random -# estimators (we average their error over several trials), while ``qsw`` is -# deterministic (a single value per number of projections). +# 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 -Xs = rng.multivariate_normal([0, 0, 0], np.eye(3), n_samples) -Xt = rng.multivariate_normal([1.5, 1.0, -0.5], np.eye(3), n_samples) +delta = np.array([1.5, 1.0, -0.5]) +Xs = rng.uniform(-2, 2, (n_samples, d)) +Xt = Xs + delta -# "Ground truth" SWD, estimated with a very large number of uniform -# projections so that the residual Monte-Carlo error is negligible. -sw_true = ot.sliced_wasserstein_distance(Xs, Xt, n_projections=200000, seed=0) +# 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, 1000] -n_trials = 20 +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))) @@ -161,7 +178,7 @@ pl.xscale("log") pl.yscale("log") pl.xlabel("Number of projections") -pl.ylabel("Absolute error to reference SWD") +pl.ylabel("Absolute error to the true SWD") pl.title("Convergence of the Sliced Wasserstein estimate (3D)") pl.legend() pl.show()