Add checkpoint/resume functionality to EMD solver

Implements pause and resume capabilities for the EMD (Earth Mover's Distance)
solver, allowing long-running optimizations to be interrupted and continued
from their exact state.

Changes:
- Modified EMD_wrap() signature to accept checkpoint parameters for saving
  and restoring complete internal solver state (flow, potentials, tree structure)
- Added saveCheckpoint(), restoreCheckpoint(), and runFromCheckpoint() methods
  to NetworkSimplexSimple class in network_simplex_simple.h
- Extended emd_c() Cython wrapper to handle checkpoint dictionaries with 12 fields
  (10 arrays + 2 scalar arc counts)
- Added 'checkpoint' and 'return_checkpoint' parameters to emd() Python API
- Includes search_arc_num and all_arc_num scalars in checkpoint to preserve
  initialization state required by start() method

Tests:
- Added test_emd_checkpoint() for basic save/resume functionality
- Added test_emd_checkpoint_multiple() for multiple pause/resume cycles
- Added test_emd_checkpoint_structure() to verify checkpoint field integrity

Author: nathanneike

ot/lp/EMD.h

@@ -29,7 +29,16 @@ enum ProblemType {
     MAX_ITER_REACHED
 };
 
-int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter);
+int EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G, 
+             double* alpha, double* beta, double *cost, uint64_t maxIter,
+             int resume_mode=0, int return_checkpoint=0,
+             double* flow_state=nullptr, double* pi_state=nullptr, 
+             signed char* state_state=nullptr, int* parent_state=nullptr, 
+             int64_t* pred_state=nullptr, int* thread_state=nullptr, 
+             int* rev_thread_state=nullptr, int* succ_num_state=nullptr, 
+             int* last_succ_state=nullptr, signed char* forward_state=nullptr,
+             int64_t* search_arc_num_out=nullptr, int64_t* all_arc_num_out=nullptr);
+
 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);
 
 

ot/lp/EMD_wrapper.cpp

@@ -20,7 +20,14 @@
 
 
 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, double* beta, double *cost, uint64_t maxIter,
