Skip to content

[MRG] Add QSW sampling for sliced Wasserstein - #838

Open
Samuel-Vangu wants to merge 3 commits into
PythonOT:masterfrom
Samuel-Vangu:feature/add-qsw-sampling
Open

[MRG] Add QSW sampling for sliced Wasserstein#838
Samuel-Vangu wants to merge 3 commits into
PythonOT:masterfrom
Samuel-Vangu:feature/add-qsw-sampling

Conversation

@Samuel-Vangu

Copy link
Copy Markdown

Types of changes

  • New feature
  • Documentation update
  • Tests

Motivation and context / Related issue

Closes #835

This PR adds Quasi-Monte Carlo (QMC) sampling of projection directions to the Sliced Wasserstein module.

Currently, sliced_wasserstein_distance samples projection directions uniformly at random, corresponding to standard Monte Carlo sampling. This PR adds two alternatives based on the generalized spiral point construction described in [Nguyen, Bariletto & Ho (2024)](https://arxiv.org/abs/2309.11713):

  • sampling_slices="qsw": deterministic Quasi-Sliced Wasserstein (QSW) projection directions.
  • sampling_slices="rqsw": Randomized QSW (RQSW), obtained by applying a random rotation to the deterministic spiral point set.

The new sampling methods are currently limited to 3D, while the existing "uniform" sampling remains the default.

The implementation also exposes get_projections_spiral and updates the documentation, README references, release notes, and adds a 3D example.

How has this been tested (if it applies)

The changes have been tested with:

  • pre-commit run --all-files — all checks pass.

  • pytest test/sliced/test_sliced_distances.py62 tests passed.

  • Added tests covering:

    • deterministic spiral projections;
    • randomized projections and sphere preservation;
    • seed reproducibility;
    • invalid dimensions and sampling methods;
    • QSW/RQSW with NumPy, JAX, PyTorch and TensorFlow backends;
    • consistency of deterministic QSW across backends;
    • QSW approximation compared with uniform Monte Carlo sampling in 3D.

PR checklist

@Samuel-Vangu

Copy link
Copy Markdown
Author

Hi @rflamary, @clbonet,

This one is ready for your review whenever you have time. Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Quasi-Monte Carlo point sets for the Sliced Wasserstein module

1 participant