[MRG] Stereographic Spherical Sliced Wasserstein

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.
+

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).

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",

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",
 ]

ot/sliced/_spherical_sliced.py

@@ -5,15 +5,24 @@
 
 # Author: Nicolas Courty <ncourty@irisa.fr>
 # Author: Clément Bonet <clement.bonet.mapp@polytechnique.edu>
+# Author: continuousml <continuousml@gmail.com>
 #
 # 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] <references-s3w>`,
+    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] <references-s3w>`.
+
+    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] <references-s3w>`, 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] <references-s3w>`).
+
+    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

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] <references-get-random-rotations>`
+    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
 ):