[WIP] Add backend-agnostic semi-discrete OT module and SGD-based solver in ot.semidiscrete by Ferdinand-Genans · Pull Request #812 · PythonOT/POT

Ferdinand-Genans · 2026-05-17T17:07:40Z

Types of changes

New feature (non-breaking change which adds functionality)
Bug fix (non-breaking change which fixes an issue)
Breaking change (fix or feature that would cause existing functionality to change)
Documentation / Examples update

Motivation and context / Related issue

Semi-discrete optimal transport — continuous source, discrete target — is a setting that does not naturally fit POT's discrete-discrete solvers, yet appears in many applications: Monge–Kantorovich statistical depth and quantiles, generative modeling with a continuous prior, Brenier-map estimation, and any setup where the source is given by a sampler rather than a finite empirical distribution.

This PR adds ot.semidiscrete, a single-file backend-agnostic solver for the semi-dual of semi-discrete OT. Every public function takes a reg: float argument, so the user can switch between unregularized OT (reg=0) and the entropic formulation (reg>0) with a single parameter. The underlying algorithm is Averaged SGD on the semi-dual; optionally, the regularization can be decreased throughout the iterations via the DRAG schedule ([83] in README.md; Genans et al., NeurIPS 2025), which empirically improves convergence, especially on large scale problems, when the target regularization is small.

No existing issue tracks this addition.

How has this been tested

test/test_semidiscrete.py — 34 tests, automatically parametrized over every available POT backend (NumPy and PyTorch on the CI runner). Covers:
- convergence to a known closed-form optimum on three toy problems (regular grid, nonuniform target weights, shifted 1D);
- both plain SGD and the DRAG schedule;
- the entropic regime;
- a user-supplied custom cost;
- helper functions (atom_weights row-stochasticity, c_transform reducing to min_j c(x, y_j) at g = 0, ot_map shape and finiteness);
- solver options (warm start, projection bound, log dict, polyak_average=False, init_potential is not mutated).
One doctest on solve_semidiscrete (a 500-iter smoke run).
Full suite + doctest run locally: 35 passed in 5.88 s.
examples/others/plot_semidiscrete.py runs end-to-end in ~5 s on CPU NumPy and produces the Laguerre-cells figure that becomes the gallery page.
Also verified locally a larger scale problem on GPU with PyTorch: discrete measure with 20 000 atoms and 30 000 iterations of minibatch of 64 run in ~17 s on a single CUDA device.

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.

Introduces ot.semidiscrete: Projected Averaged SGD on the semi-dual, with an optional decreasing entropic-regularization schedule (DRAG) from Genans et al. 2025. Works with NumPy, PyTorch, JAX, CuPy and TensorFlow via ot.backend. - ot/semidiscrete.py: solve_semidiscrete, atom_weights, ot_map, c_transform. Closed-form gradient, no autograd graph through the loop; quadratic cost by default with custom-callable override. - ot/__init__.py: register the new submodule. - test/test_semidiscrete.py: convergence on three toy problems with known optimal potentials, helper-function contracts (row- stochasticity of atom_weights, identity for c_transform at g=0, shape and finiteness of ot_map), and solver options (warm-start, projection, log, polyak_average off, entropic regime, custom cost). All tests parametrized over the nx fixture (NumPy + PyTorch). - examples/others/plot_semidiscrete.py: gallery example on a small 2D toy problem with Laguerre cells, empirical cell masses and a Monte Carlo estimate of the semi-dual cost. - RELEASES.md: new-features entry under 0.9.7.dev0.

modified: RELEASES.md modified: examples/others/plot_semidiscrete.py Final small doc modifications.

solve_semidiscrete, and added a more detailed explanation of the effect of this argument on convergence in the example scipt.

"max_cost" explanation.

codecov · 2026-05-22T16:57:28Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.84%. Comparing base (3e10dd6) to head (a5aa34d).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #812      +/-   ##
==========================================
+ Coverage   96.81%   96.84%   +0.02%     
==========================================
  Files         124      126       +2     
  Lines       24276    24502     +226     
==========================================
+ Hits        23503    23729     +226     
  Misses        773      773

🚀 New features to boost your workflow:

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

rflamary · 2026-05-27T08:56:59Z

Hello @Ferdinand-Genans ,

thanks for the PR, we have since merged a paper, could you please resolve conflits and update the reference number ()there is already a 83 now in teh readme file).

