[MRG] Add QSW sampling for sliced Wasserstein - #838
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #838 +/- ##
==========================================
+ Coverage 96.86% 96.87% +0.01%
==========================================
Files 128 128
Lines 25947 26099 +152
==========================================
+ Hits 25133 25284 +151
- Misses 814 815 +1 🚀 New features to boost your workflow:
|
clbonet
left a comment
clbonet
left a comment
There was a problem hiding this comment.
Thank you @Samuel-Vangu for the great PR! Overall the code is great.
About the name of the method in "sampling_slices", I am wondering whether we should put something more precised than "qsw". The generalized_spiral seems to be the fastest method, but maybe we would like to add other Quasi Monte-Carlo methods in the future such as the minimization of the Coulomb energy or something else. Thus, I think we should put something more precised, e.g. "spiral_qmc"?
Also, for the randomized option. Since it can be applied to any Quasi-Monte Carlo method, maybe it should be "randomized_spiral_qmc", and we can get the boolean with checking "randomized" in sampling_slices?
What do you think @Samuel-Vangu, @rflamary ?
Otherwise, I have few minor comments below.
|
|
||
| By default, the projection directions :math:`\theta` are sampled uniformly | ||
| at random. Setting ``sampling_slices`` to ``"qsw"`` or ``"rqsw"`` instead | ||
| uses Quasi-Monte Carlo point sets on the sphere (generalized spiral |
There was a problem hiding this comment.
Is the generalized spiral the best QMC method? With the current API, we can't really add other QMC methods such as Equal area mappings, Coulomb energy etc...
Would it be worth it to add this in the names?
| ): | ||
| r""" | ||
| Generates n_projections points on the sphere via generalized | ||
| spiral points (Rakhmanov, Saff & Zhou, 1994). |
|
|
||
|
|
||
| def get_projections_spiral( | ||
| d, n_projections, randomized=True, seed=None, backend=None, type_as=None |
There was a problem hiding this comment.
Generalized spirals are only valid for d==3, so maybe we don't need d here
| if d != 3: | ||
| raise ValueError( | ||
| f"get_projections_spiral is only implemented for d=3, got d={d}" | ||
| ) |
| elif sampling_slices == "qsw": | ||
| projections = get_projections_spiral( | ||
| d, n_projections, randomized=False, backend=nx, type_as=X_s | ||
| ) | ||
| elif sampling_slices == "rqsw": | ||
| projections = get_projections_spiral( | ||
| d, n_projections, randomized=True, seed=seed, backend=nx, type_as=X_s | ||
| ) |
There was a problem hiding this comment.
Maybe put "randomized_qsw" in the name, and then use a variable randomized="randomized" in sampling_slices. It would be better if we want to add more QMC sampling methods later
| # sin/cos are not exposed by the backend abstraction (only arccos, atan2 | ||
| # exist), so the deterministic point construction is done in plain NumPy | ||
| # and converted to the target backend at the end. |
There was a problem hiding this comment.
You can add sin and cos to the backend and use nx instead of np
| ## 0.9.8dev | ||
| *August 2026* | ||
|
|
||
| #### New features |
| - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized | ||
| spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`, | ||
| as described in [93] (PR #838) |
| between two point clouds -- known here in closed form, with no | ||
| approximation error left except from the number of projections itself. | ||
|
|
||
| .. [93] Nguyen, K., Bariletto, N., & Ho, N. (2024). Quasi-Monte Carlo for |
3 participants
Types of changes
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_distancesamples 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_spiraland 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.py— 62 tests passed.Added tests covering:
PR checklist