[MRG] Partial optimal transport 1d solver

.github/labeler.yml

@@ -74,7 +74,7 @@ ot.solvers:
 
 ot.partial:
 - changed-files:
-  - any-glob-to-any-file: ot/partial.py
+  - any-glob-to-any-file: ot/partial/**
 
 ot.sliced:
 - changed-files:
@@ -94,4 +94,4 @@ ot.dr:
 
 ot.gnn:
 - changed-files:
-  - any-glob-to-any-file: ot/gnn/**
+  - any-glob-to-any-file: ot/gnn/**

.github/workflows/build_doc.yml

@@ -26,7 +26,7 @@ jobs:
         python -m pip install --user --upgrade --progress-bar off pip
         python -m pip install --user --upgrade --progress-bar off -r requirements_all.txt
         python -m pip install --user --upgrade --progress-bar off -r docs/requirements.txt
-        python -m pip install --user --upgrade --progress-bar off ipython "https://api.github.com/repos/sphinx-gallery/sphinx-gallery/zipball/master" memory_profiler
+        python -m pip install --user --upgrade --progress-bar off ipython sphinx-gallery memory_profiler
         python -m pip install --user -e .
     # Look at what we have and fail early if there is some library conflict
     - name: Check installation

.gitignore

@@ -15,6 +15,7 @@ docs/modules/
 
 # Cython output
 ot/lp/emd_wrap.cpp
+ot/partial/partial_cython.cpp
 
 # Distribution / packaging
 .Python

README.md

@@ -389,3 +389,5 @@ Artificial Intelligence.
 [74] Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). [Gradient descent algorithms for Bures-Wasserstein barycenters](https://proceedings.mlr.press/v125/chewi20a.html). In Conference on Learning Theory (pp. 1276-1304). PMLR.
 
 [75] Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://papers.neurips.cc/paper_files/paper/2021/hash/b9acb4ae6121c941324b2b1d3fac5c30-Abstract.html). Advances in Neural Information Processing Systems, 34, 22132-22145.
+
+[76] Chapel, L., Tavenard, R. (2025). [One for all and all for one: Efficient computation of partial Wasserstein distances on the line](https://iclr.cc/virtual/2025/poster/28547). In International Conference on Learning Representations.

RELEASES.md

@@ -17,6 +17,7 @@
 - Backend implementation of `ot.dist` for (PR #701)
 - Updated documentation Quickstart guide and User guide with new API (PR #726)
 - Fix jax version for auto-grad (PR #732)
+- Implement 1d solver for partial optimal transport (PR #741)
 - Fix reg_div function compatibility with numpy in `ot.unbalanced.lbfgsb_unbalanced` via new function `ot.utils.fun_to_numpy` (PR #731)
 - Added to each example in the examples gallery the information about the release version in which it was introduced (PR #743)
 - Removed release information from quickstart guide (PR #744)

examples/unbalanced-partial/plot_partial_1d.py

@@ -0,0 +1,85 @@
+"""
+=========================
+Partial Wasserstein in 1D
+=========================
+
+This script demonstrates how to compute and visualize the Partial Wasserstein distance between two 1D discrete distributions using `ot.partial.partial_wasserstein_1d`.
+
+We illustrate the intermediate transport plans for all `k = 1...n`, where `n = min(len(x_a), len(x_b))`.
+"""
+
+# sphinx_gallery_thumbnail_number = 5
+
+import numpy as np
+import matplotlib.pyplot as plt
+from ot.partial import partial_wasserstein_1d
+
+
+def plot_partial_transport(
+    ax, x_a, x_b, indices_a=None, indices_b=None, marginal_costs=None
+):
+    y_a = np.ones_like(x_a)
+    y_b = -np.ones_like(x_b)
+    min_min = min(x_a.min(), x_b.min())
+    max_max = max(x_a.max(), x_b.max())
+
+    ax.plot([min_min - 1, max_max + 1], [1, 1], "k-", lw=0.5, alpha=0.5)
+    ax.plot([min_min - 1, max_max + 1], [-1, -1], "k-", lw=0.5, alpha=0.5)
+
+    # Plot transport lines
+    if indices_a is not None and indices_b is not None:
+        subset_a = np.sort(x_a[indices_a])
+        subset_b = np.sort(x_b[indices_b])
+
+        for x_a_i, x_b_j in zip(subset_a, subset_b):
+            ax.plot([x_a_i, x_b_j], [1, -1], "k--", alpha=0.7)
+
+    # Plot all points
+    ax.plot(x_a, y_a, "o", color="C0", label="x_a", markersize=8)
+    ax.plot(x_b, y_b, "o", color="C1", label="x_b", markersize=8)
+
+    if marginal_costs is not None:
+        k = len(marginal_costs)
+        ax.set_title(
+            f"Partial Transport - k = {k}, Cumulative Cost = {sum(marginal_costs):.2f}",
+            fontsize=16,
+        )
+    else:
+        ax.set_title("Original 1D Discrete Distributions", fontsize=16)
+    ax.legend(loc="upper right", fontsize=14)
+    ax.set_yticks([])
+    ax.set_xticks([])
+    ax.set_ylim(-2, 2)
+    ax.set_xlim(min(x_a.min(), x_b.min()) - 1, max(x_a.max(), x_b.max()) + 1)
+    ax.axis("off")
+
+
+# Simulate two 1D discrete distributions
+np.random.seed(0)
+n = 6
+x_a = np.sort(np.random.uniform(0, 10, size=n))
+x_b = np.sort(np.random.uniform(0, 10, size=n))
+
+# Plot original distributions
+plt.figure(figsize=(6, 2))
+plot_partial_transport(plt.gca(), x_a, x_b)
+plt.show()
+
+# %%
+indices_a, indices_b, marginal_costs = partial_wasserstein_1d(x_a, x_b)
+
+# Compute cumulative cost
+cumulative_costs = np.cumsum(marginal_costs)
+
+# Visualize all partial transport plans
+for k in range(n):
+    plt.figure(figsize=(6, 2))
+    plot_partial_transport(
+        plt.gca(),
+        x_a,
+        x_b,
+        indices_a[: k + 1],
+        indices_b[: k + 1],
+        marginal_costs[: k + 1],
+    )
+    plt.show()

ot/partial/__init__.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+"""
+Efficient 1D solver for the partial optimal transport problem.
+"""
+
+# Author: Romain Tavenard <[email protected]>
+#
+# License: MIT License
+
+# import compiled emd
+from .partial_solvers import (
+    partial_wasserstein_lagrange,
+    partial_wasserstein,
+    partial_wasserstein2,
+    entropic_partial_wasserstein,
+    gwgrad_partial,
+    gwloss_partial,
+    partial_gromov_wasserstein,
+    partial_gromov_wasserstein2,
+    entropic_partial_gromov_wasserstein,
+    entropic_partial_gromov_wasserstein2,
+    partial_wasserstein_1d,
+)
+
+__all__ = [
+    "partial_wasserstein_1d",
+    "partial_wasserstein_lagrange",
+    "partial_wasserstein",
+    "partial_wasserstein2",
+    "entropic_partial_wasserstein",
+    "gwgrad_partial",
+    "gwloss_partial",
+    "partial_gromov_wasserstein",
+    "partial_gromov_wasserstein2",
+    "entropic_partial_gromov_wasserstein",
+    "entropic_partial_gromov_wasserstein2",
+]