[MRG] Stereographic Spherical Sliced Wasserstein - #836
Open
huytransformer wants to merge 4 commits into
Open
Conversation
Collaborator
|
Thank you @huytransformer for this PR! This is a very good work, and we were waiting for more sliced distances on the sphere! I did a review. The code seems good to me and I only have minor comments. |
clbonet
reviewed
Aug 25, 2026
| return res | ||
|
|
||
|
|
||
| def _projection_sphere_stereographic(x, eps, nx): |
Collaborator
There was a problem hiding this comment.
Can be moved to _utils.py?
Collaborator
There was a problem hiding this comment.
Please add the description of parameters and returned values
Collaborator
There was a problem hiding this comment.
Since it uses the composition of h with the stereographic projection, the name might be a bit misleading?
Author
There was a problem hiding this comment.
Done, done, and done (I renamed it to projection_sphere_to_ball)
| 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>`, |
| 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 |
|
|
||
| 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\}` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Types of changes
This PR aims to add the Stereographic Spherical Sliced Wasserstein distance and its rotationally invariant extension introduced in Stereographic Spherical Sliced Wasserstein Distances.
Changes:
ot.sliced.stereographic_sliced_wasserstein_spherefunction to compute the S3W distance, and the RI-S3W distance with then_rotationsargument. The amortized version ARI-S3W can be computed by pregenerating a pool of rotations withot.sliced.get_random_rotationsand passing a random subset with therotationsargument.ot.sliced.get_random_rotationsfunction to sample rotations uniformly on SO(d), used for RI-S3W.test/sliced/test_spherical_sliced.py.Motivation and context / Related issue
The S3W distances are fast alternatives to
ot.sliced_wasserstein_spherefor comparing distributions on the sphere: after a stereographic projection composed with a near-isometric map, they only require computing sliced Wasserstein distances in R^{d-1}, and they support all backends including tensorflow.How has this been tested (if it applies)
I added tests of these functions in
test/sliced/test_spherical_sliced.py.PR checklist