[MRG] Add QSW sampling for sliced Wasserstein by Samuel-Vangu · Pull Request #838 · PythonOT/POT

Samuel-Vangu · 2026-08-25T17:34:49Z

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.py — 62 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

I have read the [CONTRIBUTING](https://pythonot.github.io/contributing.html) 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.

Samuel-Vangu · 2026-08-25T17:44:09Z

Hi @rflamary, @clbonet,

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

codecov · 2026-08-26T09:48:28Z

Codecov Report

❌ Patch coverage is 93.71429% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.84%. Comparing base (9439b8a) to head (90f62f3).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #838      +/-   ##
==========================================
- Coverage   96.86%   96.84%   -0.03%     
==========================================
  Files         128      128              
  Lines       25947    26119     +172     
==========================================
+ Hits        25133    25294     +161     
- Misses        814      825      +11

🚀 New features to boost your workflow:

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

clbonet

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.

clbonet · 2026-08-31T15:27:04Z


+    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


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?

clbonet · 2026-08-31T15:31:01Z

+):
+    r"""
+    Generates n_projections points on the sphere via generalized
+    spiral points (Rakhmanov, Saff & Zhou, 1994).


Please add the ref

clbonet · 2026-08-31T15:34:13Z



+def get_projections_spiral(
+    d, n_projections, randomized=True, seed=None, backend=None, type_as=None


Generalized spirals are only valid for d==3, so maybe we don't need d here

clbonet · 2026-08-31T15:35:35Z

+    if d != 3:
+        raise ValueError(
+            f"get_projections_spiral is only implemented for d=3, got d={d}"
+        )


Not needed?

clbonet · 2026-08-31T15:38:51Z

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


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

clbonet · 2026-08-31T15:40:19Z

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


You can add sin and cos to the backend and use nx instead of np

clbonet · 2026-08-31T15:46:38Z

+## 0.9.8dev
+*August 2026*
+
+#### New features


clbonet · 2026-08-31T15:46:47Z

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


To move below

clbonet · 2026-08-31T16:08:07Z

+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


Update [93] to [95]

Samuel-Vangu · 2026-09-01T07:48:16Z

Hi @clbonet,

Thanks for the detailed review! Here's a summary of what I changed in response to each point:

Naming (sampling_slices values)
Agreed — "qsw"/"rqsw" were too generic given that spiral is just one of several possible QMC constructions (equal-area, Coulomb energy, etc. could come later). Renamed to "spiral_qmc" and "randomized_spiral_qmc", and the dispatcher now derives the boolean via sampling_slices.startswith("randomized_") plus a method = sampling_slices.removeprefix("randomized_"), so adding a new QMC construction later only needs one new elif, not two.

Missing reference
Added Rakhmanov, Saff & Zhou (1994) as [96] in the README, and cited it in get_projections_spiral's docstring alongside [95] (Nguyen et al.) which is now correctly numbered.

d parameter
Kept it in the signature — for consistency with get_projections_sphere and so future d-generic constructions (Coulomb, equal-area) fit the same dispatch pattern — but added if d != 3: raise ValueError(...) inside get_projections_spiral itself, not just in the dispatcher. That way a user calling the function directly (not through sliced_wasserstein_distance) is still protected, not only callers going through sampling_slices.

sin/cos on the backend
Added sin/cos to Backend and all five implementations. I kept get_projections_spiral's deterministic point construction in plain NumPy though, rather than switching to nx: most backends' arange() doesn't propagate type_as's dtype (e.g. JAX defaults to float32 without x64 enabled), so building natively per-backend would likely break test_qsw_matches_across_backends's exact cross-backend equality check.

RELEASES.md
Removed the duplicate ## 0.9.8dev header and merged the entry under the existing section.

[93] → [95]
Fixed in plot_qsw_3d.py and everywhere else it appeared.

Let me know if any of these choices don't sit right.

Add QSW sampling for sliced Wasserstein

dcfe739

github-actions Bot added Tests Examples ot.sliced labels Aug 25, 2026

Add the PR number in the RELEASES.mf file

57a5a19

Samuel-Vangu and others added 2 commits August 25, 2026 22:46

modifying the structure of the example file

a055d12

Merge branch 'master' into feature/add-qsw-sampling

6a834f4

Merge branch 'master' into feature/add-qsw-sampling

8dd4c6b

clbonet self-assigned this Aug 28, 2026

Resolve merge conflicts with recently merged changes

4eb176c

clbonet reviewed Aug 31, 2026

View reviewed changes

Push some fixes as requested in the review

90f62f3

github-actions Bot added the ot.backend label Sep 1, 2026

Navigation Menu

Uh oh!

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

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

Samuel-Vangu commented Aug 25, 2026

Uh oh!

Samuel-Vangu commented Aug 25, 2026

Uh oh!

codecov Bot commented Aug 26, 2026 •

edited

Loading

Uh oh!

clbonet left a comment

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

clbonet Aug 31, 2026

Uh oh!

Samuel-Vangu commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants



		def get_projections_spiral(
		d, n_projections, randomized=True, seed=None, backend=None, type_as=None

Uh oh!

Conversation

Samuel-Vangu commented Aug 25, 2026

Types of changes

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

Uh oh!

Samuel-Vangu commented Aug 25, 2026

Uh oh!

codecov Bot commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

clbonet left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Samuel-Vangu commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Aug 26, 2026 •

edited

Loading