diff --git a/README.md b/README.md index d4e4653a4..759125506 100644 --- a/README.md +++ b/README.md @@ -488,3 +488,6 @@ 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. (2025). [Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance](https://arxiv.org/abs/2509.10166). arXiv:2509.10166. \ No newline at end of file diff --git a/RELEASES.md b/RELEASES.md index d2b3acdc1..6fd59cee4 100644 --- a/RELEASES.md +++ b/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 diff --git a/examples/sliced-wasserstein/plot_uniortho_high_dim.py b/examples/sliced-wasserstein/plot_uniortho_high_dim.py new file mode 100644 index 000000000..488230601 --- /dev/null +++ b/examples/sliced-wasserstein/plot_uniortho_high_dim.py @@ -0,0 +1,198 @@ +# -*- 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 +orthogonal group :math:`O(\\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. 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. (2025). Repulsive + Monte Carlo on the sphere for the sliced Wasserstein distance. + arXiv:2509.10166. +""" + +# 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_uniortho + +############################################################################## +# 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. +# +# 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_projections_uniortho(d, n_projections, seed=seed) + +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.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 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(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_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, where QSW/RQSW are the better choice. +# As a rule of thumb from the literature [98]: 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. + +# %% diff --git a/ot/sliced/__init__.py b/ot/sliced/__init__.py index 439abd88d..0b06c64c2 100644 --- a/ot/sliced/__init__.py +++ b/ot/sliced/__init__.py @@ -17,6 +17,7 @@ projection_sphere_to_circle, get_projections_spiral, projection_sphere_to_ball, + get_projections_uniortho, ) from ._sliced_distances import ( sliced_wasserstein_distance, @@ -46,4 +47,5 @@ "linear_sliced_wasserstein_sphere", "get_projections_spiral", "stereographic_sliced_wasserstein_sphere", + "get_projections_uniortho", ] diff --git a/ot/sliced/_sliced_distances.py b/ot/sliced/_sliced_distances.py index c3574dd08..42a9505d6 100644 --- a/ot/sliced/_sliced_distances.py +++ b/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_projections_uniortho, +) 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. + + A recent numerical and theoretical study [98] recommends + ``"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:`O(\mathrm{dim})`. Defined for any ``dim``, and + recommended in particular for large ``dim`` [98]. See + :any:`get_projections_uniortho` for details, including how + ``n_projections`` not being a multiple of ``dim`` is handled. Returns ------- @@ -118,6 +135,8 @@ 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. (2025). "Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance." arXiv:2509.10166. """ X_s, X_t = list_to_array(X_s, X_t) @@ -160,10 +179,18 @@ def sliced_wasserstein_distance( backend=nx, type_as=X_s, ) + elif method == "unif_ortho": + projections = get_projections_uniortho( + 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] diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index e5f926fda..0c85220bf 100644 --- a/ot/sliced/_utils.py +++ b/ot/sliced/_utils.py @@ -364,3 +364,91 @@ def projection_sphere_to_ball(x, eps=1e-6, backend=None): norm2 = nx.sum(x_azimuth**2, axis=-1, keepdims=True) radius = nx.arccos(-(3.0 + 5.0 * x_d) / (5.0 + 3.0 * x_d)) / np.pi return radius * x_azimuth / nx.sqrt(norm2) + + +def get_projections_uniortho(d, n_projections, seed=None, backend=None, type_as=None): + r""" + Generates n_projections directions on the sphere via UnifOrtho + (Rowland et al., 2019) [97]: independent blocks of d mutually + orthogonal, unit-norm directions, each block drawn from the Haar + measure on the orthogonal group :math:`O(d)`. + + Unlike the generalized spiral points (:any:`get_projections_spiral`), + which only cover :math:`d=3`, UnifOrtho is defined for any dimension + d, and is recommended specifically for large d: Petrovic, Bardenet & + Desolneux (2025) [98] show both empirically and theoretically that it + reduces the variance of the Sliced Wasserstein estimator in high + dimension, while quasi-Monte Carlo methods such as the spiral points + remain preferable in low dimension (:math:`d \in \{2, 3\}`). + + Parameters + ---------- + d : int + dimension of the space. Any d >= 1 is supported. + n_projections : int + number of samples requested + seed: int or RandomState, optional + Seed used for the underlying random rotations + backend: + Backend to use for random generation + type_as: type, optional + Type of the returned array + + Returns + ------- + out: ndarray, shape (d, n_projections) + The UnifOrtho directions on the sphere + + Notes + ----- + ``n_projections`` need not be a multiple of ``d``: internally, + ``ceil(n_projections / d)`` independent orthogonal blocks are drawn + and concatenated, then truncated to the requested length. Every + individual direction is still marginally uniform on the sphere, but + if the last block is truncated, the directions coming from that + specific block are no longer guaranteed to be mutually orthogonal + with each other. + + Examples + -------- + >>> n_projections = 100 + >>> d = 5 + >>> projs = get_projections_uniortho(d, n_projections, seed=0) + >>> np.allclose(np.sum(np.square(projs), 0), 1.) # doctest: +NORMALIZE_WHITESPACE + True + >>> first_block = projs[:, :d] + >>> np.allclose(first_block.T @ first_block, np.eye(d)) # doctest: +NORMALIZE_WHITESPACE + True + + References + ---------- + + .. [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. (2025). "Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance." arXiv:2509.10166. + """ + if backend is None: + nx = NumpyBackend() + else: + nx = backend + + # Number of independent orthogonal blocks needed so that k * d covers + # at least n_projections directions (integer ceiling of n_projections / d). + k = -(-n_projections // d) + + # k independent Haar-uniform (d, d) orthogonal matrices. Within each + # (d, d) slice, the ROWS form d mutually orthogonal, unit-norm + # directions (see get_random_rotations for the QR-based construction + # and its Mezzadri (2007) sign correction, which is what makes the + # rotation genuinely Haar-uniform rather than merely orthogonal). + rotations = get_random_rotations(d, k, seed=seed, backend=nx, type_as=type_as) + + # Flatten the k blocks into a single list of k * d directions, one per row. + directions = nx.reshape(rotations, (k * d, d)) + + # (d, k * d): match the (dim, n_projections) convention. + directions = nx.transpose(directions) + + # Trim down to exactly n_projections when n_projections is not a + # multiple of d (see Notes above for what this means for the last, + # possibly incomplete, orthogonal block). + return directions[:, :n_projections] diff --git a/test/sliced/test_sliced_distances.py b/test/sliced/test_sliced_distances.py index 51e9cd34a..3af0de2aa 100644 --- a/test/sliced/test_sliced_distances.py +++ b/test/sliced/test_sliced_distances.py @@ -11,7 +11,11 @@ import pytest import ot -from ot.sliced import get_random_projections, get_projections_spiral +from ot.sliced import ( + get_random_projections, + get_projections_spiral, + get_projections_uniortho, +) from ot.backend import tf, torch @@ -658,3 +662,231 @@ def test_sliced_qsw_beats_uniform_d3(): assert mean_rqsw_error < mean_uniform_error assert qsw_error < mean_uniform_error + + +# ============================================================================= +# UnifOrtho: get_projections_uniortho-specific tests, mirroring the +# get_projections_spiral section above where the same property applies, and +# adding dedicated tests for what is specific to UnifOrtho (orthogonality +# within a block, no dimension restriction, exact behaviour when +# n_projections is a multiple of d). +# ============================================================================= + + +def test_get_projections_uniortho(): + """UnifOrtho directions must lie on the unit sphere, + regardless of n_projections being a multiple of d.""" + projections = get_projections_uniortho(7, 50, seed=0) + np.testing.assert_almost_equal(np.sum(projections**2, 0), 1.0) + + +def test_get_projections_uniortho_seed_reproducibility(): + """Same seed must give identical rotations; different seeds must differ.""" + p1 = get_projections_uniortho(7, 50, seed=42) + p2 = get_projections_uniortho(7, 50, seed=42) + p3 = get_projections_uniortho(7, 50, seed=43) + np.testing.assert_allclose(p1, p2) + assert not np.allclose(p1, p3) + + +@pytest.mark.parametrize("d", [1, 2, 7, 13, 50]) +def test_get_projections_uniortho_works_for_any_dimension(d): + """UnifOrtho places no restriction on the dimension. + This is its main advantage over the spiral points.""" + projections = get_projections_uniortho(d, 20, seed=0) + assert projections.shape == (d, 20) + np.testing.assert_almost_equal(np.sum(projections**2, 0), 1.0) + + +def test_get_projections_uniortho_block_is_orthogonal(): + """The defining property of UnifOrtho : within one full block of d directions, + they must be EXACTLY mutually orthogonal (not just individually uniform on the + sphere, as plain i.i.d. sampling already gives).""" + d = 6 + projections = get_projections_uniortho(d, d, seed=0) # exactly one block + gram = projections.T @ projections + np.testing.assert_allclose(gram, np.eye(d), atol=1e-10) + + +def test_get_projections_uniortho_handles_non_multiple_of_d(): + """n_projections need not be a multiple of d (see Notes in the + docstring): the returned shape must still be exactly what was + requested, and every direction must still be unit-norm, even though + the last block is truncated and thus not fully orthogonal internally.""" + d, n_projections = 5, 13 # 13 is not a multiple of 5 + projections = get_projections_uniortho(d, n_projections, seed=0) + assert projections.shape == (d, n_projections) + np.testing.assert_almost_equal(np.sum(projections**2, 0), 1.0) + + +def test_get_projections_uniortho_exact_when_n_projections_equals_d(): + """When n_projections == d, UnifOrtho draws a single COMPLETE + orthonormal basis of R^d. For any fixed vector delta, summing the + squared coefficients of delta in a complete orthonormal basis recovers + ||delta||^2 exactly (Parseval), regardless of which orthonormal basis + is used. Verified here directly: this is not an approximation that + happens to be good, it is an exact identity up to floating point. + """ + d = 10 + rng = np.random.RandomState(0) + delta = rng.randn(d) + + projections = get_projections_uniortho(d, d, seed=1) + measured = np.mean((projections.T @ delta) ** 2) + exact = np.sum(delta**2) / d + + np.testing.assert_allclose(measured, exact, atol=1e-9) + + +@pytest.mark.parametrize("sampling_slices", ["unif_ortho"]) +def test_sliced_unif_ortho_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, 7) + u = ot.utils.unif(n) + + res = ot.sliced_wasserstein_distance( + x, x, u, u, 20, seed=0, sampling_slices=sampling_slices + ) + np.testing.assert_almost_equal(res, 0.0) + + +def test_sliced_unif_ortho_different_dists(): + """Different distributions -> SWD > 0, mirrors test_sliced_different_dists.""" + n = 100 + rng = np.random.RandomState(0) + x = rng.randn(n, 7) + y = rng.randn(n, 7) + 2.0 + u = ot.utils.unif(n) + + res = ot.sliced_wasserstein_distance( + x, y, u, u, 20, seed=0, sampling_slices="unif_ortho" + ) + assert res > 0.0 + + +def test_sliced_unif_ortho_ignores_randomized_prefix(): + """Documents current, intended behaviour of the generic + 'randomized_' prefix stripping used for sampling_slices: since + get_projections_uniortho has no randomized/deterministic distinction, + 'randomized_unif_ortho' is accepted and behaves exactly like + 'unif_ortho' (the randomized flag it would imply is simply unused). + This is not a bug, but it is worth pinning down explicitly so a future + refactor cannot silently change it without a test failing.""" + n = 30 + rng = np.random.RandomState(0) + x = rng.randn(n, 6) + y = rng.randn(n, 6) + 1.0 + + val_plain = ot.sliced_wasserstein_distance( + x, y, n_projections=18, seed=7, sampling_slices="unif_ortho" + ) + val_prefixed = ot.sliced_wasserstein_distance( + x, y, n_projections=18, seed=7, sampling_slices="randomized_unif_ortho" + ) + assert val_plain == val_prefixed + + +def test_sliced_unif_ortho_backend(nx): + """UnifOrtho must work identically across backends, mirrors test_sliced_backend.""" + n = 100 + rng = np.random.RandomState(0) + x = rng.randn(n, 7) + y = rng.randn(2 * n, 7) + + xb, yb = nx.from_numpy(x, y) + + val = ot.sliced_wasserstein_distance( + xb, yb, n_projections=20, seed=0, sampling_slices="unif_ortho" + ) + val2 = ot.sliced_wasserstein_distance( + xb, yb, n_projections=20, seed=0, sampling_slices="unif_ortho" + ) + + assert nx.to_numpy(val) > 0 + assert val == val2 + + +def test_unif_ortho_seed_does_not_match_across_backends(nx): + """Mirrors test_rqsw_seed_does_not_match_across_backends: UnifOrtho's + orthogonal blocks are built from each backend's OWN native Gaussian + draw (via get_random_rotations), so 'the same' integer seed does not + produce the same rotation, or the same SW value, across backends. + NumPy is skipped as the trivial reference, exactly as for RQSW.""" + if nx.__name__ == "numpy": + pytest.skip("NumPy is the reference backend; it trivially matches itself") + + d = 7 + n_projections = 20 + rng = np.random.RandomState(0) + X_s = rng.normal(0, 1, (30, d)) + X_t = rng.normal(1, 1, (30, d)) + + val_np = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=n_projections, seed=0, sampling_slices="unif_ortho" + ) + + X_s_b, X_t_b = nx.from_numpy(X_s, X_t) + val_b = ot.sliced_wasserstein_distance( + X_s_b, X_t_b, n_projections=n_projections, seed=0, sampling_slices="unif_ortho" + ) + + assert not np.isclose(nx.to_numpy(val_b), val_np), ( + f"Expected UnifOrtho to differ between numpy and '{nx.__name__}' " + f"with 'the same' seed (different RNG algorithms), but they " + f"agreed: {val_np} vs {nx.to_numpy(val_b)}" + ) + + +def test_sliced_unif_ortho_beats_uniform_high_dim(): + """UnifOrtho should reduce the SW approximation error compared to + uniform random sampling in HIGH dimension -- the regime it is + recommended for (unlike spiral_qmc/RQSW, tested for d=3 in + test_sliced_qsw_beats_uniform_d3; see get_projections_uniortho and + sliced_wasserstein_distance docstrings for the literature recommending + this dimension-dependent choice). + + Uses the exact same closed-form reference construction as + test_sliced_qsw_beats_uniform_d3 (a pure translation, whose SW is + known exactly, with zero finite-sample error). n_projections is + deliberately NOT a multiple of d, to avoid the degenerate case where + UnifOrtho draws exactly one complete orthonormal basis (see + test_get_projections_uniortho_exact_when_n_projections_equals_d), + which would make this test measure that special identity rather than + UnifOrtho's typical behaviour. + + Measured during development, averaged over 20 seeds at d=30, + n_projections=47: UnifOrtho was ~2.3x more accurate than uniform + sampling. The threshold below is set well below that measurement. + """ + d = 30 + n_projections = 47 # not a multiple of d, see docstring above + n_trials = 20 + + rng = np.random.RandomState(0) + delta = rng.normal(size=d) + X_s = rng.normal(0, 1, (200, d)) + X_t = X_s + delta + reference = np.linalg.norm(delta) / np.sqrt(d) + + uniform_errors = [] + uniortho_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_uniortho = ot.sliced_wasserstein_distance( + X_s, + X_t, + n_projections=n_projections, + seed=seed, + sampling_slices="unif_ortho", + ) + uniform_errors.append(abs(val_uniform - reference)) + uniortho_errors.append(abs(val_uniortho - reference)) + + mean_uniform_error = np.mean(uniform_errors) + mean_uniortho_error = np.mean(uniortho_errors) + + assert mean_uniortho_error < mean_uniform_error