[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 99.40476% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.88%. Comparing base (32ed906) to head (55c58e0).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #838      +/-   ##
==========================================
+ Coverage   96.86%   96.88%   +0.01%     
==========================================
  Files         128      128              
  Lines       25972    26137     +165     
==========================================
+ Hits        25158    25322     +164     
- Misses        814      815       +1

🚀 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: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: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

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.

clbonet

Thank you @Samuel-Vangu. A last small comment, if you add sin and cos to the backend, can you also please add the tests.

…backends

…Vangu/POT into feature/add-qsw-sampling

Samuel-Vangu · 2026-09-07T07:59:51Z

Hi @clbonet ,

Added sin/cos coverage to both test_empty_backend and test_func_backends,
as requested. All green (test_backend.py + test_sliced_distances.py, 62
tests).

rflamary

Hello @Samuel-Vangu

Thanks for the contribution. I detected a few remeianing somments thta need sto be adressed but we are close to merging the PR.

rflamary · 2026-09-07T08:41:49Z

+
+    results = {}
+    projections_by_backend = {}
+    for nx in backends:


no loop inside the etsts please pass nx as a parameter of the functin and pytest will loop over available backends

rflamary · 2026-09-07T08:42:00Z

+    X_t = rng.normal(1, 1, (30, d))
+
+    results = {}
+    for nx in backends:


rflamary · 2026-09-07T08:42:33Z

+    mean_uniform_error = np.mean(uniform_errors)
+    mean_rqsw_error = np.mean(rqsw_errors)
+
+    print(f"\n[DEBUG] mean uniform error={mean_uniform_error:.6e}")


remove the debugging print

rflamary · 2026-09-07T08:46:41Z

+    if not randomized:
+        return theta
+
+    if isinstance(seed, np.random.RandomState) and str(nx) == "numpy":


this test should not be there, nx.seed and nx.randn shoudl work with numpy backend too

…n the last review

Samuel-Vangu · 2026-09-07T09:23:39Z

Hi @rflamary, addressed all three test/sliced/test_sliced_distances.py
comments: both loops replaced with nx as a function parameter (pytest
now loops over backends automatically), and the debug print removed --
the thresholds it was showing are documented in the test's docstring
instead.

On the isinstance(seed, np.random.RandomState) check in _utils.py: I
tested this before removing it, and I believe it's actually needed.
RandomState.seed() doesn't accept another RandomState instance as
input:

import numpy as np
rng_self = np.random.RandomState()
rng_passed = np.random.RandomState(42)
rng_self.seed(rng_passed)
# TypeError: Cannot cast scalar from dtype('O') to dtype('int64')
#            according to the rule 'safe'

So the generic nx.seed(seed) path raises when someone passes an
already-instantiated RandomState rather than an int. This is also the
exact same pattern already used in get_random_projections,
get_projections_sphere, and get_random_rotations just above in this
same file, which I didn't touch.

rflamary · 2026-09-07T09:30:29Z

Hello @Samuel-Vangu thanks for the fixeds.

About nx.seed : if it breaks then it means that we need to fix the numpy backend (and do the test in the nx.seed function). We need a backend that can run this without if/then tests. Coiuld you do that please? And thank you for identifying this bug.

Samuel-Vangu · 2026-09-07T10:01:46Z

@rflamary ,
Thanks! Grep found 6 hits: the same check in all 4 functions here, plus
one in _spherical_sliced.py and one in ot/utils.py that look like
different logic, worth checking separately.

I'd rather open a follow-up PR for the actual fix -- ot/backend.py is
shared code deserving its own review, and this PR is close to merging,
so I don't want to reopen a broader cycle now. Leaving the checks
untouched here for that reason. What do you think ?

rflamary · 2026-09-07T11:31:37Z

OK we can merge but please before open an issue with precise description of what needs to be done and where (and hopefully indeed do the PR after ;) ).

The whole point of the backend is to have shared code and this is a bug so it should be listed in Issues.

Samuel-Vangu · 2026-09-07T13:23:48Z

@rflamary ,

I just opened the issue #848

clbonet · 2026-09-08T20:04:04Z

@Samuel-Vangu Don't forget to add your name in the Contributors file.

Samuel-Vangu · 2026-09-09T03:28:55Z

Hi @clbonet ,

Added my name to CONTRIBUTORS.md.

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

Comment thread examples/sliced-wasserstein/plot_qsw_3d.py Outdated

Push some fixes as requested in the review

90f62f3

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

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

5ac4940

clbonet reviewed Sep 7, 2026

View reviewed changes

Comment thread ot/backend.py

Samuel-Vangu added 2 commits September 7, 2026 09:52

test(backend): cover sin and cos in test_empty_backend and test_func_…

99df88a

…backends

Merge branch 'feature/add-qsw-sampling' of https://github.com/Samuel-…

5c108db

…Vangu/POT into feature/add-qsw-sampling

rflamary reviewed Sep 7, 2026

View reviewed changes

Modified the test/sliced/test_sliced_distances.py file as suggested i…

b319e65

…n the last review

Samuel-Vangu mentioned this pull request Sep 7, 2026

NumpyBackend.seed() rejects a np.random.RandomState instance, forcing callers to special-case it #848

Open

CunjieLee mentioned this pull request Sep 7, 2026

[MRG] Support RandomState instances in NumpyBackend.seed #849

Open

9 tasks

Added my name in the contributors file

55c58e0

rflamary merged commit 45b3c21 into PythonOT:master Sep 9, 2026
21 checks passed

Samuel-Vangu mentioned this pull request Sep 9, 2026

Add UnifOrtho sampling of slicing directions for Sliced Wasserstein #850

Open



		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!

Uh oh!

Uh oh!

clbonet Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

clbonet Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

clbonet Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Samuel-Vangu commented Sep 1, 2026

Uh oh!

clbonet left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Samuel-Vangu commented Sep 7, 2026

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

rflamary Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

rflamary Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

rflamary Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

rflamary Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

Samuel-Vangu commented Sep 7, 2026

Uh oh!

rflamary commented Sep 7, 2026

Uh oh!

Samuel-Vangu commented Sep 7, 2026

Uh oh!

rflamary commented Sep 7, 2026

Uh oh!

Samuel-Vangu commented Sep 7, 2026

Uh oh!

clbonet commented Sep 8, 2026

Uh oh!

Samuel-Vangu commented Sep 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Aug 26, 2026 •

edited

Loading