I will have a look and do a proper code review ASAP.

- README.md: keep new [83] (Spectral-Grassmann) and [84] (BSP-OT) from upstream; renumber the DRAG citation to [85]. - RELEASES.md: keep both feature bullets (semidiscrete + sgot). - ot/__init__.py: keep both new submodule imports (semidiscrete + sgot). - examples/others/plot_semidiscrete.py: update local refs [83] -> [85] consistently with the README.

Ferdinand-Genans · 2026-05-29T08:13:16Z

Hi @rflamary,
I think it is good now.
Best regards

…rete

rflamary

Hello @Ferdinand-Genans ,

This is a very nice PR, I did a code reveiw and have a few comments mainly on the API to make i more consistent with the other POT functions. Also some suggetsion for the example.

rflamary · 2026-06-01T12:57:22Z

+    return nx.where(mask, one, zero)
+
+
+def atom_weights(


Suggested change

def atom_weights(

def semidiscrete_atom_weights(

rflamary · 2026-06-01T12:57:37Z

+    return _atom_weights(score, reg, log_b, nx)
+
+
+def ot_map(


Suggested change

def ot_map(

def semidiscrete_ot_map(

rflamary · 2026-06-15T12:33:06Z

+    cost=None,
+    reg=0.0,
+):
+    r"""Row-stochastic atom-assignment weights induced by ``semi_dual_potential``.


please do a proper documentation with detailed input and output description. define what is computed with math equations and refer to eq. in papers/books.

rflamary · 2026-06-15T12:33:24Z

+    cost=None,
+    reg=0.0,
+):
+    r"""Transport map :math:`T(x) = \sum_j w_j(x)\, y_j` induced by the potential."""


Same here alos add detailed documentation

rflamary · 2026-06-16T06:23:55Z

+# A single call to :func:`solve_semidiscrete` runs DRAG with the default
+# arguments (``decreasing_reg=True``). We show the initial Voronoi cells
+# (:math:`g = 0`) next to the Laguerre cells at the optimum.
+# In this problem, the maximum cost between samples is 1.0, so we pass it as


is it 1? I would say since it is the sqare it would be \sqrt(2) no? what max cost is it between the continuous dist and the discrete one?

rflamary · 2026-06-16T06:27:16Z

+plot_laguerre_cells(target_positions, g_drag, axes[1], "DRAG")
+plt.tight_layout()
+plt.show()
+


maybe do a visualization of the map from a grid on source with arrows? it would e a nice complementary visualization than the cells and would show an example of the use of the mapping function

rflamary · 2026-06-16T06:28:26Z

+# .. math::
+#     \mathcal{S}(g) = \langle g, b\rangle + \mathbb{E}_X[\varphi_g(X)]
+#
+# estimated by Monte Carlo. The solver maximises :math:`\mathcal{S}`.


visualize the mass with a bar plot to see the small discrepancy (with ylim around 1/15 and betwen 0 and 1/15 to show clearly)?

rflamary · 2026-06-16T06:42:19Z

+def solve_semidiscrete(
+    target_positions,
+    source_sampler,
+    target_weights=None,
+    cost=None,
+    reg=0.0,
+    n_iter=10_000,
+    batch_size=32,
+    lr0=None,
+    lr_exponent=2.0 / 3.0,
+    init_potential=None,
+    decreasing_reg=True,
+    decreasing_reg_initial_eps=0.1,
+    decreasing_reg_exponent=0.5,
+    max_cost=None,
+    polyak_average=True,
+    log=False,
+):


Suggested change

def solve_semidiscrete(

target_positions,

source_sampler,

target_weights=None,

cost=None,

reg=0.0,

n_iter=10_000,

batch_size=32,

lr0=None,

lr_exponent=2.0 / 3.0,

init_potential=None,

decreasing_reg=True,

decreasing_reg_initial_eps=0.1,

decreasing_reg_exponent=0.5,

max_cost=None,

polyak_average=True,

log=False,

):

def solve_semidiscrete(

X_target,

sampler_source,

a_target=None,

metric=None,

reg=0.0,

max_iter=10_000,

batch_size=32,

lr0=None,

lr_exponent=2.0 / 3.0,

init_potential=None,

decreasing_reg=True,

decreasing_reg_initial_eps=0.1,

decreasing_reg_exponent=0.5,

max_metric=None,

polyak_average=True,

log=False,

):

Renamed parameter to have a closer API to ot.solve_sample. all other public functions shoudl also have renamed params (map, weight, c_transform)

cost-> metric should accept 'sqeuclidean' by default but also euclidean and other stuff from ot.dist (maybe should call ot.dist if not a function?)

sampler_source should accept 'unif', 'unit_ball' et 'normal' in addition to a function that takes the size of eh batch in parameter

rflamary · 2026-06-16T06:47:19Z

+    return float(np.linalg.norm(estimated - reference))
+
+
+def lift(nx, target_np, weights_np):


you can replace this function by
target, weights = nx.from_numpy(target_np,weights)

no need for a new function there

rflamary · 2026-06-16T06:55:18Z

+def solve_semidiscrete(
+    target_positions,
+    source_sampler,
+    target_weights=None,
+    cost=None,
+    reg=0.0,
+    n_iter=10_000,
+    batch_size=32,
+    lr0=None,
+    lr_exponent=2.0 / 3.0,
+    init_potential=None,
+    decreasing_reg=True,
+    decreasing_reg_initial_eps=0.1,
+    decreasing_reg_exponent=0.5,
+    max_cost=None,
+    polyak_average=True,
+    log=False,
+):


Suggested change

def solve_semidiscrete(

target_positions,

source_sampler,

target_weights=None,

cost=None,

reg=0.0,

n_iter=10_000,

batch_size=32,

lr0=None,

lr_exponent=2.0 / 3.0,

init_potential=None,

decreasing_reg=True,

decreasing_reg_initial_eps=0.1,

decreasing_reg_exponent=0.5,

max_cost=None,

polyak_average=True,

log=False,

):

def solve_semidiscrete(

X_target,

sampler_source,

a_target=None,

metric=None,

reg=0.0,

max_iter=10_000,

batch_size=32,

lr0=None,

lr_exponent=2.0 / 3.0,

init_potential=None,

decreasing_reg=True,

decreasing_reg_initial_eps=0.1,

decreasing_reg_exponent=0.5,

max_cost=None,

polyak_average=True,

log=False,

):

few change in name of parameters to be more consistent with ot.solve_sample. please change them in all public facing function (atom_weight, map, c_transform)

sampler_source should accept also strings for 'unif' (or 'unif_cube'), 'ball' or 'unif_ball' and 'normal' if not callable , should be unif by default.

cost -> metric should also accept any existing string from ot.dist (it should use ot.dist in the function) if not callable. shoudl be 'sqeuclidean' by default.

Ferdinand-Genans added 4 commits May 17, 2026 17:06

modified: README.md

835c959

modified: RELEASES.md modified: examples/others/plot_semidiscrete.py Final small doc modifications.

Modified "proj_bound" to "max_cost" in the

f8549f0

solve_semidiscrete, and added a more detailed explanation of the effect of this argument on convergence in the example scipt.

Edited the semi-discrete example script

b5383f5

"max_cost" explanation.

github-actions Bot added Tests Examples labels May 17, 2026

Ferdinand-Genans added 2 commits May 17, 2026 19:52

Changed "proj_bound" to "max_cost" in test_semidiscrete.py

3b8b330

Add PR's number .

bb15a06

rflamary changed the title ~~[MRG] Add backend-agnostic semi-discrete OT module and SGD-based solver in ot.semidiscrete~~ [WIP] Add backend-agnostic semi-discrete OT module and SGD-based solver in ot.semidiscrete May 27, 2026

Ferdinand-Genans force-pushed the feature/semi-discrete branch from bc2b63c to 50f4094 Compare May 29, 2026 08:04

rflamary and others added 3 commits June 1, 2026 14:28

Merge branch 'master' into feature/semi-discrete

8cf415b

Merge branch 'master' into feature/semi-discrete

53563cc

Merge remote-tracking branch 'upstream/master' into feature/semi-disc…

ad4277b

…rete

rflamary reviewed Jun 16, 2026

View reviewed changes

Merge branch 'master' into feature/semi-discrete

a5aa34d

		return float(np.linalg.norm(estimated - reference))


		def lift(nx, target_np, weights_np):

Conversation

Ferdinand-Genans commented May 17, 2026

Types of changes

Motivation and context / Related issue

How has this been tested

PR checklist

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

rflamary commented May 27, 2026

Uh oh!

Ferdinand-Genans commented May 29, 2026

Uh oh!

rflamary 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!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented May 22, 2026 •

edited

Loading