feat: Add backend-agnostic sparse EMD support

Refactor sparse optimal transport implementation to work across
different backends (NumPy/scipy.sparse, PyTorch/torch.sparse).

Key changes:
- Add `sparse_coo_data()` method to backend layer for uniform sparse
  matrix handling across NumPy, PyTorch, JAX, and TensorFlow backends
- Update `emd()` and `emd2()` to return transport plans in backend-native
  sparse format (scipy.sparse for NumPy, torch.sparse for PyTorch)
- Refactor `plot2D_samples_mat()` to efficiently visualize both dense
  and sparse transport plans by detecting format and iterating only
  over non-zero entries for sparse matrices
- Update `plot_sparse_emd.py` example to use new plotting function
- Add comprehensive tests for sparse EMD across backends
- Update documentation to reflect backend-agnostic sparse support

Author: nathanneike

examples/plot_sparse_emd.py

@@ -4,17 +4,16 @@
 Sparse Optimal Transport
 ============================================
 
-In many real-world optimal transport (OT) problems, the transport plan is naturally sparse: only a small fraction of all possible source-target pairs actually exchange mass. In such cases, using sparse OT solvers can provide significant computational speedups and memory savings compared to dense solvers, which compute and store the full transport matrix.
-
-The figure below illustrates the advantages of sparse OT solvers over dense ones in terms of speed and memory usage for different sparsity levels of the transport plan.
-
-.. image:: /_static/images/comparison.png
-    :align: center
-    :width: 80%
-    :alt: Dense vs Sparse OT: Speed and Memory Advantages
+In many real-world optimal transport (OT) problems, the transport plan is
+naturally sparse: only a small fraction of all possible source-target pairs
+actually exchange mass. Using sparse OT solvers can provide significant
+computational speedups and memory savings compared to dense solvers.
+
+This example demonstrates how to use sparse cost matrices with POT's EMD solver,
+comparing sparse and dense formulations on both a minimal example and a larger
+concentric circles dataset.
 """
 
-
 # Author: Nathan Neike <nathan.neike@example.com>
 # License: MIT License
 # sphinx_gallery_thumbnail_number = 2
@@ -26,279 +25,199 @@
 
 
 ##############################################################################
-# Generate minimal example data
+# Minimal example with 4 points
 # ------------------------------
-#
-# We create a simple example with 2 source points and 2 target points to
-# illustrate the concept of sparse optimal transport.
 
 # %%
 
-X = np.array([[0, 0], [1, 0]])
-Y = np.array([[0, 1], [1, 1]])
-a = np.array([0.5, 0.5])
-b = np.array([0.5, 0.5])
-
-
-##############################################################################
-# Build sparse cost matrix
-# -------------------------
-#
-# Instead of allowing all possible edges (dense OT), we only allow two edges:
-# source 0 -> target 0 and source 1 -> target 1. This is specified using a
-# sparse matrix format (COO).
-
-# %%
+X = np.array([[0, 0], [1, 0], [0.5, 0], [1.5, 0]])
+Y = np.array([[0, 1], [1, 1], [0.5, 1], [1.5, 1]])
+a = np.array([0.25, 0.25, 0.25, 0.25])
+b = np.array([0.25, 0.25, 0.25, 0.25])
 
-# Only allow two edges: source 0 -> target 0, source 1 -> target 1
-rows = [0, 1]
-cols = [0, 1]
-vals = [np.linalg.norm(X[0] - Y[0]), np.linalg.norm(X[1] - Y[1])]
-M_sparse = coo_matrix((vals, (rows, cols)), shape=(2, 2))
+# Build sparse cost matrix allowing only selected edges
+rows = [0, 1, 2, 3]
+cols = [0, 1, 2, 3]
+vals = [np.linalg.norm(X[i] - Y[j]) for i, j in zip(rows, cols)]
+M_sparse = coo_matrix((vals, (rows, cols)), shape=(4, 4))
 
 
 ##############################################################################
-# Solve sparse OT problem
-# ------------------------
-#
-# When passing a sparse cost matrix to ot.emd with log=True, the solution
-# is returned in the log dictionary with fields 'flow_sources', 'flow_targets',
-# and 'flow_values' containing the edge information.
+# Solve and display sparse OT solution
+# -------------------------------------
 
 # %%
 
 G, log = ot.emd(a, b, M_sparse, log=True)
 
 print("Sparse OT cost:", log["cost"])
-print("Edges:")
-for i, j, v in zip(log["flow_sources"], log["flow_targets"], log["flow_values"]):
-    print(f"  source {i} -> target {j}, flow={v:.3f}")
+print("Solution format:", type(G))
+print("Non-zero edges:", G.nnz)
+print("\nEdges:")
+G_coo = G if isinstance(G, coo_matrix) else G.tocoo()
+for i, j, v in zip(G_coo.row, G_coo.col, G_coo.data):
+    if v > 1e-10:
+        print(f"  source {i} -> target {j}, flow={v:.3f}")
 
 
 ##############################################################################
-# Visualize allowed edges
-# ---------------------------------
-#
-# The sparse cost matrix only allows transport along specific edges.
+# Visualize sparse vs dense edge structure
+# -----------------------------------------
 
 # %%
 
-
 plt.figure(figsize=(8, 4))
 
-# Sparse OT: allowed edges only
 plt.subplot(1, 2, 1)
 plt.scatter(X[:, 0], X[:, 1], c="r", marker="o", s=100, zorder=3)
 plt.scatter(Y[:, 0], Y[:, 1], c="b", marker="x", s=100, zorder=3)
 for i, j in zip(rows, cols):
-    plt.plot([X[i, 0], Y[j, 0]], [X[i, 1], Y[j, 1]], "b-", linewidth=2, alpha=0.6)
+    plt.plot([X[i, 0], Y[j, 0]], [X[i, 1], Y[j, 1]], "b-", linewidth=1, alpha=0.6)
 plt.title("Sparse OT: Allowed Edges Only")
-
-plt.xlim(-0.5, 1.5)
+plt.xlim(-0.5, 2.0)
 plt.ylim(-0.5, 1.5)
-plt.xticks([0, 1])
-plt.yticks([0, 1])
 
-# Dense OT: all possible edges
 plt.subplot(1, 2, 2)
 plt.scatter(X[:, 0], X[:, 1], c="r", marker="o", s=100, zorder=3)
 plt.scatter(Y[:, 0], Y[:, 1], c="b", marker="x", s=100, zorder=3)
-for i in range(2):
-    for j in range(2):
-        plt.plot([X[i, 0], Y[j, 0]], [X[i, 1], Y[j, 1]], "b-", linewidth=2, alpha=0.3)
+for i in range(len(X)):
+    for j in range(len(Y)):
+        plt.plot([X[i, 0], Y[j, 0]], [X[i, 1], Y[j, 1]], "b-", linewidth=1, alpha=0.3)
 plt.title("Dense OT: All Possible Edges")
-plt.xlim(-0.5, 1.5)
+plt.xlim(-0.5, 2.0)
 plt.ylim(-0.5, 1.5)
-plt.xticks([0, 1])
-plt.yticks([0, 1])
 
 plt.tight_layout()
 
 
 ##############################################################################
-# Larger example with clusters
-# --------------------------------------
-#
-# Now we create a more realistic example with multiple clusters of sources
-# and targets, where transport is only allowed within each cluster.
+# Larger example: concentric circles
+# -----------------------------------
 
 # %%
 
-grid_size = 4
-n_clusters = grid_size * grid_size
-points_per_cluster = 2
-cluster_spacing = 15.0
-intra_cluster_spacing = 1.5
-cluster_centers = (
-    np.array([[i, j] for i in range(grid_size) for j in range(grid_size)])
-    * cluster_spacing
-)
+n_clusters = 8
+points_per_cluster = 25
+n = n_clusters * points_per_cluster
+k_neighbors = 8
+rng = np.random.default_rng(0)
 
-X_large = []
-Y_large = []
-a_large = []
-b_large = []
+r_source = 1.0
+r_target = 2.0
+noise_scale = 0.06
 
-for idx, (cx, cy) in enumerate(cluster_centers):
-    for i in range(points_per_cluster):
-        X_large.append(
-            [cx + intra_cluster_spacing * (i - 1), cy - intra_cluster_spacing]
-        )
-        a_large.append(1.0 / (n_clusters * points_per_cluster))
-
-    for i in range(points_per_cluster):
-        Y_large.append(
-            [cx + intra_cluster_spacing * (i - 1), cy + intra_cluster_spacing]
-        )
-        b_large.append(1.0 / (n_clusters * points_per_cluster))
+theta = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False)
+cluster_labels = np.repeat(np.arange(n_clusters), points_per_cluster)
 
-X_large = np.array(X_large)
-Y_large = np.array(Y_large)
-a_large = np.array(a_large)
-b_large = np.array(b_large)
-
-nA = nB = n_clusters * points_per_cluster
-source_labels = np.repeat(np.arange(n_clusters), points_per_cluster)
-sink_labels = np.repeat(np.arange(n_clusters), points_per_cluster)
-
-
-##############################################################################
-# Build sparse cost matrix (intra-cluster only)
-# ----------------------------------------------
-#
-# We construct a sparse cost matrix that only includes edges within each cluster.
+X_large = np.column_stack(
+    [r_source * np.cos(theta), r_source * np.sin(theta)]
+) + rng.normal(scale=noise_scale, size=(n, 2))
+Y_large = np.column_stack(
+    [r_target * np.cos(theta), r_target * np.sin(theta)]
+) + rng.normal(scale=noise_scale, size=(n, 2))
 
-# %%
+a_large = np.zeros(n)
+b_large = np.zeros(n)
+for k in range(n_clusters):
+    idx = np.where(cluster_labels == k)[0]
+    a_large[idx] = 1.0 / n_clusters / points_per_cluster
+    b_large[idx] = 1.0 / n_clusters / points_per_cluster
 
 M_full = ot.dist(X_large, Y_large, metric="euclidean")
 
+# Build sparse cost matrix: intra-cluster k-nearest neighbors
+angles_X = np.arctan2(X_large[:, 1], X_large[:, 0])
+angles_Y = np.arctan2(Y_large[:, 1], Y_large[:, 0])
+
 rows = []
 cols = []
 vals = []
 for k in range(n_clusters):
-    src_idx = np.where(source_labels == k)[0]
-    sink_idx = np.where(sink_labels == k)[0]
+    src_idx = np.where(cluster_labels == k)[0]
+    tgt_idx = np.where(cluster_labels == k)[0]
     for i in src_idx:
-        for j in sink_idx:
+        diff = np.angle(np.exp(1j * (angles_Y[tgt_idx] - angles_X[i])))
+        idx = np.argsort(np.abs(diff))[:k_neighbors]
+        for j_local in idx:
+            j = tgt_idx[j_local]
             rows.append(i)
             cols.append(j)
             vals.append(M_full[i, j])
-M_sparse_large = coo_matrix((vals, (rows, cols)), shape=(nA, nB))
 
+M_sparse_large = coo_matrix((vals, (rows, cols)), shape=(n, n))
+allowed_sparse = set(zip(rows, cols))
 
 ##############################################################################
-# Visualize allowed edges structure
-# ----------------------------------
-#
-# Dense OT allows all connections, while sparse OT restricts to intra-cluster edges.
+# Visualize edge structures
+# --------------------------
 
 # %%
 
 plt.figure(figsize=(16, 6))
 
-# Dense OT: all possible edges
 plt.subplot(1, 2, 1)
-for i in range(nA):
-    for j in range(nB):
+for i in range(n):
+    for j in range(n):
         plt.plot(
             [X_large[i, 0], Y_large[j, 0]],
             [X_large[i, 1], Y_large[j, 1]],
             color="blue",
-            alpha=0.1,
-            linewidth=0.7,
+            alpha=0.2,
+            linewidth=0.05,
         )
 plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20)
 plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20)
 plt.axis("equal")
 plt.title("Dense OT: All Possible Edges")
 
-# Sparse OT: only intra-cluster edges
 plt.subplot(1, 2, 2)
-for k in range(n_clusters):
-    src_idx = np.where(source_labels == k)[0]
-    sink_idx = np.where(sink_labels == k)[0]
-    for i in src_idx:
-        for j in sink_idx:
-            plt.plot(
-                [X_large[i, 0], Y_large[j, 0]],
-                [X_large[i, 1], Y_large[j, 1]],
-                color="blue",
-                alpha=0.7,
-                linewidth=1.5,
-            )
+for i, j in allowed_sparse:
+    plt.plot(
+        [X_large[i, 0], Y_large[j, 0]],
+        [X_large[i, 1], Y_large[j, 1]],
+        color="blue",
+        alpha=1,
+        linewidth=0.05,
+    )
 plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20)
 plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20)
 plt.axis("equal")
-plt.title("Sparse OT: Only Intra-Cluster Edges")
+plt.title("Sparse OT: Intra-Cluster k-NN Edges")
 
 plt.tight_layout()
-
+plt.show()
 
 ##############################################################################
-# Solve and compare sparse vs dense OT
-# -------------------------------------
-#
-# We solve both dense and sparse OT problems and verify that they produce
-# the same optimal solution when the sparse edges include the optimal paths.
+# Solve and visualize transport plans
+# ------------------------------------
 
 # %%
 
-# Solve dense OT (full cost matrix)
 G_dense = ot.emd(a_large, b_large, M_full)
 cost_dense = np.sum(G_dense * M_full)
 print(f"Dense OT cost: {cost_dense:.6f}")
 
-# Solve sparse OT (intra-cluster only)
 G_sparse, log_sparse = ot.emd(a_large, b_large, M_sparse_large, log=True)
 cost_sparse = log_sparse["cost"]
 print(f"Sparse OT cost: {cost_sparse:.6f}")
 
-
-##############################################################################
-# Visualize optimal transport plans
-# ----------------------------------
-#
-# Plot the edges that carry flow in the optimal solutions.
-
-# %%
-
 plt.figure(figsize=(16, 6))
 
-# Dense OT
 plt.subplot(1, 2, 1)
-for i in range(nA):
-    for j in range(nB):
-        if G_dense[i, j] > 1e-10:
-            plt.plot(
-                [X_large[i, 0], Y_large[j, 0]],
-                [X_large[i, 1], Y_large[j, 1]],
-                color="blue",
-                alpha=0.7,
-                linewidth=1.5,
-            )
-plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20)
-plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20)
+ot.plot.plot2D_samples_mat(
+    X_large, Y_large, G_dense, thr=1e-10, c=[0.5, 0.5, 1], alpha=0.5
+)
+plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20, zorder=3)
+plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20, zorder=3)
 plt.axis("equal")
 plt.title("Dense OT: Optimal Transport Plan")
 
-# Sparse OT
 plt.subplot(1, 2, 2)
-if log_sparse["flow_sources"] is not None:
-    for i, j, v in zip(
-        log_sparse["flow_sources"],
-        log_sparse["flow_targets"],
-        log_sparse["flow_values"],
-    ):
-        if v > 1e-10:
-            plt.plot(
-                [X_large[i, 0], Y_large[j, 0]],
-                [X_large[i, 1], Y_large[j, 1]],
-                color="blue",
-                alpha=0.7,
-                linewidth=1.5,
-            )
-plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20)
-plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20)
+ot.plot.plot2D_samples_mat(
+    X_large, Y_large, G_sparse, thr=1e-10, c=[0.5, 0.5, 1], alpha=0.5
+)
+plt.scatter(X_large[:, 0], X_large[:, 1], c="r", marker="o", s=20, zorder=3)
+plt.scatter(Y_large[:, 0], Y_large[:, 1], c="b", marker="x", s=20, zorder=3)
 plt.axis("equal")
 plt.title("Sparse OT: Optimal Transport Plan")