parallel execution

tools/annotate_roadmap.py

@@ -1,9 +1,13 @@
 #!/usr/bin/env python3
 import argparse
+from multiprocessing import Pool
+from typing import List, Tuple
 
 import numpy as np
 import yaml
 
+N_PROCESSES = 8
+
 if __name__ != "__main__":
     from . import collision
 
@@ -54,27 +58,41 @@ def add_self_edges(roadmap):
     return roadmap
 
 
+def check_proxy(args):
+    """Prox method to call the collision checker with the right arguments."""
+    _, _, E, p0, p1, q0, q1 = args
+    return collision.ellipsoid_collision_motion(E, p0, p1, q0, q1)
+
+
 def compute_edge_conflicts(radius, map):
     # compute the pairwise collisions and add them to the map
     E = np.diag([radius, radius])
     num_edges = len(map["roadmap"]["edges"])
     v_dict = map["roadmap"]["vertices"]
     edges = map["roadmap"]["edges"]
     conflicts = [[] for _ in range(num_edges)]
+    edges_to_check: List[Tuple[
+        int, int, np.ndarray,  # i, j, E
+        np.ndarray, np.ndarray, np.ndarray, np.ndarray  # p0, p1, q0, q1
+    ]] = []
     for i in range(0, num_edges):
         p0 = np.asarray(v_dict[edges[i][0]])
         p1 = np.asarray(v_dict[edges[i][1]])
         for j in range(i+1, num_edges):
             q0 = np.asarray(v_dict[edges[j][0]])
             q1 = np.asarray(v_dict[edges[j][1]])
             if collision.precheck_bounding_box(E, p0, p1, q0, q1):
-                collides = collision.ellipsoid_collision_motion(
-                    E, p0, p1, q0, q1)
-            else:
-                collides = False
-            if collides:
-                conflicts[i].append(j)
-                conflicts[j].append(i)
+                edges_to_check.append((i, j, E, p0, p1, q0, q1))
+
+    # check all edges in parallel
+
+    with Pool(N_PROCESSES) as p:
+        results = p.map(check_proxy, edges_to_check)
+    for result, (i, j, _, _, _, _, _) in zip(results, edges_to_check):
+        if result:
+            conflicts[i].append(j)
+            conflicts[j].append(i)
+
     return conflicts