Merge branch 'master' into dev_sliced_plans

Author: lchapel

.github/workflows/build_tests.yml

@@ -27,7 +27,7 @@ jobs:
 
 
     - name: Checking Out Repository
-      uses: actions/checkout@v2
+      uses: actions/checkout@v4
     # Install Python & Packages
     - uses: actions/setup-python@v4
       with:
@@ -39,6 +39,20 @@ jobs:
         pre-commit install --install-hooks
         pre-commit run --all-files
 
+  build_from_source:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+    - name: Build from source
+      run: |
+        python -m pip install --upgrade pip setuptools wheel
+        python -m pip install cython numpy
+        python setup.py sdist bdist_wheel
+        pip install dist/*.tar.gz
 
   linux:
 

CITATION.cff

@@ -72,6 +72,9 @@ authors:
     family-names: Fernandes Montesuma
     affiliation: Université Paris-Saclay & CEA-List
     orcid: 'https://orcid.org/0000-0003-3850-4602'
+  - given-names: Nathan
+    family-names: Neike
+    affiliation: Hi! PARIS
 identifiers:
   - type: url
     value: 'https://github.com/PythonOT/POT'

CONTRIBUTORS.md

@@ -58,6 +58,7 @@ The contributors to this library are:
 * [Laurène David](https://github.com/laudavid) (Low rank sinkhorn, Low rank Gromov-Wasserstein samples)
 * [Julie Delon](https://judelo.github.io/) (GMM OT)
 * [Samuel Boïté](https://samuelbx.github.io/) (GMM OT)
+* [Nathan Neike](https://github.com/nathanneike) (Sparse EMD solver)
 
 
 ## Acknowledgments

MANIFEST.in

@@ -10,4 +10,5 @@ include ot/lp/full_bipartitegraph.h
 include ot/lp/full_bipartitegraph_omp.h
 include ot/lp/network_simplex_simple.h
 include ot/lp/network_simplex_simple_omp.h
+include ot/lp/sparse_bipartitegraph.h
 include ot/partial/partial_cython.pyx

examples/plot_sparse_emd.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+"""
+============================================
+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. 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
+#
+# License: MIT License
+# sphinx_gallery_thumbnail_number = 2
+
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.sparse import coo_array
+import ot
+
+##############################################################################
+# Example: concentric circles
+# -----------------------------------
+
+# %%
+
+n_clusters = 8
+points_per_cluster = 25
+n = n_clusters * points_per_cluster
+k_neighbors = 8
+rng = np.random.default_rng(0)
+
+r_source = 1.0
+r_target = 2.0
+noise_scale = 0.06
+
+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.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(cluster_labels == k)[0]
+    tgt_idx = np.where(cluster_labels == k)[0]
+    for i in src_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_array((vals, (rows, cols)), shape=(n, n))
+allowed_sparse = set(zip(rows, cols))
+
+##############################################################################
+# Visualize edge structures
+# --------------------------
+
+# %%
+
+plt.figure(figsize=(16, 6))
+
+plt.subplot(1, 2, 1)
+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.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")
+
+plt.subplot(1, 2, 2)
+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: Intra-Cluster k-NN Edges")
+
+plt.tight_layout()
+plt.show()
+
+##############################################################################
+# Solve and visualize transport plans
+# ------------------------------------
+
+# %%
+
+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}")
+
+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}")
+
+plt.figure(figsize=(16, 6))
+
+plt.subplot(1, 2, 1)
+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")
+
+plt.subplot(1, 2, 2)
+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")
+
+plt.tight_layout()
+plt.show()

ot/__init__.py

@@ -79,7 +79,7 @@
 # utils functions
 from .utils import dist, unif, tic, toc, toq
 
-__version__ = "0.9.6.post1"
+__version__ = "0.9.7.dev0"
 
 __all__ = [
     "emd",

ot/backend.py

@@ -119,7 +119,7 @@
         import jax
         import jax.numpy as jnp
         import jax.scipy.special as jspecial
-        from jax.lib import xla_bridge
+        from jax.extend.backend import get_backend as _jax_get_backend
 
         jax_type = jax.numpy.ndarray
         jax_new_version = float(".".join(jax.__version__.split(".")[1:])) > 4.24
@@ -178,7 +178,16 @@ def _get_backend_instance(backend_impl):
 
 
 def _check_args_backend(backend_impl, args):
-    is_instance = set(isinstance(arg, backend_impl.__type__) for arg in args)
+    # Get backend instance to use issparse method
+    backend = _get_backend_instance(backend_impl)
+
+    # Check if each arg is either:
+    # 1. An instance of backend.__type__ (e.g., np.ndarray for NumPy)
+    # 2. A sparse matrix recognized by backend.issparse() (e.g., scipy.sparse for NumPy)
+    is_instance = set(
+        isinstance(arg, backend_impl.__type__) or backend.issparse(arg) for arg in args
+    )
+
     # check that all arguments matched or not the type
     if len(is_instance) == 1:
         return is_instance.pop()
@@ -839,6 +848,31 @@ def todense(self, a):
         """
         raise NotImplementedError()
 
+    def sparse_coo_data(self, a):
+        r"""
+        Extracts COO format data (row, col, data, shape) from a sparse matrix.
+
+        Returns row indices, column indices, data values, and shape as numpy arrays/tuple.
+        This is used to interface with C++ solvers that require explicit edge lists.
+
+        Parameters
+        ----------
+        a : sparse matrix
+            Sparse matrix in backend's COO format
+
+        Returns
+        -------
+        row : numpy.ndarray
+            Row indices (1D array)
+        col : numpy.ndarray
+            Column indices (1D array)
+        data : numpy.ndarray
+            Data values (1D array)
+        shape : tuple
+            Shape of the matrix (n_rows, n_cols)
+        """
+        raise NotImplementedError()
+
     def where(self, condition, x, y):
         r"""
         Returns elements chosen from x or y depending on condition.
@@ -1349,6 +1383,15 @@ def todense(self, a):
         else:
             return a
 
+    def sparse_coo_data(self, a):
+        # Convert to COO format if needed
+        if not isinstance(a, coo_matrix):
+            a_coo = coo_matrix(a)
+        else:
+            a_coo = a
+
+        return a_coo.row, a_coo.col, a_coo.data, a_coo.shape
+
     def where(self, condition, x=None, y=None):
         if x is None and y is None:
             return np.where(condition)
@@ -1509,7 +1552,7 @@ def __init__(self):
         self.__type_list__ = []
         # available_devices = jax.devices("cpu")
         available_devices = []
-        if xla_bridge.get_backend().platform == "gpu":
+        if _jax_get_backend().platform == "gpu":
             available_devices += jax.devices("gpu")
         for d in available_devices:
             self.__type_list__ += [
@@ -1768,6 +1811,15 @@ def todense(self, a):
         # Currently, JAX does not support sparse matrices
         return a
 
+    def sparse_coo_data(self, a):
+        # JAX doesn't support sparse matrices, so this shouldn't be called
+        # But if it is, convert the dense array to sparse using scipy
+        a_np = self.to_numpy(a)
+        from scipy.sparse import coo_matrix
+
+        a_coo = coo_matrix(a_np)
+        return a_coo.row, a_coo.col, a_coo.data, a_coo.shape
+
     def where(self, condition, x=None, y=None):
         if x is None and y is None:
             return jnp.where(condition)
@@ -1938,6 +1990,7 @@ def __init__(self):
             self.rng_cuda_ = torch.Generator("cpu")
 
         from torch.autograd import Function
+        from torch.autograd.function import once_differentiable
 
         # define a function that takes inputs val and grads
         # ad returns a val tensor with proper gradients
@@ -1952,7 +2005,31 @@ def backward(ctx, grad_output):
                 # the gradients are grad
                 return (None, None) + tuple(g * grad_output for g in ctx.grads)
 
+        # define a differentiable SPD matrix sqrt
+        # with closed-form VJP
+        class MatrixSqrtFunction(Function):
+            @staticmethod
+            def forward(ctx, a):
+                a_sym = 0.5 * (a + a.transpose(-2, -1))
+                L, V = torch.linalg.eigh(a_sym)
+                s = L.clamp_min(0).sqrt()
+                y = (V * s.unsqueeze(-2)) @ V.transpose(-2, -1)
+                ctx.save_for_backward(s, V)
+                return y
+
+            @staticmethod
+            @once_differentiable
+            def backward(ctx, g):
+                s, V = ctx.saved_tensors
+                g_sym = 0.5 * (g + g.transpose(-2, -1))
+                ghat = V.transpose(-2, -1) @ g_sym @ V
+                d = s.unsqueeze(-1) + s.unsqueeze(-2)
+                xhat = ghat / d
+                xhat = xhat.masked_fill(d == 0, 0)
+                return V @ xhat @ V.transpose(-2, -1)
+
         self.ValFunction = ValFunction
+        self.MatrixSqrtFunction = MatrixSqrtFunction
 
     def _to_numpy(self, a):
         if isinstance(a, float) or isinstance(a, int) or isinstance(a, np.ndarray):
@@ -2315,6 +2392,20 @@ def todense(self, a):
         else:
             return a
 
+    def sparse_coo_data(self, a):
+        # For torch sparse tensors, coalesce first to ensure unique indices
+        a_coalesced = a.coalesce()
+        indices = a_coalesced._indices()
+        values = a_coalesced._values()
+
+        # Convert to numpy
+        row = self.to_numpy(indices[0])
+        col = self.to_numpy(indices[1])
+        data = self.to_numpy(values)
+        shape = tuple(a_coalesced.shape)
+
+        return row, col, data, shape
+
     def where(self, condition, x=None, y=None):
         if x is None and y is None:
             return torch.where(condition)
@@ -2395,12 +2486,7 @@ def pinv(self, a, hermitian=False):
         return torch.linalg.pinv(a, hermitian=hermitian)
 
     def sqrtm(self, a):
-        L, V = torch.linalg.eigh(a)
-        L = torch.sqrt(L)
-        # Q[...] = V[...] @ diag(L[...])
-        Q = torch.einsum("...jk,...k->...jk", V, L)
-        # R[...] = Q[...] @ V[...].T
-        return torch.einsum("...jk,...kl->...jl", Q, torch.transpose(V, -1, -2))
+        return self.MatrixSqrtFunction.apply(a)
 
     def eigh(self, a):
         return torch.linalg.eigh(a)