+                int resume_mode, int return_checkpoint,
+                double* flow_state, double* pi_state, signed char* state_state,
+                int* parent_state, int64_t* pred_state,
+                int* thread_state, int* rev_thread_state,
+                int* succ_num_state, int* last_succ_state,
+                signed char* forward_state,
+                int64_t* search_arc_num_out, int64_t* all_arc_num_out)  {
     // beware M and C are stored in row major C style!!!
 
     using namespace lemon;
@@ -93,8 +100,29 @@ int EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G,
 
 
     // Solve the problem with the network simplex algorithm
-
-    int ret=net.run();
+    // If resume_mode=1 and checkpoint data provided, resume from checkpoint
+    // Otherwise do normal run
+    
+    int64_t search_arc_num_in = 0, all_arc_num_in = 0;
+    if (resume_mode == 1 && search_arc_num_out != nullptr && all_arc_num_out != nullptr) {
+        search_arc_num_in = *search_arc_num_out;
+        all_arc_num_in = *all_arc_num_out;
+    }
+    
+    int ret;
+    if (resume_mode == 1 && flow_state != nullptr) {
+        // Resume from checkpoint
+        ret = net.runFromCheckpoint(
+            flow_state, pi_state, state_state,
+            parent_state, pred_state,
+            thread_state, rev_thread_state,
+            succ_num_state, last_succ_state, forward_state,
+            search_arc_num_in, all_arc_num_in);
+    } else {
+        // Normal run
+        ret = net.run();
+    }
+    
     uint64_t i, j;
     if (ret==(int)net.OPTIMAL || ret==(int)net.MAX_ITER_REACHED) {
         *cost = 0;
@@ -111,16 +139,22 @@ int EMD_wrap(int n1, int n2, double *X, double *Y, double *D, double *G,
 
     }
 
+    // Save checkpoint if requested and arrays provided
+    if (return_checkpoint == 1 && flow_state != nullptr) {
+        net.saveCheckpoint(
+            flow_state, pi_state, state_state,
+            parent_state, pred_state,
+            thread_state, rev_thread_state,
+            succ_num_state, last_succ_state, forward_state,
+            search_arc_num_out, all_arc_num_out);
+    }
 
     return ret;
 }
 
 
 
 
-
-
-
 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)  {
     // beware M and C are stored in row major C style!!!

ot/lp/_network_simplex.py

@@ -172,6 +172,8 @@ def emd(
     center_dual=True,
     numThreads=1,
     check_marginals=True,
+    checkpoint=None,
+    return_checkpoint=False,
 ):
     r"""Solves the Earth Movers distance problem and returns the OT matrix
 
@@ -232,6 +234,15 @@ def emd(
     check_marginals: bool, optional (default=True)
         If True, checks that the marginals mass are equal. If False, skips the
         check.
+    checkpoint: dict, optional (default=None)
+        Checkpoint data from a previous emd() call to resume computation.
+        The checkpoint must contain internal solver state including flow,
+        potentials, and tree structure. Obtain by calling emd() with
+        return_checkpoint=True.
+    return_checkpoint: bool, optional (default=False)
+        If True and log=True, includes complete internal solver state in the
+        returned log dictionary for checkpointing. This enables pausing and
+        resuming the optimization.
 
 
     Returns
@@ -241,7 +252,8 @@ def emd(
         parameters
     log: dict, optional
         If input log is true, a dictionary containing the
-        cost and dual variables and exit status
+        cost and dual variables and exit status. If return_checkpoint=True,
+        also contains internal solver state for resuming computation.
 
 
     Examples
@@ -321,7 +333,43 @@ def emd(
 
     numThreads = check_number_threads(numThreads)
 
-    G, cost, u, v, result_code = emd_c(a, b, M, numItermax, numThreads)
+    checkpoint_data = None
+    if checkpoint is not None:
+        # Extract checkpoint arrays and convert to numpy (strip leading underscore)
+        checkpoint_data = {
+            "flow": nx.to_numpy(checkpoint["_flow"]) if "_flow" in checkpoint else None,
+            "pi": nx.to_numpy(checkpoint["_pi"]) if "_pi" in checkpoint else None,
+            "state": nx.to_numpy(checkpoint["_state"])
+            if "_state" in checkpoint
+            else None,
+            "parent": nx.to_numpy(checkpoint["_parent"])
+            if "_parent" in checkpoint
+            else None,
+            "pred": nx.to_numpy(checkpoint["_pred"]) if "_pred" in checkpoint else None,
+            "thread": nx.to_numpy(checkpoint["_thread"])
+            if "_thread" in checkpoint
+            else None,
+            "rev_thread": nx.to_numpy(checkpoint["_rev_thread"])
+            if "_rev_thread" in checkpoint
+            else None,
+            "succ_num": nx.to_numpy(checkpoint["_succ_num"])
+            if "_succ_num" in checkpoint
+            else None,
+            "last_succ": nx.to_numpy(checkpoint["_last_succ"])
+            if "_last_succ" in checkpoint
+            else None,
+            "forward": nx.to_numpy(checkpoint["_forward"])
+            if "_forward" in checkpoint
+            else None,
+            "search_arc_num": int(checkpoint.get("search_arc_num", 0)),
+            "all_arc_num": int(checkpoint.get("all_arc_num", 0)),
+        }
+        # Filter out None values
+        checkpoint_data = {k: v for k, v in checkpoint_data.items() if v is not None}
+
+    G, cost, u, v, result_code, checkpoint_out = emd_c(
+        a, b, M, numItermax, numThreads, checkpoint_data, int(return_checkpoint)
+    )
 
     if center_dual:
         u, v = center_ot_dual(u, v, a, b)
@@ -345,6 +393,22 @@ def emd(
         log["v"] = nx.from_numpy(v, type_as=type_as)
         log["warning"] = result_code_string
         log["result_code"] = result_code
+
+        # Add checkpoint data if requested (preserve original dtypes, don't cast)
+        if return_checkpoint and checkpoint_out is not None:
+            log["_flow"] = checkpoint_out["flow"]
+            log["_pi"] = checkpoint_out["pi"]
+            log["_state"] = checkpoint_out["state"]
+            log["_parent"] = checkpoint_out["parent"]
+            log["_pred"] = checkpoint_out["pred"]
+            log["_thread"] = checkpoint_out["thread"]
+            log["_rev_thread"] = checkpoint_out["rev_thread"]
+            log["_succ_num"] = checkpoint_out["succ_num"]
+            log["_last_succ"] = checkpoint_out["last_succ"]
+            log["_forward"] = checkpoint_out["forward"]
+            log["search_arc_num"] = int(checkpoint_out["search_arc_num"])
+            log["all_arc_num"] = int(checkpoint_out["all_arc_num"])
+
         return nx.from_numpy(G, type_as=type_as), log
     return nx.from_numpy(G, type_as=type_as)