[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 97.06%. Comparing base (5cb82c9) to head (8c33b48).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #863      +/-   ##
==========================================
+ Coverage   97.00%   97.06%   +0.05%     
==========================================
  Files         130      132       +2     
  Lines       26842    27342     +500     
==========================================
+ Hits        26039    26539     +500     
  Misses        803      803              

🚀 New features to boost your workflow:

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

[Copilot]

🟡 Changes recommended

Incorrect mass handling and unchecked native buffer geometry leave correctness and memory-safety issues unresolved.

3 open issues

• Reconcile float32 mass rounding before solving · New
• Pin CDF endpoints to the actual shared mass · New
• Validate grid shape and buffer lengths before native calls · New

What changed in this PR

Adds ot.lp.emd_grid_l1 for exact cityblock transport between histograms on a shared Cartesian grid, reusing POT’s network simplex solver on sparse grid adjacency rather than a full bipartite graph.

Changes:

• Adds the grid solver, optional transport plans, and dual-potential gradients.
• Provides a backend-native 1D path and updates backend arange dtype/device handling.
• Adds correctness, gradient, backend, and dispatch tests, plus packaging and release notes.

File	Description
test/​test_grid.py	Tests grid costs, plans, gradients, backends, and dispatch.
test/​test_backend.py	Tests arange dtype/device preservation.
RELEASES.md	Announces the grid solver.
ot/​lp/​sparse_digraph.h	Adds a general sparse directed graph.
ot/​lp/​solver_circle.py	Handles scalar shifts with updated arange behavior.
ot/​lp/​EMD.h	Declares the native grid solver interface.
ot/​lp/​EMD_wrapper.cpp	Implements grid flow solving and plan recovery.
ot/​lp/​emd_wrap.pyx	Exposes the native solver through Cython.
ot/​lp/​_grid.py	Adds the public API and backend-native 1D implementation.
ot/​lp/​__init__.py	Exports emd_grid_l1.
ot/​backend.py	Updates arange to respect type_as.
MANIFEST.in	Includes the new graph header in source distributions.

_{Review effort: Balanced}

[tvercaut]

Fixed in 0fc955a. Following ot.emd, B is now rescaled to A's total mass before solving: in float64 for the d >= 2 path (so the C++ and network simplex mass-balance checks see matching totals), and backend-natively in the 1D path, which no longer has its own tight tolerance check. Only a zero-mass B against a nonzero A is still reported infeasible. Added float32 regression tests for 1D, 2D and 3D grids that check both the cost and the plan marginals (they fail on the previous code), and updated the mass-mismatch tests, which now expect the rescaled solution when check_marginals=False.

[tvercaut]

Fixed in e96b7f0. Both CDF endpoints are now pinned to the shared total mass (A's, which B is rescaled to) rather than 1, and the comments are updated. Added plan-marginal tests for non-unit totals (including A=[2, 0], B=[0, 2]) and for zero total mass, which now gives an all-zero plan.

[tvercaut]

Fixed in dcf48c9, with checks at each layer:

• emd_grid_l1 raises ValueError for 0-d inputs (the segfault you found) and for empty grids.
• emd_c_grid_l1 raises ValueError unless the shape is non-empty with positive extents and both a and b have exactly prod(shape) entries, before any pointer is passed to C++.
• EMD_wrap_grid_l1 itself also rejects an empty shape (and node counts that would overflow its int node ids), on top of its existing check on non-positive extents.

Added tests for both the public function and the Cython wrapper.

While checking how other solvers handle empty inputs, I found that ot.emd crashes the interpreter on an empty cost matrix; filed separately as #901, since this PR doesn't touch that code.

[rflamary]

Hello @tvercaut this is a big PR. I asked for a check to copilot and will also do a coee rveiw myself but it might take some time. But it looks great and we will definitely merge it

[tvercaut]

Thanks, no worries re time to go through it. A first feedback on the user-facing python API would be helpful as I am not sure I followed the POT pattern well enough. The co-pilot review looks valid and relatively easy to fix.

[tvercaut]

Forgot to mention: this PR solves #864 "on the way" as it needs it fixed to work. If it makes reviewing easier, I can open a small PR just for #864 (i.e. commit tvercaut@d83455e) which could be merged quicker and thus remove some code from this PR.