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_1dinot.partial, backed by a Cython routine. - Registers a new Cython extension
partial_cythoninsetup.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 |
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
xsandxtwith 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):
rtavenar
changed the title
[WIP] Partial optimal transport 1d solver
[MRG] Partial optimal transport 1d solver
| 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"], |