[WIP] BSP-OT: Sparse transport plans between discrete measures in loglinear time

.circleci/config.yml

@@ -4,8 +4,8 @@ version: 2
 jobs:
     build_docs:
       docker:
-        - image: cimg/python:3.10
-      resource_class: medium
+        - image: cimg/python:3.12
+      resource_class: large
       steps:
         - checkout
         - run:
@@ -27,6 +27,7 @@ jobs:
                 git remote add upstream https://github.com/PythonOT/POT.git;
                 git pull --ff-only upstream "refs/pull/$(cat merge.txt)/merge";
                 git fetch upstream master;
+                git submodule update --init --recursive
               fi
 
         # Load our data

.github/workflows/build_doc.yml

@@ -14,6 +14,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
     # Standard drop-in approach that should work for most people.
     - name: Free Disk Space (Ubuntu)
       uses: insightsengineering/disk-space-reclaimer@v1
@@ -32,7 +34,7 @@ jobs:
         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 sphinx-gallery memory_profiler
-        python -m pip install --user -e .
+        python -m pip install -v --user -e .
     # Look at what we have and fail early if there is some library conflict
     - name: Check installation
       run: |

.github/workflows/build_tests.yml

@@ -28,6 +28,8 @@ jobs:
 
     - name: Checking Out Repository
       uses: actions/checkout@v4
+      with:
+        submodules: true
     # Install Python & Packages
     - uses: actions/setup-python@v4
       with:
@@ -43,6 +45,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
+
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -70,6 +75,9 @@ jobs:
         android: true
         dotnet: true
     - uses: actions/checkout@v4
+      with:
+        submodules: true
+
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
@@ -95,6 +103,9 @@ jobs:
     if: "!contains(github.event.head_commit.message, 'no ci')"
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
+
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -122,6 +133,8 @@ jobs:
 
      steps:
      - uses: actions/checkout@v4
+       with:
+        submodules: true
      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
@@ -150,6 +163,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:

.github/workflows/build_wheels.yml

@@ -19,6 +19,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
     - name: Set up Python 3.10
       uses: actions/setup-python@v5
       with:
@@ -54,6 +56,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
     - name: Set up Python 3.10
       uses: actions/setup-python@v5
       with:

.github/workflows/build_wheels_weekly.yml

@@ -18,6 +18,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: true
     - name: Set up Python 3.11
       uses: actions/setup-python@v5
       with:

.gitignore

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

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "deps/eigen"]
+	path = deps/eigen
+	url = https://gitlab.com/libeigen/eigen.git
+	branch = 3.4

MANIFEST.in

@@ -12,3 +12,8 @@ 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
+include ot/bsp/BSP-OT_header_only.h
+include ot/bsp/bsp_wrapper.cpp
+include ot/bsp/bsp_wrapper.h
+include ot/bsp/bsp_wrap.pyx
+graft deps/eigen/Eigen

README.md

