rtavenar · GitHub

@rtavenar

…enard 2025

@rflamary

Pull Request Overview

This PR adds an efficient 1D partial optimal transport solver implemented in Cython with a Python wrapper, updates packaging to include the new extension, and provides basic tests.

  • Introduces partial_wasserstein_1d in ot.partial, backed by a Cython routine.
  • Registers a new Cython extension partial_cython in setup.py.
  • Adds tests for dimensionality checks and correctness of 1D solver outputs.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
test/test_partial.py New tests for error on non-1D input and basic functional tests
setup.py Added partial_cython extension to cythonize call
ot/partial/partial_solvers.py Imported and exposed the Python wrapper partial_wasserstein_1d
ot/partial/partial_cython.pyx Complete Cython implementation of the 1D partial Wasserstein solver
ot/partial/init.py Exported partial_wasserstein_1d in the package API
Comments suppressed due to low confidence (2)

ot/partial/partial_solvers.py:1308

  • The docstring states it "returns the OT matrix", but the function actually returns (indices_x, indices_y, marginal_costs). Please update the brief description to reflect the actual return values.
    r"""Solves the partial Wasserstein distance problem between 1d measures and returns

test/test_partial.py:62

  • [nitpick] The test comments mention non-1D inputs but don’t define xs and xt with invalid shapes. Consider explicitly creating e.g. xs = np.zeros((5,2)), xt = np.zeros((5,2)) before calling to ensure the assertion is triggered.
    with pytest.raises(AssertionError):
Co-authored-by: lchapel <laetitia.chapel@irisa.fr>
…rly)

@rtavenar rtavenar changed the title [WIP] Partial optimal transport 1d solver [MRG] Partial optimal transport 1d solver

Aug 5, 2025

rflamary

name="ot.partial.partial_cython",
sources=["ot/partial/partial_cython.pyx"],
include_dirs=[numpy.get_include(), os.path.join(ROOT, "ot/partial")],
extra_compile_args=["-O3"],

@rflamary

Read the original on github.com ↗