diff --git a/compute_shortest_path_naive_graph.py b/compute_shortest_path_naive_graph.py
index acc08b1..281a0dd 100644
--- a/compute_shortest_path_naive_graph.py
+++ b/compute_shortest_path_naive_graph.py
@@ -97,7 +97,7 @@ def main():
 
 
     # TODO replace the big weight hack with 2 ROI nodes, a source and a target with unidirectional free edges
-    merge_w = 10000 # remove twice from computed path weight
+    merge_w = 0 # remove twice from computed path weight
 
 
     if args.target == 'COM':
@@ -108,6 +108,7 @@ def main():
         g.add_vertices(['COM_{}_source'.format(i) for i in range(1, label_map.max()+1)])
         g.add_vertices(['COM_{}_target'.format(i) for i in range(1, label_map.max()+1)])
 
+        source_vertex = []
         target_vertex = []
         edges_to_add = []
         for i in range(1, label_map.max()+1):
@@ -120,7 +121,7 @@ def main():
             # add new vertex to converge them all there
             new_vert_id_source = g.vs['name'].index('COM_{}_source'.format(i))
             new_vert_id_target = g.vs['name'].index('COM_{}_target'.format(i))
-            target_vertex.append(new_vert_id_source)
+            source_vertex.append(new_vert_id_source)
             target_vertex.append(new_vert_id_target)
             # create IN and OUT edge for all node at COM
             edges_to_add += [(new_vert_id_source, i_vert) for i_vert in COM_vertex]
@@ -132,7 +133,7 @@ def main():
         g.add_edges(edges_to_add, 
                     {'neg_log':[0]*len(edges_to_add)})
 
-        print('Temporarily added {:} nodes to graph'.format(len(target_vertex)))
+        print('Temporarily added {:} nodes to graph'.format(len(source_vertex)+len(target_vertex)))
         print('Temporarily added {:} edges to graph'.format(len(edges_to_add)))
 
 
@@ -145,27 +146,31 @@ def main():
         start_time = time()
         # compute center-of-mass -ish voxel for each roi
         # add nodes
-        g.add_vertices(['ROI_{}'.format(i) for i in range(1, label_map.max()+1)])
+        g.add_vertices(['ROI_{}_source'.format(i) for i in range(1, label_map.max()+1)])
+        g.add_vertices(['ROI_{}_target'.format(i) for i in range(1, label_map.max()+1)])
 
+        source_vertex = []
         target_vertex = []
         edges_to_add = []
         for i in range(1, label_map.max()+1):
             ROI_vertex = rois_vertex[i-1]
             # add new vertex to converge them all there
-            new_vert_id = g.vs['name'].index('ROI_{}'.format(i))
-            target_vertex.append(new_vert_id)
+            new_vert_id_source = g.vs['name'].index('ROI_{}_source'.format(i))
+            new_vert_id_target = g.vs['name'].index('ROI_{}_target'.format(i))
+            source_vertex.append(new_vert_id_source)
+            target_vertex.append(new_vert_id_target)
             # create IN and OUT edge for all node at COM
-            edges_to_add += [(new_vert_id, i_vert) for i_vert in ROI_vertex]
-            edges_to_add += [(i_vert, new_vert_id) for i_vert in ROI_vertex]
+            edges_to_add += [(new_vert_id_source, i_vert) for i_vert in ROI_vertex]
+            edges_to_add += [(i_vert, new_vert_id_target) for i_vert in ROI_vertex]
 
         # edge of zero could give loops
         # instead we put very very expensive nodes, and we can remove it when counting
         g.add_edges(edges_to_add, 
-                    {'neg_log':[merge_w]*len(edges_to_add)})
+                    {'neg_log':[0]*len(edges_to_add)})
         end_time = time()
 
         print('Making ROI-nodes = {:.2f} s'.format(end_time - start_time))
-        print('Temporarily added {:} nodes to graph'.format(len(target_vertex)))
+        print('Temporarily added {:} nodes to graph'.format(len(source_vertex)+len(target_vertex)))
         print('Temporarily added {:} edges to graph'.format(len(edges_to_add)))
 
 
@@ -173,6 +178,7 @@ def main():
     print('Computing shortest paths')
     start_time = time()
     paths, paths_length, weights = compute_shortest_paths_COM2COM(g,
+                                                                  source_vertex,
                                                                   target_vertex,
                                                                   w='neg_log')
 
diff --git a/utils.py b/utils.py
index 3cc4a2c..6bb2b31 100644
--- a/utils.py
+++ b/utils.py
@@ -129,22 +129,22 @@ def mask_COM(mask):
 
 
 
-def compute_shortest_paths_COM2COM(g, COMs, w='neg_log'):
+def compute_shortest_paths_COM2COM(g, source_COMs, target_COMs, w='neg_log'):
     # compute shortest paths between all pairs of center-of-mass voxel
 
     # compute graph weight of shortest path
-    path_weights = g.shortest_paths(source=COMs, 
-                                    target=COMs, 
+    path_weights = g.shortest_paths(source=source_COMs, 
+                                    target=target_COMs, 
                                     weights=w, 
                                     mode='out')
 
     # compute the graph shortest path
     paths = [] # paths[i_source][i_dest]
-    for i_source in range(len(COMs)):
-        source = COMs[i_source]
+    for i_source in range(len(source_COMs)):
+        source = source_COMs[i_source]
 
         path = g.get_shortest_paths(source, 
-                                    to=COMs, 
+                                    to=target_COMs, 
                                     weights=w, 
                                     mode='out', 
                                     output='vpath')
@@ -152,7 +152,7 @@ def compute_shortest_paths_COM2COM(g, COMs, w='neg_log'):
 
     # compute the lenght in term of vertex of the shortest path
     paths_length = []
-    for i_source in range(len(COMs)):
+    for i_source in range(len(source_COMs)):
         path = paths[i_source]
         paths_length.append([len(s) for s in path])