@@ -451,3 +451,5 @@ Artificial Intelligence.
 [81] Xu, H., Luo, D., & Carin, L. (2019). [Scalable Gromov-Wasserstein learning for graph partitioning and matching](https://proceedings.neurips.cc/paper/2019/hash/6e62a992c676f611616097dbea8ea030-Abstract.html). Neural Information Processing Systems (NeurIPS).
 
 [82] Bonet, C., Nadjahi, K., Séjourné, T., Fatras, K., & Courty, N. (2024). [Slicing Unbalanced Optimal Transport](https://openreview.net/forum?id=AjJTg5M0r8). Transactions on Machine Learning Research.
+
+[83] Genest, B., Bonneel, N., Nivoliers, V., & Coeurjolly, D. (2025). [BSP-OT: Sparse transport plans between discrete measures in loglinear time.](https://dl.acm.org/doi/10.1145/3763281) ACM Transactions on Graphics (TOG), 44(6), 1-15.

RELEASES.md

@@ -12,6 +12,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver
 - Migrate backend from deprecated `scipy.sparse.coo_matrix` to modern `scipy.sparse.coo_array` (PR #782)
 - Geomloss function now handles both scalar and slice indices for i and j (PR #785)
 - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397)
+- Add "BSP-OT: Sparse transport plans between discrete measures in loglinear time" (PR #768)
 - Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765)
 - Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765)
 

docs/source/all.rst

@@ -17,6 +17,7 @@ API and modules
    backend
    batch
    bregman
+   bsp
    coot
    da
    datasets

examples/plot_bsp_ot.py

@@ -0,0 +1,208 @@
+"""
+===================
+Fast and accurate bijections using BSP-OT
+===================
+
+This example shows two use cases for the bijections provided by BSP-OT,
+between two large point clouds: shape morphing (and animated morphing)
+and full color transfer (pixel permutation).
+
+[83] Genest, B., Bonneel, N., Nivoliers, V., Coeurjolly, D.
+BSP-OT: Sparse transport plans between discrete measures in log-linear time
+ACM Transactions on Graphics, Siggraph Asia (2025).
+
+"""
+
+# Author: Baptiste Genest <baptistegenest@gmail.com>
+#
+# License: MIT License
+
+import ot.utils
+import ot.bsp
+import numpy as np
+import time
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+
+##############################################################################
+# Shape interpolation and morphing
+# ----------------------------------
+
+##############################################################################
+# Data generation
+# ----------------------------------
+
+# %% Two large 2D point clouds
+# For this example, let's create two large 2D point clouds,
+# one sampled from a single ball, and the other from two smaller balls.
+
+
+def sample_ball(n, radius=1.0, center=(0.0, 0.0)):
+    np.random.seed(0)
+    theta = 2 * np.pi * np.random.rand(n)
+    r = radius * np.sqrt(np.random.rand(n))
+
+    x = r * np.cos(theta) + center[0]
+    y = r * np.sin(theta) + center[1]
+
+    return np.stack((x, y), axis=1)
+
+
+def sample_two_balls(n, radius=1.0, centers=((-1, 0), (1, 1))):
+    assert n % 2 == 0, "n must be even"
+
+    n_half = n // 2
+    X1 = sample_ball(n_half, radius, centers[0])
+    X2 = sample_ball(n_half, radius, centers[1])
+
+    return np.vstack((X1, X2))
+
+
+# ----------------------------
+# Load point clouds
+# ----------------------------
+N = 100000
+X_s = sample_ball(N)
+X_t = sample_two_balls(N, 0.5)
+
+N = X_s.shape[0]
+
+
+##############################################################################
+# Bijection computation
+# ----------------------------------
+
+ot.utils.tic()
+
+# %% call BSP-OT solver
+# The solver returns the transport cost, the final bijection and the
+# intermediary ones used to compute the final one (here we set k = 64).
+# Here we only use the final bijection.
+cost, perm, _ = ot.bsp.compute_bspot_bijection(X_s, X_t, 64, 2)
+print(
+    "Bijection computed between {} points, with cost {} in {}s".format(
+        N, cost, ot.utils.toc()
+    )
+)
+
+# %% Reordering
+# As the plan is a bijection, it is simply stored as permutation (e.g. a list of numbers)
+# such that X_s[i] is assigned to X_t[perm[i]].
+# For the sake of the animation, we reorder X_t according to the obtained bijection
+# such that the points are in correspondence along the morphing animation
+# using simply X_s*(1-t) + X_t*t
+X_t_perm = X_t[perm]
+
+##############################################################################
+# Setup animation
+# ----------------------------
+
+# Animation parameters
+FRAMES = 100
+INTERVAL = 20
+
+# Plot setup
+fig, ax = plt.subplots(figsize=(6, 6))
+ax.set_aspect("equal")
+ax.set_title("Bijective Point Cloud Morphing")
+
+scat = ax.scatter(X_s[:, 0], X_s[:, 1], s=0.05, c="tab:blue")
+
+all_pts = np.vstack([X_s, X_t_perm])
+pad = 0.5
+ax.set_xlim(all_pts[:, 0].min() - pad, all_pts[:, 0].max() + pad)
+ax.set_ylim(all_pts[:, 1].min() - pad, all_pts[:, 1].max() + pad)
+
+
+# ----------------------------
+# Animation update
+# ----------------------------
+def update(frame):
+    t = frame / (FRAMES - 1)
+    t = t * t * (3 - 2 * t)
+    P = (1 - t) * X_s + t * X_t_perm
+    scat.set_offsets(P)
+    return (scat,)
+
+
+# ----------------------------
+# Run animation
+# ----------------------------
+ani = FuncAnimation(fig, update, frames=FRAMES, interval=INTERVAL, blit=True)
+
+plt.show()
+
+
+##############################################################################
+# Other example: full color transfer (pixel permutation) using BSP-OT
+# ----------------------------
+
+import os
+from pathlib import Path
+import numpy as np
+from PIL import Image
+
+
+def im2mat(img):
+    """Converts an image to matrix (one pixel per line)"""
+    return img.reshape((img.shape[0] * img.shape[1], img.shape[2]))
+
+
+def mat2im(X, shape):
+    """Converts back a matrix to an image"""
+    return X.reshape(shape)
+
+
+def minmax(img):
+    return np.clip(img, 0, 1)
+
+
+this_file = os.path.realpath("__file__")
+data_path = os.path.join(Path(this_file).parent.parent, "data")
+
+
+I1_pil = Image.open(os.path.join(data_path, "ocean_day.jpg")).convert("RGB")
+I2_pil = Image.open(os.path.join(data_path, "ocean_sunset.jpg")).convert("RGB")
+
+# force same size to obtain bijection between all pixels
+I1_pil = I1_pil.resize(I2_pil.size, Image.BILINEAR)
+
+
+I1 = np.asarray(I1_pil).astype(np.float64) / 255.0
+I2 = np.asarray(I2_pil).astype(np.float64) / 255.0
+
+X1 = im2mat(I1)
+X2 = im2mat(I2)
+
+
+start = ot.utils.tic()
+
+_, perm, _ = ot.bsp.compute_bspot_bijection(X1, X2, 16)
+
+print("Bijection computed between {} pixels in {}s".format(X1.shape[0], ot.utils.toc()))
+
+# reorder the second image according to the obtained bijection
+X2_perm = X2[perm]
+
+# reshape back to image
+I2_perm = mat2im(X2_perm, I1.shape)
+
+plt.figure(figsize=(12, 6))
+
+plt.subplot(1, 3, 1)
+plt.imshow(I2)
+plt.title("Color image")
+plt.axis("off")
+
+plt.subplot(1, 3, 2)
+plt.imshow(I1)
+plt.title("Target image")
+plt.axis("off")
+
+plt.subplot(1, 3, 3)
+plt.imshow(I2_perm)
+plt.title("Permuted image")
+plt.axis("off")
+
+plt.tight_layout()
+plt.show()

ot/__init__.py

@@ -82,10 +82,13 @@
 from .lowrank import lowrank_sinkhorn
 
 from .batch import solve_batch, solve_sample_batch, solve_gromov_batch, dist_batch
+from .bsp import compute_bspot_bijection, merge_bijections
+
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
 
+
 __version__ = "0.9.7.dev0"
 
 __all__ = [
@@ -156,4 +159,6 @@
     "solve_gromov_batch",
     "solve_sample_batch",
     "dist_batch",
+    "compute_bspot_bijection",
+    "merge_bijections",
 ]