Skip to content

Commit bd214c8

Browse files
nathanneikerflamary
nathanneike
and
rflamary
authored
[MRG] Refactor simplex and Lazy simplex improvement (#813)
* Refactor lazy simplex storage * Add CI native import diagnostics * Add CI native import diagnostics * renamed some config variables * got rid off useless class and more renaming * updated release file with PR --------- Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
1 parent 394f3e1 commit bd214c8

9 files changed

Lines changed: 663 additions & 316 deletions

File tree

‎RELEASES.md‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,11 @@
44

55
This new release adds support for sparse cost matrices and a new lazy EMD solver that computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m). Both implementations are backend-agnostic and preserve gradient computation for automatic differentiation.
66

7-
#### New features
7+
#### New features
88

9+
- Refactor lazy EMD network simplex storage to avoid dense per-arc cost,
10+
endpoint, flow, and state storage where possible, and return sparse lazy
11+
transport plans instead of materializing dense plans internally (PR #813)
912
- Add sliced transport plans (min-pivot sliced and expected sliced) solvers (PR #767)
1013
- Add lazy EMD solver with on-the-fly distance computation from coordinates (PR #788)
1114
- Add Warmstart feature to the EMD solver for existing potentials (PR #793)
@@ -17,17 +20,12 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver
1720
- Add "BSP-OT: Sparse transport plans between discrete measures in loglinear time" (PR #768)
1821
- Added UOT1D with Frank-Wolfe in `ot.unbalanced.uot_1d` (PR #765)
1922
- Add Sliced UOT and Unbalanced Sliced OT in `ot/unbalanced/_sliced.py` (PR #765)
20-
- Add `ot.utils.DataScaler` class for backend-aware joint normalization of input
21-
distributions, with sklearn-compatible `fit`/`transform`/`fit_transform` API and
22-
support for `'standard'`, `'minmax'`, and `'l2'` methods (PR #808)
23+
- Add `ot.utils.DataScaler` class for backend-aware joint normalization of input distributions, with sklearn-compatible `fit`/`transform`/`fit_transform` API and support for `'standard'`, `'minmax'`, and `'l2'` methods (PR #808)
2324
- Add `ot.utils.apply_scaler` helper that dispatches preprocessing to a scaler object,
2425
a callable, or a no-op (PR #808)
25-
- Add optional `scaler` parameter to `sliced_wasserstein_distance` and
26-
`max_sliced_wasserstein_distance` (PR #808)
26+
- Add optional `scaler` parameter to `sliced_wasserstein_distance` and `max_sliced_wasserstein_distance` (PR #808)
2727
- Add a numerically stable log-domain solver for entropic partial Wasserstein, selectable via the new `method` parameter of `entropic_partial_wasserstein` (`method='sinkhorn_log'`) or directly through `entropic_partial_wasserstein_logscale` (Issue #723)
28-
- Add cost functions between linear operators following
29-
[A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920),
30-
implemented in `ot.sgot` (PR #792)
28+
- Add cost functions between linear operators following [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920), implemented in `ot.sgot` (PR #792)
3129
- Build wheels on ubuntu ARM to avoid QEMU emulation (PR #818)
3230

3331
#### Closed issues

‎ot/lp/EMD.h‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ int EMD_wrap_sparse(
4444
uint64_t *flow_sources_out, // Output: source indices of non-zero flows
4545
uint64_t *flow_targets_out, // Output: target indices of non-zero flows
4646
double *flow_values_out, // Output: flow values
47-
uint64_t *n_flows_out,
47+
uint64_t *n_flows_out,
48+
uint64_t max_flows_out,
4849
double *alpha, // Output: dual variables for sources (n1)
4950
double *beta, // Output: dual variables for targets (n2)
5051
double *cost, // Output: total transportation cost
@@ -62,7 +63,11 @@ int EMD_wrap_lazy(
6263
double *coords_b, // Target coordinates (n2 x dim)
6364
int dim, // Dimension of coordinates
6465
int metric, // Distance metric: 0=sqeuclidean, 1=euclidean, 2=cityblock
65-
double *G, // Output: transport plan (n1 x n2)
66+
uint64_t *flow_sources_out, // Output: source indices of non-zero flows
67+
uint64_t *flow_targets_out, // Output: target indices of non-zero flows
68+
double *flow_values_out, // Output: flow values
69+
uint64_t *n_flows_out,
70+
uint64_t max_flows_out,
6671
double *alpha, // Output: dual variables for sources (n1)
6772
double *beta, // Output: dual variables for targets (n2)
6873
double *cost, // Output: total transportation cost

‎ot/lp/EMD_wrapper.cpp‎

Lines changed: 100 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,62 @@ inline void extract_compressed_support(
147147
}
148148
}
149149

150+
template <
151+
typename NetType,
152+
typename DigraphType,
153+
typename InvalidType,
154+
typename SourceIndexVector,
155+
typename TargetIndexVector,
156+
typename CostAccessor
157+
>
158+
inline bool extract_sparse_solution(
159+
const NetType& net,
160+
DigraphType& di,
161+
InvalidType invalid,
162+
const SourceIndexVector& idx_a,
163+
const TargetIndexVector& idx_b,
164+
double* alpha,
165+
double* beta,
166+
double* cost,
167+
uint64_t* flow_sources_out,
168+
uint64_t* flow_targets_out,
169+
double* flow_values_out,
170+
uint64_t* n_flows_out,
171+
uint64_t max_flows_out,
172+
CostAccessor cost_accessor,
173+
double min_output_flow
174+
) {
175+
const int n = static_cast<int>(idx_a.size());
176+
177+
for (int i = 0; i < n; i++) {
178+
alpha[static_cast<uint64_t>(idx_a[i])] = -net.potential(i);
179+
}
180+
for (int j = 0; j < static_cast<int>(idx_b.size()); j++) {
181+
beta[static_cast<uint64_t>(idx_b[j])] = net.potential(j + n);
182+
}
183+
184+
typename DigraphType::Arc a;
185+
di.first(a);
186+
for (; a != invalid; di.next(a)) {
187+
const int i = di.source(a);
188+
const int j = di.target(a) - n;
189+
const double flow = net.flow(a);
190+
if (flow != 0) {
191+
*cost += flow * cost_accessor(a, i, j);
192+
}
193+
if (flow > min_output_flow) {
194+
if (*n_flows_out >= max_flows_out) {
195+
return false;
196+
}
197+
flow_sources_out[*n_flows_out] = static_cast<uint64_t>(idx_a[i]);
198+
flow_targets_out[*n_flows_out] = static_cast<uint64_t>(idx_b[j]);
199+
flow_values_out[*n_flows_out] = flow;
200+
++(*n_flows_out);
201+
}
202+
}
203+
return true;
204+
}
205+
150206
} // namespace
151207

152208

@@ -186,9 +242,9 @@ int EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G,
186242
std::vector<double> weights1(n), weights2(m);
187243
Digraph di(n, m);
188244
const SetupPolicy policy = make_setup_policy(n, m, n1, n2, true);
189-
NetworkSimplexSimple<Digraph,double,double, node_id_type> net(
190-
di, policy.use_arc_mixing, (int) (n + m), n * m, maxIter
191-
);
245+
typedef NetworkSimplexSimple<Digraph, double, double, node_id_type> Simplex;
246+
Simplex::SimplexOptions simplex_options(policy.use_arc_mixing);
247+
Simplex net(di, simplex_options, (int) (n + m), n * m, maxIter);
192248

193249
// Set supply and demand, don't account for 0 values (faster)
194250

@@ -341,6 +397,7 @@ int EMD_wrap_sparse(
341397
uint64_t *flow_targets_out,
342398
double *flow_values_out,
343399
uint64_t *n_flows_out,
400+
uint64_t max_flows_out,
344401
double *alpha,
345402
double *beta,
346403
double *cost,
@@ -432,9 +489,9 @@ int EMD_wrap_sparse(
432489

433490
di.buildFromEdges(edges);
434491

435-
NetworkSimplexSimple<Digraph, double, double, node_id_type> net(
436-
di, true, (int)(n + m), di.arcNum(), maxIter
437-
);
492+
typedef NetworkSimplexSimple<Digraph, double, double, node_id_type> Simplex;
493+
Simplex::SimplexOptions simplex_options(true);
494+
Simplex net(di, simplex_options, (int)(n + m), di.arcNum(), maxIter);
438495

439496
net.supplyMap(&weights1[0], (int)n, &weights2[0], (int)m);
440497

@@ -463,41 +520,27 @@ int EMD_wrap_sparse(
463520
int ret = net.run();
464521
if (ret == (int)net.OPTIMAL || ret == (int)net.MAX_ITER_REACHED) {
465522
*cost = 0;
466-
*n_flows_out = 0;
523+
*n_flows_out = 0;
467524

468-
Arc a;
469-
di.first(a);
470-
for (; a != INVALID; di.next(a)) {
471-
uint64_t i = di.source(a);
472-
uint64_t j = di.target(a);
473-
double flow = net.flow(a);
474-
475-
uint64_t orig_i = indI[i];
476-
uint64_t orig_j = indJ[j - n];
477-
478-
479-
double arc_cost = arc_costs[a];
480-
481-
*cost += flow * arc_cost;
482-
483-
484-
*(alpha + orig_i) = -net.potential(i);
485-
*(beta + orig_j) = net.potential(j);
486-
487-
if (flow > 1e-15) {
488-
flow_sources_out[*n_flows_out] = orig_i;
489-
flow_targets_out[*n_flows_out] = orig_j;
490-
flow_values_out[*n_flows_out] = flow;
491-
(*n_flows_out)++;
492-
}
525+
auto sparse_cost = [&arc_costs](Arc a, int, int) {
526+
return arc_costs[a];
527+
};
528+
if (!extract_sparse_solution(
529+
net, di, INVALID, indI, indJ, alpha, beta, cost,
530+
flow_sources_out, flow_targets_out, flow_values_out,
531+
n_flows_out, max_flows_out, sparse_cost, 1e-15)) {
532+
return (int)net.MAX_ITER_REACHED;
493533
}
494534
}
495535
return ret;
496536
}
497537

498538
int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b,
499-
int dim, int metric, double *G, double *alpha, double *beta,
500-
double *cost, uint64_t maxIter, double *alpha_init, double *beta_init) {
539+
int dim, int metric, uint64_t *flow_sources_out,
540+
uint64_t *flow_targets_out, double *flow_values_out,
541+
uint64_t *n_flows_out, uint64_t max_flows_out,
542+
double *alpha, double *beta, double *cost, uint64_t maxIter,
543+
double *alpha_init, double *beta_init) {
501544
using namespace lemon;
502545
typedef FullBipartiteDigraph Digraph;
503546
DIGRAPH_TYPEDEFS(Digraph);
@@ -552,8 +595,19 @@ int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double
552595
// Create full bipartite graph
553596
Digraph di(n, m);
554597

555-
NetworkSimplexSimple<Digraph, double, double, node_id_type> net(
556-
di, true, (int)(n + m), (uint64_t)(n) * (uint64_t)(m), maxIter
598+
typedef NetworkSimplexSimple<Digraph, double, double, node_id_type> Simplex;
599+
Simplex::SimplexOptions simplex_options(false);
600+
// Lazy mode does not store costs or endpoints for the real complete
601+
// bipartite arcs. Artificial root arcs are still explicit because the
602+
// simplex initialization assigns them costs 0 or ART_COST.
603+
simplex_options.cost_storage_mode = Simplex::CostStorageMode::ArtificialArcCosts;
604+
simplex_options.flow_storage_mode = Simplex::FlowStorageMode::SparseArcFlows;
605+
simplex_options.endpoint_storage_mode =
606+
Simplex::EndpointStorageMode::ArcEndpoints;
607+
simplex_options.state_storage_mode = Simplex::StateStorageMode::PackedArcStates;
608+
609+
Simplex net(
610+
di, simplex_options, (int)(n + m), (uint64_t)(n) * (uint64_t)(m), maxIter
557611
);
558612

559613
// Set supplies
@@ -583,32 +637,20 @@ int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double
583637

584638
if (ret == (int)net.OPTIMAL || ret == (int)net.MAX_ITER_REACHED) {
585639
*cost = 0;
640+
*n_flows_out = 0;
586641

587642
// Initialize output arrays
588-
for (int i = 0; i < n1 * n2; i++) G[i] = 0.0;
589643
for (int i = 0; i < n1; i++) alpha[i] = 0.0;
590644
for (int i = 0; i < n2; i++) beta[i] = 0.0;
591-
592-
// Extract solution
593-
Arc a;
594-
di.first(a);
595-
for (; a != INVALID; di.next(a)) {
596-
int i = di.source(a);
597-
int j = di.target(a) - n;
598-
599-
int orig_i = idx_a[i];
600-
int orig_j = idx_b[j];
601-
602-
double flow = net.flow(a);
603-
G[orig_i * n2 + orig_j] = flow;
604-
605-
alpha[orig_i] = -net.potential(i);
606-
beta[orig_j] = net.potential(j + n);
607-
608-
if (flow > 0) {
609-
double c = net.computeLazyCost(i, j);
610-
*cost += flow * c;
611-
}
645+
646+
auto lazy_cost = [&net](Arc, int i, int j) {
647+
return net.computeLazyCost(i, j);
648+
};
649+
if (!extract_sparse_solution(
650+
net, di, INVALID, idx_a, idx_b, alpha, beta, cost,
651+
flow_sources_out, flow_targets_out, flow_values_out,
652+
n_flows_out, max_flows_out, lazy_cost, 0.0)) {
653+
return (int)net.MAX_ITER_REACHED;
612654
}
613655
}
614656

‎ot/lp/_network_simplex.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1037,7 +1037,7 @@ def emd2_lazy(
10371037
alpha_init_np = np.asarray(alpha_init_np, dtype=np.float64, order="C")
10381038
beta_init_np = np.asarray(beta_init_np, dtype=np.float64, order="C")
10391039

1040-
G, cost, u, v, result_code = emd_c_lazy(
1040+
flow_sources, flow_targets, flow_values, cost, u, v, result_code = emd_c_lazy(
10411041
a_np, b_np, X_a_np, X_b_np, metric, numItermax, alpha_init_np, beta_init_np
10421042
)
10431043

@@ -1053,8 +1053,6 @@ def emd2_lazy(
10531053
stacklevel=2,
10541054
)
10551055

1056-
G_backend = nx.from_numpy(G, type_as=type_as)
1057-
10581056
cost_backend = nx.set_gradients(
10591057
nx.from_numpy(cost, type_as=type_as),
10601058
(a0, b0),
@@ -1075,6 +1073,16 @@ def emd2_lazy(
10751073
"result_code": result_code,
10761074
}
10771075
if return_matrix:
1076+
flow_values_backend = nx.from_numpy(flow_values, type_as=type_as)
1077+
flow_sources_backend = nx.from_numpy(flow_sources.astype(np.int64))
1078+
flow_targets_backend = nx.from_numpy(flow_targets.astype(np.int64))
1079+
G_backend = nx.coo_matrix(
1080+
flow_values_backend,
1081+
flow_sources_backend,
1082+
flow_targets_backend,
1083+
shape=(n1, n2),
1084+
type_as=type_as,
1085+
)
10781086
log_dict["G"] = G_backend
10791087
return cost_backend, log_dict
10801088
else:

‎ot/lp/emd_wrap.pyx‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ import warnings
2222
cdef extern from "EMD.h":
2323
int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil
2424
int EMD_wrap_omp(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter, int numThreads) nogil
25-
int EMD_wrap_sparse(int n1, int n2, double *X, double *Y, uint64_t n_edges, uint64_t *edge_sources, uint64_t *edge_targets, double *edge_costs, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, double *alpha, double *beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil
26-
int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, int dim, int metric, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil
25+
int EMD_wrap_sparse(int n1, int n2, double *X, double *Y, uint64_t n_edges, uint64_t *edge_sources, uint64_t *edge_targets, double *edge_costs, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double *alpha, double *beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil
26+
int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, int dim, int metric, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil
2727
cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED
2828

2929

@@ -306,7 +306,7 @@ def emd_c_sparse(np.ndarray[double, ndim=1, mode="c"] a,
306306
n_edges,
307307
<uint64_t*> edge_sources.data, <uint64_t*> edge_targets.data, <double*> edge_costs.data,
308308
<uint64_t*> flow_sources.data, <uint64_t*> flow_targets.data, <double*> flow_values.data,
309-
&n_flows_out,
309+
&n_flows_out, n_edges,
310310
<double*> alpha.data, <double*> beta.data, &cost, max_iter,
311311
alpha_init_ptr, beta_init_ptr
312312
)
@@ -329,6 +329,8 @@ def emd_c_lazy(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1
329329
cdef int result_code = 0
330330
cdef double cost = 0
331331
cdef int metric_code
332+
cdef uint64_t n_flows_out = 0
333+
cdef uint64_t max_flows_out = n1 + n2
332334

333335
# Validate dimension consistency
334336
if coords_b.shape[1] != dim:
@@ -345,9 +347,11 @@ def emd_c_lazy(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1
345347
except KeyError:
346348
raise ValueError(f"Unknown metric: '{metric}'. Supported metrics are: {list(metric_map.keys())}")
347349

350+
cdef np.ndarray[uint64_t, ndim=1, mode="c"] flow_sources = np.zeros(max_flows_out, dtype=np.uint64)
351+
cdef np.ndarray[uint64_t, ndim=1, mode="c"] flow_targets = np.zeros(max_flows_out, dtype=np.uint64)
352+
cdef np.ndarray[double, ndim=1, mode="c"] flow_values = np.zeros(max_flows_out, dtype=np.float64)
348353
cdef np.ndarray[double, ndim=1, mode="c"] alpha = np.zeros(n1)
349354
cdef np.ndarray[double, ndim=1, mode="c"] beta = np.zeros(n2)
350-
cdef np.ndarray[double, ndim=2, mode="c"] G = np.zeros([n1, n2])
351355
if not len(a):
352356
a = np.ones((n1,)) / n1
353357
if not len(b):
@@ -360,5 +364,10 @@ def emd_c_lazy(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1
360364
beta_init_ptr = <double*> beta_init.data
361365

362366
with nogil:
363-
result_code = EMD_wrap_lazy(n1, n2, <double*> a.data, <double*> b.data, <double*> coords_a.data, <double*> coords_b.data, dim, metric_code, <double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, max_iter, alpha_init_ptr, beta_init_ptr)
364-
return G, cost, alpha, beta, result_code
367+
result_code = EMD_wrap_lazy(n1, n2, <double*> a.data, <double*> b.data, <double*> coords_a.data, <double*> coords_b.data, dim, metric_code, <uint64_t*> flow_sources.data, <uint64_t*> flow_targets.data, <double*> flow_values.data, &n_flows_out, max_flows_out, <double*> alpha.data, <double*> beta.data, <double*> &cost, max_iter, alpha_init_ptr, beta_init_ptr)
368+
369+
flow_sources = flow_sources[:n_flows_out]
370+
flow_targets = flow_targets[:n_flows_out]
371+
flow_values = flow_values[:n_flows_out]
372+
373+
return flow_sources, flow_targets, flow_values, cost, alpha, beta, result_code

0 commit comments

Comments
 (0)