[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,3 @@
+[submodule "deps/eigen"]
+	path = deps/eigen
+	url = https://gitlab.com/libeigen/eigen.git

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

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)
 
@@ -27,6 +28,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver
 - Clean documentation (PR #787)
 - Fix code coverage (PR #791)
 - Fix test of the version of jax in `ot.backend` (PR #794)
+- Reverting the openmp fix on macOS (PR #789) for macOS (PR #796)
 - Reverting the openmp fix on macOS (PR #789) for macOS (PR #797)
 - Align documentation build dependencies and doc extras (PR #801)
 

examples/others/plot_bsp_ot.py

@@ -0,0 +1,120 @@
+"""
+===================
+Fast and accurate bijections using BSP-OT example
+===================
+
+This example shows how to use the BSP-OT solver to compute a bijection
+between two large point clouds, and animate the morphing between them.
+
+[?] 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.bsp
+import numpy as np
+import time
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+
+
+
+##############################################################################
+# Data ganeration
+# ----------------------------------
+
+def sample_ball(n, radius=1.0, center=(0.0, 0.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
+A = sample_ball(N)
+B = sample_two_balls(N, 0.5)
+
+N = A.shape[0]
+
+
+##############################################################################
+# Bijection computation
+# ----------------------------------
+
+start = time.time()
+
+# %% 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.bsp_solve(A, B, 64,2)
+print(
+    "Bijection computed between {} points, with cost {} in {}s".format(
+        N, cost, time.time() - start
+    )
+)
+
+# %% Reorder B according to bijection
+# such that the points are in correspondence along the morphing animation
+# using simply A*(1-t) + B*t
+B_perm = B[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(A[:, 0], A[:, 1], s=0.05, c="tab:blue")
+
+all_pts = np.vstack([A, B_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) * A + t * B_perm
+    scat.set_offsets(P)
+    return (scat,)
+
+
+# ----------------------------
+# Run animation
+# ----------------------------
+ani = FuncAnimation(fig, update, frames=FRAMES, interval=INTERVAL, blit=True)
+
+plt.show()