Merge pull request #68 from francois-drielsma/develop

Massive edge selector speed-up + vertexer bug fix

Author: francois-drielsma

bin/larcv_count_entries.py

@@ -60,7 +60,7 @@ def main(source, source_list, tree_name):
                        help='Path to a text file of data file paths',
                        type=str)
 
-    parser.add_argument('--tree_name',
+    parser.add_argument('--tree-name',
                         help='TTree name used to count the entries.',
                         type=str)
 

bin/larcv_find_duplicates.py

@@ -100,7 +100,7 @@ def main(source, source_list, output, tree_name):
                         help='Path to the output text file with the duplicate list',
                         type=str, required=True)
 
-    parser.add_argument('--tree_name',
+    parser.add_argument('--tree-name',
                         help='TTree name used to count the entries.',
                         type=str)
 

bin/larcv_find_run.py

@@ -83,7 +83,7 @@ def main(source, source_list, output, run_number, tree_name):
                         help='Run number to look for',
                         type=int, required=True)
 
-    parser.add_argument('--tree_name',
+    parser.add_argument('--tree-name',
                         help='TTree name used to count the entries.',
                         type=str)
 

spine/post/reco/vertex.py

@@ -115,6 +115,8 @@ def reconstruct_vertex_single(self, inter):
                     elif (np.linalg.norm(part.start_point - inter.vertex) 
                           < self.touching_threshold):
                         part.is_primary = True
-                    elif (part.shape == SHOWR_SHP and
-                          np.dot(part.start_point, inter.vertex) < self.angle_threshold):
-                        part.is_primary = True
+                    elif part.shape == SHOWR_SHP:
+                        vec = part.start_point - inter.vertex
+                        cos = np.dot(vec/np.linalg.norm(vec), part.start_dir)
+                        if cos < self.angle_threshold:
+                            part.is_primary = True

spine/utils/gnn/evaluation.py

@@ -642,12 +642,15 @@ def edge_assignment_score(edge_index: nb.int64[:,:],
     best_groups = np.arange(num_nodes, dtype=np.int64)
     track_used  = np.zeros(num_nodes, dtype=np.bool_)
     best_loss   = grouping_loss(pred_adj.flatten(), empty_adj.flatten())
+    group_set   = {np.int64(-1)} # Cannot be empty (typing)
     for k, e in enumerate(ord_index):
         # If the edge connect two nodes already in the same group, proceed
         a, b = e
         group_a, group_b = best_groups[a], best_groups[b]
         if group_a == group_b:
             continue
+        if group_a in group_set and group_b in group_set:
+            continue
 
         # If requested, check whether there is already a track connection
         if track_node is not None:
@@ -677,6 +680,9 @@ def edge_assignment_score(edge_index: nb.int64[:,:],
             best_groups[best_groups == group_b] = group_a
             best_loss += combined_loss - current_loss
             best_ids = np.append(best_ids, k)
+            group_set -= {group_a, group_b}
+        else:
+            group_set.update({group_a, group_b})
 
     # Build the edge index
     best_index = ord_index[best_ids]