[MRG] Stereographic Spherical Sliced Wasserstein

[huytransformer]

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:

• Added the ot.sliced.stereographic_sliced_wasserstein_sphere function to compute the S3W distance, and the RI-S3W distance with the n_rotations argument. The amortized version ARI-S3W can be computed by pregenerating a pool of rotations with ot.sliced.get_random_rotations and passing a random subset with the rotations argument.
• Added the ot.sliced.get_random_rotations function to sample rotations uniformly on SO(d), used for RI-S3W.
• Added tests for these functions in test/sliced/test_spherical_sliced.py.

Motivation and context / Related issue

The S3W distances are fast alternatives to ot.sliced_wasserstein_sphere for 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

• I have read the CONTRIBUTING document.
• The documentation is up-to-date with the changes I made (check build artifacts).
• All tests passed, and additional code has been covered with new tests.
• I have added the PR and Issue fix to the RELEASES.md file.

[clbonet]

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.88%. Comparing base (62706ab) to head (80ea606).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #836      +/-   ##
==========================================
+ Coverage   96.86%   96.88%   +0.02%     
==========================================
  Files         128      128              
  Lines       25709    25909     +200     
==========================================
+ Hits        24902    25102     +200     
  Misses        807      807              

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[clbonet]

Thank you @huytransformer for doing the modifications. The code now seems good to me. Don't forget to add your name in the Contributors file.

Also, it could be nice to add a call to the function in an example (e.g. with a comparison with SSW). There is an example showing the variance with respect to the number of projections. You could add a comparison here, or do another example if you have a better idea. If you don't have time to add this, we can also merge the code directly.

[huytransformer]

Thank you @huytransformer for doing the modifications. The code now seems good to me. Don't forget to add your name in the Contributors file.

Also, it could be nice to add a call to the function in an example (e.g. with a comparison with SSW). There is an example showing the variance with respect to the number of projections. You could add a comparison here, or do another example if you have a better idea. If you don't have time to add this, we can also merge the code directly.

Thank you very much @clbonet! I added my name to CONTRIBUTORS.md and added $S3W_1$ next to $SSW_1$ in the variance example as you suggested. If there's anything else blocking merge please let me know :)

[rflamary]

Hello @huytransformer , Im' a bit curious about teh difference in scale betweet SSW and S3W. is that suppsoed to happen or are we missing a scaling/division in the implementation ?

[huytransformer]

Hello @huytransformer , Im' a bit curious about teh difference in scale betweet SSW and S3W. is that suppsoed to happen or are we missing a scaling/division in the implementation ?

Thank you so much for the good catch, I was indeed missing a factor of $\pi$ (fixed). The remaining gap I think is expected by construction. Also S3W depends on where the data live on the sphere which is why we usually use the rotationally invariant version by setting n_rotations > 1 :)