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

MANIFEST.in

@@ -11,6 +11,7 @@ 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/lp/sparse_digraph.h
 include ot/partial/partial_cython.pyx
 include ot/bsp/BSP-OT_header_only.h
 include ot/bsp/bsp_wrapper.cpp

RELEASES.md

@@ -5,6 +5,7 @@
 #### New features
 
 
+- Add `ot.lp.emd_grid_l1`, an exact EMD-L1 solver for histograms sharing a multi-dimensional Cartesian grid support. It solves a min-cost flow on the grid's adjacency graph rather than the full bipartite graph (Ling & Okada, 2007), which is one to two orders of magnitude faster for this case (PR #863)
 - `ot.dr.wda` gains `solver='torch'`, a PyTorch autodiff solver with Riemannian gradient descent, usable on installations without autograd or pymanopt, and accepting torch tensors directly (PR #858, Issue #806)
 - `ot.dr` dependencies (autograd, pymanopt, scikit-learn, torch) are now imported optionally, so importing `ot.dr` no longer requires all of them; each function raises an explicit `ImportError` naming what it needs (PR #858)
 - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856)

ot/backend.py

@@ -1292,7 +1292,10 @@ def ones(self, shape, type_as=None):
             return np.ones(shape, dtype=type_as.dtype)
 
     def arange(self, stop, start=0, step=1, type_as=None):
-        return np.arange(start, stop, step)
+        if type_as is None:
+            return np.arange(start, stop, step)
+        else:
+            return np.arange(start, stop, step, dtype=type_as.dtype)
 
     def full(self, shape, fill_value, type_as=None):
         if type_as is None:
@@ -1732,7 +1735,12 @@ def ones(self, shape, type_as=None):
             return self._change_device(jnp.ones(shape, dtype=type_as.dtype), type_as)
 
     def arange(self, stop, start=0, step=1, type_as=None):
-        return jnp.arange(start, stop, step)
+        if type_as is None:
+            return jnp.arange(start, stop, step)
+        else:
+            return self._change_device(
+                jnp.arange(start, stop, step, dtype=type_as.dtype), type_as
+            )
 
     def full(self, shape, fill_value, type_as=None):
         if type_as is None:
@@ -2239,7 +2247,9 @@ def arange(self, stop, start=0, step=1, type_as=None):
         if type_as is None:
             return torch.arange(start, stop, step)
         else:
-            return torch.arange(start, stop, step, device=type_as.device)
+            return torch.arange(
+                start, stop, step, dtype=type_as.dtype, device=type_as.device
+            )
 
     def full(self, shape, fill_value, type_as=None):
         if isinstance(shape, int):
@@ -2789,7 +2799,11 @@ def ones(self, shape, type_as=None):
                 return cp.ones(shape, dtype=type_as.dtype)
 
     def arange(self, stop, start=0, step=1, type_as=None):
-        return cp.arange(start, stop, step)
+        if type_as is None:
+            return cp.arange(start, stop, step)
+        else:
+            with cp.cuda.Device(type_as.device):
+                return cp.arange(start, stop, step, dtype=type_as.dtype)
 
     def full(self, shape, fill_value, type_as=None):
         if isinstance(shape, (list, tuple)):
@@ -3242,7 +3256,10 @@ def ones(self, shape, type_as=None):
             return tnp.ones(shape, dtype=type_as.dtype)
 
     def arange(self, stop, start=0, step=1, type_as=None):
-        return tnp.arange(start, stop, step)
+        if type_as is None:
+            return tnp.arange(start, stop, step)
+        else:
+            return tnp.arange(start, stop, step, dtype=type_as.dtype)
 
     def full(self, shape, fill_value, type_as=None):
         if type_as is None:

ot/lp/EMD.h

@@ -52,6 +52,27 @@ int EMD_wrap_sparse(
     double *beta_init            // Initial dual variables for targets (warmstart)
 );
 
+int EMD_wrap_grid_l1(
+    int ndim,                    // Number of grid dimensions
+    int64_t *shape,              // Grid shape (ndim entries)
+    double *X,                   // Source histogram, flattened C-order (prod(shape))
+    double *Y,                   // Target histogram, flattened C-order (prod(shape))
+    bool return_plan,            // If false, skip decomposing the flow into a
+                                  // transportation plan and only compute cost
+    uint64_t *plan_sources_out,  // Output: source bin index of each plan entry
+    uint64_t *plan_targets_out,  // Output: target bin index of each plan entry
+    double *plan_values_out,     // Output: mass moved by each plan entry
+    uint64_t *n_plan_entries_out,
+    uint64_t max_plan_entries,
+    double *alpha,               // Output: node potentials / dual variables
+                                  // (n_nodes = prod(shape)); the Beckmann
+                                  // formulation uses a single graph (not
+                                  // bipartite), so there is one potential
+                                  // array, not alpha/beta; beta = -alpha
+    double *cost,                // Output: total transportation cost
+    uint64_t maxIter             // Maximum iterations for solver
+);
+
 int EMD_wrap_lazy(
     int n1,                      // Number of source points
     int n2,                      // Number of target points