[WIP] Partial optimal transport 1d solver

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)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)

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",
+]