[WIP] Add exact EMD-L1 solver for Cartesian grid histograms

[tvercaut]

Summary

Follow-up to discussion #862: adds ot.lp.emd_grid_l1(A, B), an exact solver for the Earth Mover's Distance with the cityblock (L1) ground metric between two histograms sharing the same d-dimensional Cartesian grid support.

Instead of solving the min-cost flow on the full bipartite graph between the n = prod(A.shape) source and target bins (as ot.emd/ot.emd2 would), this reduces the problem to a min-cost flow on the much sparser grid adjacency graph (O(d*n) arcs instead of O(n^2)), following the graph formulation of Ling & Okada [1]. Unlike that paper, which introduces a bespoke tree-based solver for the reduced graph, this reuses POT's existing (off-the-shelf) NetworkSimplexSimple LP solver on it — as suggested in the discussion, this is a small, well-contained addition on top of infrastructure POT already ships.

[1] Ling, H., & Okada, K. (2007). An efficient earth mover's distance algorithm for robust histogram comparison. IEEE TPAMI, 29(5), 840-853.

What's in this PR

• ot/lp/sparse_digraph.h: a new general (non-bipartite) sparse digraph. The existing sparse_bipartitegraph.h (from [MRG] Sparse emd implementation #778) splits nodes into a source half and a target half; that doesn't work here since every grid cell can carry both supply and demand at once.
• EMD_wrap_grid_l1 in EMD_wrapper.cpp/EMD.h: builds the grid adjacency graph directly from a shape array, runs NetworkSimplexSimple on it, and (only when requested) decomposes the resulting Beckmann-style arc flow into an explicit transportation plan — including the "self-mass" A and B already share at the same bin, which the flow decomposition alone misses (more on this below).
• ot/lp/_grid.py (emd_grid_l1): the Python-facing entry point.
  • A, B are passed as actual d-dimensional arrays (not flattened + a separate shape argument) so they naturally carry their own grid geometry.
  • return_plan=False by default: recovering the plan has a real cost of its own (network simplex flow decomposition, or an O(n) merge in 1D), so it's opt-in.
  • Dedicated 1D fast path. A 1D grid is just a shared, sorted, unit-spaced support, so:
    • the plan (when requested) is recovered via POT's own emd_1d_sorted, skipping the network-simplex setup entirely;
    • the cost-only case (the default) has a closed form — the L1 norm of the difference of cumulative sums — computed with generic backend reductions (nx.cumsum/nx.abs/nx.sum) and no CPU round-trip at all, verified end-to-end on an MPS GPU tensor.
  • Backend-compatible throughout (numpy/torch/jax/tf/cupy via ot.backend). The general (ndim >= 2) path and the 1D-with-plan path do need a CPU round-trip for the compiled solver, same as ot.emd/ot.emd2_lazy.
  • The sparse plan is returned as log["G"], a sparse matrix built via the backend's coo_matrix — the same mechanism and convention ot.emd2_lazy's return_matrix already uses (real sparse type for NumPy/PyTorch/TensorFlow/CuPy, densified for JAX, which has no sparse array type).
• test/test_grid.py: correctness against the dense solver on random grids (1D-4D), plan/coupling marginal checks (row/col sums of G match A/B exactly, not just the net residual), backend round-trips (including dtype/device preservation), and dispatch checks (mocking confirms the 1D path never touches the general C++ solver, and the plan-less 1D path never touches the O(n) merge either).

Benchmarks

Script attached at the bottom of this description (local_sandbox/bench_grid_l1.py, not committed). Compares emd_grid_l1 against ot.emd2 (dense) and ot.emd2_lazy for ndim 2-4, and against ot.emd2_1d (POT's own dedicated 1D solver) for ndim=1, across a range of grid resolutions. All costs match exactly wherever compared (asserted in the script); a couple of the largest dense/lazy runs needed numItermax raised from the default 100000 to 2,000,000 to actually converge on the full bipartite graph at that size (noted below).

ndim	bins/dim	n	reference solver	grid time	lazy time	reference time	lazy/grid	reference/grid
1	10	10	emd2_1d	0.05 ms	-	0.06 ms	-	1.4x
1	100	100	emd2_1d	0.04 ms	-	0.04 ms	-	1.0x
1	1000	1000	emd2_1d	0.05 ms	-	0.06 ms	-	1.3x
1	10000	10000	emd2_1d	0.13 ms	-	0.36 ms	-	2.8x
1	100000	100000	emd2_1d	1.00 ms	-	3.51 ms	-	3.5x
2	10	100	emd2 (dense)	0.09 ms	1.14 ms	0.49 ms	13.1x	5.6x
2	20	400	emd2 (dense)	0.37 ms	19.79 ms	8.87 ms	53.4x	23.9x
2	32	1024	emd2 (dense)	1.73 ms	133.31 ms	61.72 ms	77.2x	35.7x
2	50	2500	emd2 (dense)	7.89 ms	1204.34 ms	615.41 ms	152.7x	78.0x
2	100	10000	emd2 (dense)	153.00 ms	-	-	-	-
3	5	125	emd2 (dense)	0.12 ms	1.46 ms	0.58 ms	12.0x	4.8x
3	8	512	emd2 (dense)	0.69 ms	38.65 ms	15.32 ms	56.1x	22.2x
3	10	1000	emd2 (dense)	2.02 ms	153.81 ms	62.38 ms	76.0x	30.8x
3	20	8000	emd2 (dense)	74.13 ms	22870.76 ms	10820.37 ms	308.5x	146.0x
3	30	27000	emd2 (dense)	902.89 ms	-	-	-	-
4	4	256	emd2 (dense)	0.32 ms	7.99 ms	3.05 ms	25.3x	9.6x
4	5	625	emd2 (dense)	1.27 ms	72.92 ms	29.14 ms	57.6x	23.0x
4	6	1296	emd2 (dense)	3.97 ms	351.34 ms	141.86 ms	88.4x	35.7x
4	10	10000	emd2 (dense)	129.96 ms	-	-	-	-

Notes:

• - entries: dense/lazy skipped above n=2000 nodes to keep the benchmark's runtime bounded (full bipartite network simplex gets slow, and dense M becomes memory-prohibitive); emd_grid_l1 is still timed there to show scaling. The (3, 30) row (27,000 nodes) mirrors the 30×30×30 example from the discussion.
• In 1D, the speedup over emd2_1d is modest (1.0x-3.5x) rather than the 1000x+ margin seen against the generic bipartite solvers — expected, since emd2_1d is already the right O(n log n) tool for a shared sorted support, and emd_grid_l1's own 1D plan path literally calls the same underlying routine. The gap that remains is emd_grid_l1's closed-form, network-simplex-free cost-only path.
• For ndim >= 2, speedup over dense/lazy grows with n as expected, up to ~150x/300x at n in the low thousands.

Deliberately out of scope for this PR

This PR is scoped to L1 only, with no gradient support. If this direction is OK for integration, natural follow-ups would be:

• L2 support via the (d+1)-partite graph construction of Auricchio et al. (also discussed in Faster exact EMD for distributions on shared multi-dimensional cartesian grids (with cityblock or sqeuclidean ground metric) #862), as a separate solver alongside this one.
• Gradient support, most likely via the envelope theorem using dual potentials (net.potential() isn't currently plumbed through EMD_wrap_grid_l1), mirroring how ot.emd/ot.emd2_lazy do it. [EDIT: Gradient computation now supported]

Question: batched support?

Should emd_grid_l1 support a batch dimension (e.g. A, B of shape (B, n_1, ..., n_d))? The generic exact solvers in POT (ot.emd/ot.emd2/ot.emd2_lazy) don't batch — network simplex is a sequential pivoting algorithm per problem, so there's no way to vectorize it the way Sinkhorn/proximal iterations are batched in ot.batch. That's still true here: the Ling & Okada reduction is inherently per-pair and CPU-oriented. But since each pair in a batch is fully independent, one option specific to this solver would be to parallelize across the batch in C++ (e.g. an OpenMP parallel for over batch entries, each running its own NetworkSimplexSimple instance), rather than trying to vectorize within a single solve. Happy to prototype this as a follow-up if there's interest, but wanted to flag it now rather than bake in an API that doesn't anticipate it.

Test plan

• pytest test/test_grid.py (32 tests: correctness vs. dense solver on 1D-4D grids, plan/coupling marginals, backend round-trips incl. GPU, dispatch mocking)
• Doctest in ot/lp/_grid.py
• Full existing test suite (pytest test/, excluding slow gromov/unbalanced/sliced/batch dirs): 1719 passed, no regressions
• ruff check/ruff format clean
• Manual cross-check against ot.emd2/ot.emd2_lazy/ot.emd2_1d on random grids up to 27,000 nodes (see benchmarks above)

Benchmarking script (not committed, for reference)

local_sandbox/bench_grid_l1.py

"""
Benchmark ot.lp.emd_grid_l1 against reference EMD solvers on random grid
histograms of increasing dimension and resolution.

For ndim >= 2, the references are ot.emd2_lazy and ot.emd2 (dense), the two
generic bipartite solvers emd_grid_l1 is meant to replace for this problem.
For ndim == 1, the fair reference is POT's own dedicated 1D solver,
ot.emd2_1d (an O(n log n) sort-based algorithm -- the closest existing POT
counterpart to emd_grid_l1's O(n) closed form in that case), so the generic
dense/lazy bipartite solvers are skipped there.

Not part of the test suite: a development artefact used to produce numbers
for the PR description. Run with:

    python local_sandbox/bench_grid_l1.py
"""

import time

import numpy as np

import ot
from ot.lp import emd_grid_l1

# (ndim, bins_per_dim) configs. Kept small for ndim >= 3 so the whole script
# runs in a couple of minutes; dense/lazy are skipped above SKIP_DENSE_ABOVE
# / SKIP_LAZY_ABOVE nodes to avoid blowing up the runtime, while emd_grid_l1
# is still timed there to show how far it scales.
CONFIGS = [
    (1, 10),
    (1, 100),
    (1, 1000),
    (1, 10000),
    (1, 100000),
    (2, 10),
    (2, 20),
    (2, 32),
    (2, 50),
    (2, 100),
    (3, 5),
    (3, 8),
    (3, 10),
    (3, 20),
    (3, 30),  # matches the 30x30x30 example in the GitHub discussion
    (4, 4),
    (4, 5),
    (4, 6),
    (4, 10),
]

SKIP_DENSE_ABOVE = 2000
SKIP_LAZY_ABOVE = 2000
REPEATS = 3

# Explicitly requested larger reference points (ndim >= 2): run dense/lazy
# here too even though they exceed the skip thresholds above. Single repeat
# only, since these are the expensive end of the range. The default
# numItermax=100000 is not enough for network simplex to converge on a full
# bipartite graph at this size, so it is raised here too.
FORCE_FULL = {(2, 50), (3, 20)}
FORCE_FULL_REPEATS = 1
FORCE_FULL_NUM_ITER_MAX = 2_000_000


def make_grid_coords(shape):
    axes = [np.arange(s) for s in shape]
    return (
        np.array(np.meshgrid(*axes, indexing="ij"))
        .reshape(len(shape), -1)
        .T.astype(np.float64)
    )


def random_histograms(n, rng):
    a = rng.rand(n)
    a /= a.sum()
    b = rng.rand(n)
    b /= b.sum()
    return a, b


def time_call(fn, repeats=REPEATS):
    times = []
    result = None
    for _ in range(repeats):
        t0 = time.perf_counter()
        result = fn()
        times.append(time.perf_counter() - t0)
    return result, min(times)


def bench_config(ndim, bins_per_dim, seed=0):
    shape = (bins_per_dim,) * ndim
    n = bins_per_dim**ndim
    rng = np.random.RandomState(seed)
    a, b = random_histograms(n, rng)
    A, B = a.reshape(shape), b.reshape(shape)
    forced = (ndim, bins_per_dim) in FORCE_FULL
    repeats = FORCE_FULL_REPEATS if forced else REPEATS

    row = {"ndim": ndim, "bins_per_dim": bins_per_dim, "n": n}

    numItermax = FORCE_FULL_NUM_ITER_MAX if forced else 100000
    row["numItermax"] = numItermax
    row["lazy_converged"] = None
    row["dense_converged"] = None

    row["grid_cost"], row["grid_time"] = time_call(lambda: emd_grid_l1(A, B))

    if ndim == 1:
        # For 1D, the fair "reference solver" is POT's own dedicated 1D
        # solver (also what emd_grid_l1's own 1D path calls under the hood
        # for the plan, via emd_1d_sorted) rather than the generic
        # dense/lazy bipartite network simplex. No lazy column here.
        row["dense_label"] = "emd2_1d"
        row["lazy_cost"], row["lazy_time"] = None, None
        x = np.arange(n, dtype=np.float64)
        t0 = time.perf_counter()
        row["dense_cost"], row["dense_time"] = time_call(
            lambda: ot.emd2_1d(x, x, a, b, metric="cityblock"), repeats=repeats
        )
        print(
            f"  [emd2_1d ndim=1 bins={bins_per_dim}] took "
            f"{time.perf_counter() - t0:.1f}s"
        )
        return row

    row["dense_label"] = "emd2 (dense)"

    if n <= SKIP_LAZY_ABOVE or forced:
        coords = make_grid_coords(shape)
        t0 = time.perf_counter()

        def _run_lazy():
            cost, lazy_log = ot.emd2_lazy(
                coords,
                coords,
                a,
                b,
                metric="cityblock",
                return_matrix=False,
                numItermax=numItermax,
                log=True,
            )
            row["lazy_converged"] = lazy_log["result_code"] == 1
            return cost

        row["lazy_cost"], row["lazy_time"] = time_call(_run_lazy, repeats=repeats)
        print(
            f"  [lazy  ndim={ndim} bins={bins_per_dim}] took "
            f"{time.perf_counter() - t0:.1f}s, converged={row['lazy_converged']}"
        )
    else:
        row["lazy_cost"], row["lazy_time"] = None, None

    if n <= SKIP_DENSE_ABOVE or forced:
        coords = make_grid_coords(shape)
        M = ot.dist(coords, coords, metric="cityblock")
        t0 = time.perf_counter()

        def _run_dense():
            cost, dense_log = ot.emd2(a, b, M, numItermax=numItermax, log=True)
            row["dense_converged"] = dense_log["result_code"] == 1
            return cost

        row["dense_cost"], row["dense_time"] = time_call(_run_dense, repeats=repeats)
        print(
            f"  [dense ndim={ndim} bins={bins_per_dim}] took "
            f"{time.perf_counter() - t0:.1f}s, converged={row['dense_converged']}"
        )
    else:
        row["dense_cost"], row["dense_time"] = None, None

    return row


def fmt_time(t):
    return "-" if t is None else f"{t * 1000:8.2f} ms"


def fmt_cost(c):
    return "-" if c is None else f"{c:10.6f}"


def main():
    rows = []
    t_start = time.perf_counter()
    header = (
        f"{'ndim':>4} {'bins':>5} {'n':>7} | "
        f"{'grid cost':>10} {'lazy cost':>10} {'dense cost':>10} | "
        f"{'grid time':>11} {'lazy time':>11} {'dense time':>11} | speedup(lazy/grid) speedup(dense/grid)"
    )
    print(header)
    print("-" * len(header))
    for ndim, bins in CONFIGS:
        row = bench_config(ndim, bins)
        rows.append(row)

        speedup_lazy = (
            "-"
            if row["lazy_time"] is None
            else f"{row['lazy_time'] / row['grid_time']:8.1f}x"
        )
        speedup_dense = (
            "-"
            if row["dense_time"] is None
            else f"{row['dense_time'] / row['grid_time']:8.1f}x"
        )

        print(
            f"{row['ndim']:>4} {row['bins_per_dim']:>5} {row['n']:>7} | "
            f"{fmt_cost(row['grid_cost'])} {fmt_cost(row['lazy_cost'])} {fmt_cost(row['dense_cost'])} | "
            f"{fmt_time(row['grid_time'])} {fmt_time(row['lazy_time'])} {fmt_time(row['dense_time'])} | "
            f"{speedup_lazy:>18} {speedup_dense:>19}  (reference: {row['dense_label']})"
        )

        # Sanity checks: costs must agree across solvers, but only when the
        # reference solver actually reports having converged -- an
        # unconverged network simplex returns a feasible but suboptimal
        # (too high) cost, which is expected, not a bug.
        if row["dense_cost"] is not None:
            if row["dense_converged"] is False:
                print(
                    f"  NOTE: dense did not converge (numItermax={row['numItermax']}); "
                    f"cost {row['dense_cost']:.6f} vs grid {row['grid_cost']:.6f} "
                    "(dense is expected to be >= grid's exact cost)"
                )
                assert row["dense_cost"] >= row["grid_cost"] - 1e-6
            else:
                np.testing.assert_allclose(
                    row["grid_cost"], row["dense_cost"], rtol=1e-5, atol=1e-7
                )
        if row["lazy_cost"] is not None:
            if row["lazy_converged"] is False:
                print(
                    f"  NOTE: lazy did not converge (numItermax={row['numItermax']}); "
                    f"cost {row['lazy_cost']:.6f} vs grid {row['grid_cost']:.6f} "
                    "(lazy is expected to be >= grid's exact cost)"
                )
                assert row["lazy_cost"] >= row["grid_cost"] - 1e-6
            else:
                np.testing.assert_allclose(
                    row["grid_cost"], row["lazy_cost"], rtol=1e-5, atol=1e-7
                )

    print(f"\nTotal wall time: {time.perf_counter() - t_start:.1f} s")

    print("\nMarkdown table:\n")
    print(
        "| ndim | bins/dim | n | reference solver | grid cost | lazy cost | "
        "reference cost | grid time | lazy time | reference time | lazy/grid | "
        "reference/grid |"
    )
    print("|---|---|---|---|---|---|---|---|---|---|---|---|")
    for row in rows:
        speedup_lazy = (
            "-"
            if row["lazy_time"] is None
            else f"{row['lazy_time'] / row['grid_time']:.1f}x"
        )
        speedup_dense = (
            "-"
            if row["dense_time"] is None
            else f"{row['dense_time'] / row['grid_time']:.1f}x"
        )
        lazy_mark = "&dagger;" if row["lazy_converged"] is False else ""
        dense_mark = "&dagger;" if row["dense_converged"] is False else ""
        print(
            f"| {row['ndim']} | {row['bins_per_dim']} | {row['n']} "
            f"| {row['dense_label']} "
            f"| {fmt_cost(row['grid_cost']).strip()} "
            f"| {fmt_cost(row['lazy_cost']).strip()}{lazy_mark} "
            f"| {fmt_cost(row['dense_cost']).strip()}{dense_mark} "
            f"| {fmt_time(row['grid_time']).strip()} "
            f"| {fmt_time(row['lazy_time']).strip()} "
            f"| {fmt_time(row['dense_time']).strip()} "
            f"| {speedup_lazy} | {speedup_dense} |"
        )
    print(
        "\n&dagger; did not fully converge within its numItermax "
        f"({FORCE_FULL_NUM_ITER_MAX}); cost shown is an upper bound on the "
        "true optimum (which emd_grid_l1's cost matches exactly)."
    )


if __name__ == "__main__":
    main()

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.92%. Comparing base (98d09a1) to head (c35d037).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #863      +/-   ##
==========================================
+ Coverage   96.86%   96.92%   +0.05%     
==========================================
  Files         128      130       +2     
  Lines       26304    26755     +451     
==========================================
+ Hits        25480    25931     +451     
  Misses        824      824              

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.