diff --git a/README.md b/README.md index 70092d9ad..522a7c16d 100644 --- a/README.md +++ b/README.md @@ -480,3 +480,7 @@ 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] Tran, H., Bai, Y., Kothapalli, A., Shahbazi, A., Liu, X., Diaz Martin, R., & Kolouri, S. (2024). [Stereographic Spherical Sliced Wasserstein Distances](https://proceedings.mlr.press/v235/tran24a.html). International Conference on Machine Learning. + +\[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. + diff --git a/RELEASES.md b/RELEASES.md index f91b42bc0..f1f19a919 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -1,5 +1,11 @@ # Releases +## 0.9.8dev + +#### New features + +- Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836) + ## 0.9.7.post1 This release is identical to 0.9.7 but will allow the upload of a source distribution to PyPI and release on conda-forge (that requires a source distribution). diff --git a/ot/__init__.py b/ot/__init__.py index bdd2aadd7..9b53c10aa 100644 --- a/ot/__init__.py +++ b/ot/__init__.py @@ -70,6 +70,7 @@ sliced_wasserstein_sphere, sliced_wasserstein_sphere_unif, linear_sliced_wasserstein_sphere, + stereographic_sliced_wasserstein_sphere, min_sliced_transport_plan, expected_sliced_plan, ) @@ -133,6 +134,7 @@ "unbalanced_sliced_ot", "sliced_unbalanced_ot", "linear_sliced_wasserstein_sphere", + "stereographic_sliced_wasserstein_sphere", "min_sliced_transport_plan", "expected_sliced_plan", "gromov_wasserstein", diff --git a/ot/sliced/__init__.py b/ot/sliced/__init__.py index 709693bb7..d8659f36b 100644 --- a/ot/sliced/__init__.py +++ b/ot/sliced/__init__.py @@ -13,6 +13,7 @@ from ._utils import ( get_random_projections, get_projections_sphere, + get_random_rotations, projection_sphere_to_circle, ) from ._sliced_distances import ( @@ -23,12 +24,14 @@ sliced_wasserstein_sphere, sliced_wasserstein_sphere_unif, linear_sliced_wasserstein_sphere, + stereographic_sliced_wasserstein_sphere, ) from ._sliced_plans import min_sliced_transport_plan, expected_sliced_plan, sliced_plans __all__ = [ "get_random_projections", "get_projections_sphere", + "get_random_rotations", "projection_sphere_to_circle", "min_sliced_transport_plan", "expected_sliced_plan", @@ -38,4 +41,5 @@ "sliced_wasserstein_sphere", "sliced_wasserstein_sphere_unif", "linear_sliced_wasserstein_sphere", + "stereographic_sliced_wasserstein_sphere", ] diff --git a/ot/sliced/_spherical_sliced.py b/ot/sliced/_spherical_sliced.py index 2deaa5d2e..b78781a37 100644 --- a/ot/sliced/_spherical_sliced.py +++ b/ot/sliced/_spherical_sliced.py @@ -5,15 +5,24 @@ # Author: Nicolas Courty # Author: Clément Bonet +# Author: continuousml # # License: MIT License +import numpy as np + from ..backend import get_backend -from ._utils import get_projections_sphere, projection_sphere_to_circle +from ._utils import ( + get_projections_sphere, + get_random_projections, + get_random_rotations, + projection_sphere_to_circle, +) from ..lp import ( wasserstein_circle, semidiscrete_wasserstein2_unif_circle, linear_circular_ot, + wasserstein_1d, ) @@ -301,3 +310,187 @@ def linear_sliced_wasserstein_sphere( if log: return res, {"projections": projections, "projected_emds": projected_lcot} return res + + +def _projection_sphere_stereographic(x, eps, nx): + r"""Maps points on :math:`S^{d-1}` to :math:`\mathbb{R}^{d-1}` with :math:`h\circ\phi_\epsilon`, + where :math:`\phi_\epsilon` is the stereographic projection with an :math:`\epsilon`-cap + around the north pole and :math:`h` the defining function of :ref:`[93] `, + using the closed form :math:`(h\circ\phi_\epsilon)(x) = \mathrm{arccos}(-x_d)\frac{x_{1:d-1}}{\|x_{1:d-1}\|}`. + """ + x_d = nx.clip(x[..., -1:], -1.0, 1.0 - eps) + x_azimuth = x[..., :-1] + norm2 = nx.sum(x_azimuth**2, axis=-1, keepdims=True) + # the azimuth of the poles is arbitrary, fix it for reproducibility + x_azimuth = nx.where( + norm2 > 0, x_azimuth, nx.ones(x_azimuth.shape, type_as=x_azimuth) + ) + norm2 = nx.sum(x_azimuth**2, axis=-1, keepdims=True) + return nx.arccos(-x_d) * x_azimuth / nx.sqrt(norm2) + + +def stereographic_sliced_wasserstein_sphere( + X_s, + X_t, + a=None, + b=None, + n_projections=50, + p=2, + projections=None, + n_rotations=0, + rotations=None, + eps=1e-6, + seed=None, + log=False, +): + r"""Computes the stereographic spherical sliced Wasserstein distance from :ref:`[93] `. + + General loss returned: + + .. math:: + S3W_p(\mu,\nu) = \left(\int_{\mathbb{S}^{d-2}} W_p^p(\theta_\# (h\circ\phi_\epsilon)_\#\mu, \theta_\# (h\circ\phi_\epsilon)_\#\nu)\ \mathrm{d}\sigma(\theta)\right)^{\frac{1}{p}} + + where :math:`\mu,\nu\in\mathcal{P}(S^{d-1})` are two probability measures on the + sphere, :math:`\theta_\# \mu` stands for the pushforwards of the projection + :math:`X \in \mathbb{R}^{d-1} \mapsto \langle \theta, X \rangle`, + :math:`\phi_\epsilon` is the stereographic projection + :math:`\phi(x) = \frac{x_{1:d-1}}{1-x_d}` restricted to the sphere without the + :math:`\epsilon`-cap around the north pole (points with :math:`x_d > 1-\epsilon` + are first mapped to the circle :math:`x_d = 1-\epsilon`), and + :math:`h(x) = \mathrm{arccos}\left(\frac{1-\|x\|^2}{1+\|x\|^2}\right)\frac{x}{\|x\|}` + is the injective defining function of :ref:`[93] `, such that + :math:`(h\circ\phi_\epsilon)(x) = \mathrm{arccos}(-x_d)\frac{x_{1:d-1}}{\|x_{1:d-1}\|}` + maps each point to its geodesic distance to the south pole times its azimuth. + + If ``n_rotations >= 1`` or ``rotations`` is provided, computes instead a + Monte-Carlo approximation of the rotationally invariant extension + + .. math:: + RI\text{-}S3W_p(\mu,\nu) = \int_{\mathrm{SO}(d)} S3W_p(R_\#\mu, R_\#\nu)\ \mathrm{d}\omega(R) + + where :math:`\omega` is the normalized Haar measure on :math:`\mathrm{SO}(d)`. + The generation cost of the rotations can be amortized over several calls by + pregenerating a pool of rotations with + :any:`ot.sliced.get_random_rotations` and passing a random subset of it as + ``rotations`` at each call (ARI-S3W :ref:`[93] `). + + Parameters + ---------- + X_s: ndarray, shape (n_samples_a, dim) + Samples in the source domain + X_t: ndarray, shape (n_samples_b, dim) + Samples in the target domain + a : ndarray, shape (n_samples_a,), optional + samples weights in the source domain + b : ndarray, shape (n_samples_b,), optional + samples weights in the target domain + n_projections : int, optional + Number of projections used for the Monte-Carlo approximation + p: float, optional (default=2) + Power p used for computing the stereographic spherical sliced Wasserstein + projections: shape (dim-1, n_projections), optional + Projection matrix (n_projections and seed are not used in this case) + n_rotations : int, optional (default=0) + Number of rotations used for the Monte-Carlo approximation of + :math:`RI\text{-}S3W_p`. If 0, no rotation is applied and + :math:`S3W_p` is computed. + rotations: shape (n_rotations, dim, dim), optional + Rotation matrices (n_rotations is not used in this case) + eps: float, optional (default=1e-6) + Size of the cap around the north pole excluded from the stereographic + projection to ensure numerical stability + seed: int or RandomState or None, optional + Seed used for random number generator + log: bool, optional + if True, stereographic_sliced_wasserstein_sphere returns the projections + and rotations used and the associated EMDs. + + Returns + ------- + cost: float + Stereographic Spherical Sliced Wasserstein Cost + log: dict, optional + log dictionary return only if log==True in parameters + + Examples + -------- + >>> import ot + >>> import numpy as np + >>> n_samples_a = 20 + >>> X = np.random.normal(0., 1., (n_samples_a, 5)) + >>> X = X / np.sqrt(np.sum(X**2, -1, keepdims=True)) + >>> ot.stereographic_sliced_wasserstein_sphere(X, X, seed=0) # doctest: +NORMALIZE_WHITESPACE + 0.0 + + + .. _references-s3w: + References + ---------- + .. [93] Tran, H., Bai, Y., Kothapalli, A., Shahbazi, A., Liu, X., + Diaz Martin, R., & Kolouri, S. (2024). Stereographic Spherical Sliced + Wasserstein Distances. International Conference on Machine Learning. + """ + d = X_s.shape[-1] + + nx = get_backend(X_s, X_t, a, b, projections, rotations) + + if X_s.shape[1] != X_t.shape[1]: + raise ValueError( + "X_s and X_t must have the same number of dimensions {} and {} respectively given".format( + X_s.shape[1], X_t.shape[1] + ) + ) + if nx.any(nx.abs(nx.sum(X_s**2, axis=-1) - 1) > 10 ** (-4)): + raise ValueError("X_s is not on the sphere.") + if nx.any(nx.abs(nx.sum(X_t**2, axis=-1) - 1) > 10 ** (-4)): + raise ValueError("X_t is not on the sphere.") + + if projections is None: + projections = get_random_projections( + d - 1, n_projections, seed=seed, backend=nx, type_as=X_s + ) + if seed is not None and not isinstance(seed, np.random.RandomState): + # draw the rotations from the stream advanced by the projections + seed = None + else: + n_projections = projections.shape[1] + + if rotations is None and n_rotations > 0: + rotations = get_random_rotations( + d, n_rotations, seed=seed, backend=nx, type_as=X_s + ) + elif rotations is not None: + n_rotations = rotations.shape[0] + + if rotations is not None: + Xps = nx.einsum("kij, nj -> kni", rotations, X_s) + Xpt = nx.einsum("kij, nj -> kni", rotations, X_t) + else: + n_rotations = 1 + Xps = X_s[None, :, :] + Xpt = X_t[None, :, :] + + Xps = _projection_sphere_stereographic(Xps, eps, nx) + Xpt = _projection_sphere_stereographic(Xpt, eps, nx) + + Xps = nx.reshape( + nx.einsum("kni, il -> nkl", Xps, projections), + (X_s.shape[0], n_rotations * n_projections), + ) + Xpt = nx.reshape( + nx.einsum("kni, il -> nkl", Xpt, projections), + (X_t.shape[0], n_rotations * n_projections), + ) + + projected_emd = nx.reshape( + wasserstein_1d(Xps, Xpt, a, b, p=p), (n_rotations, n_projections) + ) + res = nx.mean(nx.mean(projected_emd, axis=-1) ** (1.0 / p)) + + if log: + return res, { + "projections": projections, + "rotations": rotations, + "projected_emds": projected_emd, + } + return res diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index fc280d52c..968386ee6 100644 --- a/ot/sliced/_utils.py +++ b/ot/sliced/_utils.py @@ -110,6 +110,72 @@ def get_projections_sphere(d, n_projections, seed=None, backend=None, type_as=No return projections +def get_random_rotations(d, n_rotations, seed=None, backend=None, type_as=None): + r""" + Generates n_rotations samples from the uniform (Haar) distribution on the special orthogonal group :math:`\mathrm{SO}(d)=\{R \in \mathbb{R}^{d\times d}, R^TR=I_d, \mathrm{det}(R)=1\}` + + The rotations are obtained from the QR factorization of Gaussian matrices, + with the sign correction of :ref:`[94] ` + and a sign flip of the first column of the matrices with negative + determinant. + + Parameters + ---------- + d : int + dimension of the space + n_rotations : int + number of samples requested + seed: int or RandomState, optional + Seed used for numpy random number generator + backend: + Backend to use for random generation + type_as: optional + Type to use for random generation + + Returns + ------- + out: ndarray, shape (n_rotations, d, d) + + Examples + -------- + >>> n_rotations = 100 + >>> d = 5 + >>> rotations = get_random_rotations(d, n_rotations) + >>> np.allclose(np.einsum("nij, nkj -> nik", rotations, rotations), np.eye(d)) # doctest: +NORMALIZE_WHITESPACE + True + >>> np.allclose(np.linalg.det(rotations), 1.) # doctest: +NORMALIZE_WHITESPACE + True + + + .. _references-get-random-rotations: + References + ---------- + .. [94] Mezzadri, F. (2007). How to generate random matrices from the + classical compact groups. Notices of the American Mathematical + Society, 54(5), 592-604. + """ + if backend is None: + nx = NumpyBackend() + else: + nx = backend + + if isinstance(seed, np.random.RandomState) and str(nx) == "numpy": + Z = seed.randn(n_rotations, d, d) + else: + if seed is not None: + nx.seed(seed) + Z = nx.randn(n_rotations, d, d, type_as=type_as) + + Q, R = nx.qr(Z) + diagonal = nx.sum(R * nx.eye(d, type_as=R)[None, :, :], axis=-1) + Q = Q * nx.sign(diagonal)[:, None, :] + flip = nx.sign(nx.det(Q)) + rotations = nx.concatenate( + (Q[:, :, :1] * flip[:, None, None], Q[:, :, 1:]), axis=-1 + ) + return rotations + + def projection_sphere_to_circle( x, n_projections=50, projections=None, seed=None, backend=None ): diff --git a/test/sliced/test_spherical_sliced.py b/test/sliced/test_spherical_sliced.py index 7b1f4a8b8..bcadff7b0 100644 --- a/test/sliced/test_spherical_sliced.py +++ b/test/sliced/test_spherical_sliced.py @@ -438,3 +438,226 @@ def test_linear_sliced_sphere_backend_type_devices(nx): nx.assert_same_dtype_device(xb, valb) np.testing.assert_almost_equal(sw_np, nx.to_numpy(valb)) + + +def test_get_random_rotations(): + rng = np.random.RandomState(0) + + n_rotations = 100 + rotations = ot.sliced.get_random_rotations(3, n_rotations, seed=rng) + + np.testing.assert_almost_equal( + np.matmul(rotations, np.transpose(rotations, [0, 2, 1])), + np.array([np.eye(3) for k in range(n_rotations)]), + ) + np.testing.assert_almost_equal(np.linalg.det(rotations), np.ones(n_rotations)) + + +def test_stereographic_sliced_sphere_same_dist(): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + u = ot.utils.unif(n) + + res = ot.stereographic_sliced_wasserstein_sphere(x, x, u, u, 10, seed=rng) + np.testing.assert_almost_equal(res, 0.0) + + +def test_stereographic_sliced_sphere_same_proj(): + n_projections = 10 + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + y = rng.randn(n, 3) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + + seed = 42 + + cost1, log1 = ot.stereographic_sliced_wasserstein_sphere( + x, y, seed=seed, n_projections=n_projections, n_rotations=3, log=True + ) + cost2, log2 = ot.stereographic_sliced_wasserstein_sphere( + x, y, seed=seed, n_projections=n_projections, n_rotations=3, log=True + ) + + assert np.allclose(log1["projections"], log2["projections"]) + assert np.allclose(log1["rotations"], log2["rotations"]) + assert np.isclose(cost1, cost2) + + +def test_stereographic_sliced_sphere_bad_shapes(): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + y = rng.randn(n, 4) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + + u = ot.utils.unif(n) + + with pytest.raises(ValueError): + _ = ot.stereographic_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng) + + +def test_stereographic_sliced_sphere_values_on_the_sphere(): + n = 100 + rng = np.random.RandomState(0) + + u = ot.utils.unif(n) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + # not on the sphere + y = rng.randn(n, 3) + + with pytest.raises(ValueError): + _ = ot.stereographic_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng) + + with pytest.raises(ValueError): + _ = ot.stereographic_sliced_wasserstein_sphere(y, x, u, u, 10, seed=rng) + + +def test_stereographic_sliced_sphere_log(): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 4) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + y = rng.randn(n, 4) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + u = ot.utils.unif(n) + + res, log = ot.stereographic_sliced_wasserstein_sphere( + x, y, u, u, 10, p=1, n_rotations=3, seed=rng, log=True + ) + assert len(log) == 3 + projections = log["projections"] + rotations = log["rotations"] + projected_emds = log["projected_emds"] + + assert projections.shape[1] == projected_emds.shape[1] == 10 + assert rotations.shape[0] == projected_emds.shape[0] == 3 + assert np.all(projected_emds > 0) + + +def test_stereographic_sliced_sphere_different_dists(): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + u = ot.utils.unif(n) + y = rng.randn(n, 3) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + + res = ot.stereographic_sliced_wasserstein_sphere(x, y, u, u, 10, seed=rng) + assert res > 0.0 + + +def test_stereographic_sliced_sphere_geodesic(): + # Diracs on a common meridian: h projects them on a common line at their + # geodesic distance to the south pole, so S3W_2 with the canonical basis + # as projections has a closed form + d = 4 + t1, t2 = 0.3, 1.9 + x = np.array([[np.sin(t1), 0.0, 0.0, np.cos(t1)]]) + y = np.array([[np.sin(t2), 0.0, 0.0, np.cos(t2)]]) + projections = np.eye(d - 1) + + res = ot.stereographic_sliced_wasserstein_sphere(x, y, projections=projections) + np.testing.assert_almost_equal(res, np.abs(t1 - t2) / np.sqrt(d - 1)) + + # the poles are mapped without nan: the north pole on the boundary of the + # eps-cap, the south pole on the origin + north = np.zeros((1, d)) + north[0, -1] = 1.0 + south = np.zeros((1, d)) + south[0, -1] = -1.0 + + res = ot.stereographic_sliced_wasserstein_sphere( + north, south, projections=projections + ) + np.testing.assert_almost_equal(res, (np.pi - np.arccos(1 - 1e-6)) / np.sqrt(d - 1)) + + +def test_stereographic_sliced_sphere_rotations(): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + y = rng.randn(n, 3) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + + rotations = ot.sliced.get_random_rotations(3, 5, seed=rng) + R0 = ot.sliced.get_random_rotations(3, 1, seed=rng)[0] + + # rotating both measures amounts to composing the rotations + cost1 = ot.stereographic_sliced_wasserstein_sphere( + x @ R0.T, y @ R0.T, rotations=rotations, seed=42 + ) + cost2 = ot.stereographic_sliced_wasserstein_sphere( + x, y, rotations=rotations @ R0, seed=42 + ) + np.testing.assert_almost_equal(cost1, cost2) + + +def test_stereographic_sliced_sphere_backend_type_devices(nx): + n = 100 + rng = np.random.RandomState(0) + + x = rng.randn(n, 3) + x = x / np.sqrt(np.sum(x**2, -1, keepdims=True)) + + y = rng.randn(2 * n, 3) + y = y / np.sqrt(np.sum(y**2, -1, keepdims=True)) + + sw_np, log = ot.stereographic_sliced_wasserstein_sphere( + x, y, n_rotations=3, log=True + ) + P = log["projections"] + R = log["rotations"] + + for tp in nx.__type_list__: + print(nx.dtype_device(tp)) + + xb, yb = nx.from_numpy(x, y, type_as=tp) + + valb = ot.stereographic_sliced_wasserstein_sphere( + xb, + yb, + projections=nx.from_numpy(P, type_as=tp), + rotations=nx.from_numpy(R, type_as=tp), + ) + + nx.assert_same_dtype_device(xb, valb) + np.testing.assert_almost_equal(sw_np, nx.to_numpy(valb)) + + +def test_stereographic_sliced_sphere_gradient(): + if torch: + import torch.nn.functional as F + + X0 = torch.randn((20, 3)) + X0 = F.normalize(X0, p=2, dim=-1) + X0.requires_grad_(True) + + X1 = torch.randn((20, 3)) + X1 = F.normalize(X1, p=2, dim=-1) + + sw = ot.stereographic_sliced_wasserstein_sphere( + X1, X0, n_projections=100, n_rotations=3, p=2 + ) + grad_x0 = torch.autograd.grad(sw, X0)[0] + + assert not torch.any(torch.isnan(grad_x0))