diff --git a/README.md b/README.md index 522a7c16d..d4e4653a4 100644 --- a/README.md +++ b/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. + diff --git a/RELEASES.md b/RELEASES.md index 6ead2ce34..72dc47677 100644 --- a/RELEASES.md +++ b/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 diff --git a/examples/sliced-wasserstein/plot_qsw_3d.py b/examples/sliced-wasserstein/plot_qsw_3d.py new file mode 100644 index 000000000..262f72286 --- /dev/null +++ b/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 +# +# 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. + +# %% diff --git a/ot/backend.py b/ot/backend.py index af622734d..1b7d24fd8 100644 --- a/ot/backend.py +++ b/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) diff --git a/ot/sliced/__init__.py b/ot/sliced/__init__.py index 84b4b58c2..439abd88d 100644 --- a/ot/sliced/__init__.py +++ b/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", ] diff --git a/ot/sliced/_sliced_distances.py b/ot/sliced/_sliced_distances.py index 1a6c300b2..c3574dd08 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 ``"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``. 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="spiral_qmc"`` (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). + - ``"spiral_qmc"``: deterministic Quasi-Sliced Wasserstein directions via + generalized spiral points. Only implemented for ``dim == 3``. + - ``"randomized_spiral_qmc"``: Randomized Quasi-Sliced Wasserstein -- the same spiral + point set as ``"spiral_qmc"``, with a random rotation applied, giving an + unbiased estimator suitable for stochastic optimization. Only + implemented for ``dim == 3``. Returns ------- @@ -94,6 +116,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. """ X_s, X_t = list_to_array(X_s, X_t) @@ -119,10 +143,28 @@ def sliced_wasserstein_distance( d = X_s.shape[1] + randomized = sampling_slices.startswith("randomized_") + method = sampling_slices.removeprefix("randomized_") + 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 method == "spiral_qmc": + projections = get_projections_spiral( + d, + n_projections, + randomized=randomized, + 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'" + ) else: n_projections = projections.shape[1] diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index cc78a5d7c..e5f926fda 100644 --- a/ot/sliced/_utils.py +++ b/ot/sliced/_utils.py @@ -235,6 +235,89 @@ 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) [96]. + + 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 + + References + ---------- + + .. [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. + """ + if d != 3: + raise ValueError( + f"Generalized spiral points are only defined for d=3, got d={d}." + ) + + if backend is None: + nx = NumpyBackend() + else: + nx = backend + + 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) + + def projection_sphere_to_ball(x, eps=1e-6, backend=None): r""" Projection of :math:`x\in S^{d-1}` on the unit ball of :math:`\mathbb{R}^{d-1}` with :math:`\frac{1}{\pi}h_1\circ\phi_\epsilon`. diff --git a/test/sliced/test_sliced_distances.py b/test/sliced/test_sliced_distances.py index 8c7998cc4..51e9cd34a 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,258 @@ 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", ["spiral_qmc", "randomized_spiral_qmc"]) +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", ["spiral_qmc", "randomized_spiral_qmc"]) +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", ["spiral_qmc", "randomized_spiral_qmc"]) +def test_sliced_qsw_wrong_dim_raises(sampling_slices): + """spiral_qmc/randomized_spiral_qmc 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", ["spiral_qmc", "randomized_spiral_qmc"]) +def test_sliced_qsw_backend(nx, sampling_slices): + """spiral_qmc/randomized_spiral_qmc 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 == "spiral_qmc": + # 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(nx): + """QSW (deterministic spiral points) must give EXACTLY the same + projections and the same SW value on ``nx`` as on NumPy -- 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)) + + val_np, log_np = ot.sliced_wasserstein_distance( + X_s, X_t, n_projections=n_projections, sampling_slices="spiral_qmc", log=True + ) + + X_s_b, X_t_b = nx.from_numpy(X_s, X_t) + val_b, log_b = ot.sliced_wasserstein_distance( + X_s_b, + X_t_b, + n_projections=n_projections, + sampling_slices="spiral_qmc", + log=True, + ) + + np.testing.assert_allclose( + nx.to_numpy(val_b), + val_np, + atol=1e-10, + err_msg=f"QSW result on '{nx.__name__}' does not match NumPy", + ) + np.testing.assert_allclose( + nx.to_numpy(log_b["projections"]), + log_np["projections"], + atol=1e-10, + err_msg=f"QSW projections on '{nx.__name__}' do not match NumPy", + ) + + +def test_rqsw_seed_does_not_match_across_backends(nx): + """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, as NumPy's. + + 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. NumPy is skipped here since it IS the + reference and would trivially match itself. + """ + if nx.__name__ == "numpy": + pytest.skip("NumPy is the reference backend; it trivially matches itself") + + 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)) + + val_np = ot.sliced_wasserstein_distance( + X_s, + X_t, + n_projections=n_projections, + seed=0, + sampling_slices="randomized_spiral_qmc", + ) + + 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="randomized_spiral_qmc", + ) + + assert not np.isclose(nx.to_numpy(val_b), val_np), ( + f"Expected RQSW to differ between numpy and '{nx.__name__}' with " + f"'the same' seed (different RNG algorithms), but they agreed: " + f"{val_np} vs {nx.to_numpy(val_b)}" + ) + + +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. + + Measured during development, averaged over 15 seeds: RQSW was ~502x + more accurate than uniform sampling, and deterministic QSW was even + closer to the reference. The thresholds below (both simply "better + than uniform") are set far below that measurement, leaving a + comfortable margin. + """ + 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="randomized_spiral_qmc", + ) + 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="spiral_qmc" + ) + qsw_error = abs(val_qsw - reference) + + mean_uniform_error = np.mean(uniform_errors) + mean_rqsw_error = np.mean(rqsw_errors) + + assert mean_rqsw_error < mean_uniform_error + assert qsw_error < mean_uniform_error diff --git a/test/test_backend.py b/test/test_backend.py index 4df918140..52188e4c6 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -177,6 +177,10 @@ def test_empty_backend(): nx.minimum(v, v) with pytest.raises(NotImplementedError): nx.abs(M) + with pytest.raises(NotImplementedError): + nx.sin(M) + with pytest.raises(NotImplementedError): + nx.cos(M) with pytest.raises(NotImplementedError): nx.log(M) with pytest.raises(NotImplementedError): @@ -445,6 +449,14 @@ def test_func_backends(nx): lst_b.append(nx.to_numpy(A)) lst_name.append("abs") + A = nx.sin(Mb) + lst_b.append(nx.to_numpy(A)) + lst_name.append("sin") + + A = nx.cos(Mb) + lst_b.append(nx.to_numpy(A)) + lst_name.append("cos") + A = nx.log(A) lst_b.append(nx.to_numpy(A)) lst_name.append("log")