Various speed-ups in the full chain routines

README.md

@@ -5,7 +5,7 @@
 [![Build Status](https://app.travis-ci.com/francois-drielsma/lartpc_mlreco3d.svg?token=WB4oxAv87vEXhuxUGH7e&branch=develop&status=passed)](https://app.travis-ci.com/github/francois-drielsma/lartpc_mlreco3d/logscans?serverType=git)
 [![Documentation Status](https://readthedocs.org/projects/lartpc-mlreco3d/badge/?version=latest)](https://lartpc-mlreco3d.readthedocs.io/en/latest/?badge=latest)
 
-The Scalable Particle Imaging with Neural Embeddings (SPINE) package leverages state-of-the-art Machine Learning (ML) algorithms -- in particular Deep Neural Networks (DNNs) -- to reconstruct particle imagaging detector data. This package was primarily developed for Liquid Argon Time-Projection Chamber (LArTPC) data and relies on Convolutional Neural Networks (CNNs) for pixel-level feature extraction and Graph Neural Networks (GNNs) for superstructure formation. The schematic below breaks down the full end-to-end reconstruction flow.
+The Scalable Particle Imaging with Neural Embeddings (SPINE) package leverages state-of-the-art Machine Learning (ML) algorithms -- in particular Deep Neural Networks (DNNs) -- to reconstruct particle imaging detector data. This package was primarily developed for Liquid Argon Time-Projection Chamber (LArTPC) data and relies on Convolutional Neural Networks (CNNs) for pixel-level feature extraction and Graph Neural Networks (GNNs) for superstructure formation. The schematic below breaks down the full end-to-end reconstruction flow.
 
 ![Full chain](https://github.com/DeepLearnPhysics/spine/blob/develop/docs/source/_static/img/spine-chain-alpha.png)
 

spine/ana/diag/track.py

@@ -1,13 +1,12 @@
 """Module to evaluate diagnostic metrics on tracks."""
 
 import numpy as np
-from scipy.spatial.distance import cdist
 
-from spine.ana.base import AnaBase
+from spine.math.distance import cdist
 
 from spine.utils.globals import TRACK_SHP
-from spine.utils.numba_local import principal_components
 
+from spine.ana.base import AnaBase
 
 __all__ = ['TrackCompletenessAna']
 
@@ -142,10 +141,11 @@ def cluster_track_chunks(points, start_point, end_point, pixel_size):
         """
         # Project and cluster on the projected axis
         direction = (end_point-start_point)/np.linalg.norm(end_point-start_point)
+        scale = pixel_size*np.max(direction)
         projs = np.dot(points - start_point, direction)
         perm = np.argsort(projs)
         seps = projs[perm][1:] - projs[perm][:-1]
-        breaks = np.where(seps > pixel_size*1.1)[0] + 1
+        breaks = np.where(seps > scale*1.1)[0] + 1
         cluster_labels = np.empty(len(projs), dtype=int)
         for i, index in enumerate(np.split(np.arange(len(projs)), breaks)):
             cluster_labels[perm[index]] = i

spine/ana/metric/segment.py

@@ -111,7 +111,7 @@ def process(self, data):
             if self.ghost:
                 # If there are ghost, must combine the predictions
                 full_seg_pred = np.full_like(seg_label, GHOST_SHP, dtype=np.int32)
-                deghost_mask = data['ghost'][:, 0] > data['ghost'][:, 1]
+                deghost_mask = np.argmax(data['ghost'], axis=1) == 0
                 full_seg_pred[deghost_mask] = seg_pred
                 seg_pred = full_seg_pred
 

spine/build/base.py

@@ -38,13 +38,13 @@ class BuilderBase(ABC):
 
     # Necessary/optional data products to load a reconstructed object
     _load_reco_keys   = (
-            ('points', True), ('depositions', True), ('sources', False)
+            ('points', False), ('depositions', False), ('sources', False)
     )
 
     # Necessary/optional data products to load a truth object
     _load_truth_keys  = (
-            ('points_label', True), ('points', False), ('points_g4', False),
-            ('depositions_label', True), ('depositions', False),
+            ('points_label', False), ('points', False), ('points_g4', False),
+            ('depositions_label', False), ('depositions', False),
             ('depositions_q_label', False), ('depositions_g4', False),
             ('sources_label', False), ('sources', False)
     )

spine/build/fragment.py

@@ -152,7 +152,7 @@ def build_truth(self, data):
         """
         return self._build_truth(**data)
 
-    def _build_truth(self, label_tensor, points_label, depositions_label, 
+    def _build_truth(self, label_tensor, points_label, depositions_label,
                      depositions_q_label=None, label_adapt_tensor=None,
                      points=None, depositions=None, label_g4_tensor=None,
                      points_g4=None, depositions_g4=None, sources_label=None,
@@ -287,16 +287,17 @@ def load_reco(self, data):
         """
         return self._load_reco(**data)
 
-    def _load_reco(self, reco_fragments, points, depositions, sources=None):
+    def _load_reco(self, reco_fragments, points=None, depositions=None,
+                   sources=None):
         """Load :class:`RecoFragment` objects from their stored versions.
 
         Parameters
         ----------
         reco_fragments : List[RecoFragment]
             (F) List of partial reconstructed fragments
-        points : np.ndarray
+        points : np.ndarray, optional
             (N, 3) Set of deposition coordinates in the image
-        depositions : np.ndarray
+        depositions : np.ndarray, optional
             (N) Set of deposition values
         sources : np.ndarray, optional
             (N, 2) Tensor which contains the module/tpc information
@@ -313,10 +314,11 @@ def _load_reco(self, reco_fragments, points, depositions, sources=None):
                     "The ordering of the stored fragments is wrong.")
 
             # Update the fragment with its long-form attributes
-            fragment.points = points[fragment.index]
-            fragment.depositions = depositions[fragment.index]
-            if sources is not None:
-                fragment.sources = sources[fragment.index]
+            if points is not None:
+                fragment.points = points[fragment.index]
+                fragment.depositions = depositions[fragment.index]
+                if sources is not None:
+                    fragment.sources = sources[fragment.index]
 
         return reco_fragments
 
@@ -335,20 +337,20 @@ def load_truth(self, data):
         """
         return self._load_truth(**data)
 
-    def _load_truth(self, truth_fragments, points_label, depositions_label,
-                    depositions_q_label=None, points=None, depositions=None,
-                    points_g4=None, depositions_g4=None, sources_label=None,
-                    sources=None):
+    def _load_truth(self, truth_fragments, points_label=None,
+                    depositions_label=None, depositions_q_label=None,
+                    points=None, depositions=None, points_g4=None,
+                    depositions_g4=None, sources_label=None, sources=None):
         """Load :class:`TruthFragment` objects from their stored versions.
 
         Parameters
         ----------
         truth_fragments : List[TruthFragment]
             (F) List of partial truth fragments
-        points_label : np.ndarray
+        points_label : np.ndarray, optional
             (N', 3) Set of deposition coordinates in the label image (identical
             for pixel TPCs, different if deghosting is involved)
-        depositions_label : np.ndarray
+        depositions_label : np.ndarray, optional
             (N') Set of true deposition values in MeV
         depositions_q_label : np.ndarray, optional
             (N') Set of true deposition values in ADC, if relevant
@@ -377,12 +379,13 @@ def _load_truth(self, truth_fragments, points_label, depositions_label,
                     "The ordering of the stored fragments is wrong.")
 
             # Update the fragment with its long-form attributes
-            fragment.points = points_label[fragment.index]
-            fragment.depositions = depositions_label[fragment.index]
-            if depositions_q_label is not None:
-                fragment.depositions_q = depositions_q_label[fragment.index]
-            if sources_label is not None:
-                fragment.sources = sources_label[fragment.index]
+            if points_label is not None:
+                fragment.points = points_label[fragment.index]
+                fragment.depositions = depositions_label[fragment.index]
+                if depositions_q_label is not None:
+                    fragment.depositions_q = depositions_q_label[fragment.index]
+                if sources_label is not None:
+                    fragment.sources = sources_label[fragment.index]
 
             if points is not None:
                 fragment.points_adapt = points[fragment.index_adapt]

spine/build/manager.py

@@ -44,7 +44,7 @@ class BuildManager:
     )
 
     def __init__(self, fragments, particles, interactions,
-                 mode='both', units='cm', sources=None):
+                 mode='both', units='cm', sources=None, lite=False):
         """Initializes the build manager.
 
         Parameters
@@ -60,6 +60,9 @@ def __init__(self, fragments, particles, interactions,
         sources : Dict[str, str], optional
             Dictionary which maps the necessary data products onto a name
             in the input/output dictionary of the reconstruction chain.
+        lite : bool, default False
+            If `True`, the objects being loaded are lite and do not map
+            to long-form attributes. Simply load the matches.
         """
         # Check on the mode, store it
         assert mode in self._run_modes, (
@@ -100,6 +103,9 @@ def __init__(self, fragments, particles, interactions,
         assert len(self.builders), (
                 "Do not call the builder unless it does anything.")
 
+        # Store whether to load the long-form attributes or not
+        self.lite = lite
+
     def __call__(self, data):
         """Build the representations for one entry.
 
@@ -115,7 +121,7 @@ def __call__(self, data):
         # If this is the first time the builders are called, build
         # the objects shared between fragments/particles/interactions
         load = True
-        if 'points' not in data and 'points_label' not in data:
+        if not self.lite and 'points' not in data and 'points_label' not in data:
             load = False
             if np.isscalar(data['index']):
                 sources = self.build_sources(data)

spine/build/particle.py

@@ -311,16 +311,17 @@ def load_reco(self, data):
         """
         return self._load_reco(**data)
 
-    def _load_reco(self, reco_particles, points, depositions, sources=None):
+    def _load_reco(self, reco_particles, points=None, depositions=None,
+                   sources=None):
         """Construct :class:`RecoParticle` objects from their stored versions.
 
         Parameters
         ----------
         reco_particles : List[RecoParticle]
             (P) List of partial reconstructed particles
-        points : np.ndarray
+        points : np.ndarray, optional
             (N, 3) Set of deposition coordinates in the image
-        depositions : np.ndarray
+        depositions : np.ndarray, optional
             (N) Set of deposition values
         sources : np.ndarray, optional
             (N, 2) Tensor which contains the module/tpc information
@@ -337,10 +338,11 @@ def _load_reco(self, reco_particles, points, depositions, sources=None):
                     "The ordering of the stored particles is wrong.")
 
             # Update the particle with its long-form attributes
-            particle.points = points[particle.index]
-            particle.depositions = depositions[particle.index]
-            if sources is not None:
-                particle.sources = sources[particle.index]
+            if points is not None:
+                particle.points = points[particle.index]
+                particle.depositions = depositions[particle.index]
+                if sources is not None:
+                    particle.sources = sources[particle.index]
 
         return reco_particles
 
@@ -354,20 +356,20 @@ def load_truth(self, data):
         """
         return self._load_truth(**data)
 
-    def _load_truth(self, truth_particles, points_label, depositions_label,
-                    depositions_q_label=None, points=None, depositions=None,
-                    points_g4=None, depositions_g4=None, sources_label=None,
-                    sources=None):
+    def _load_truth(self, truth_particles, points_label=None,
+                    depositions_label=None, depositions_q_label=None,
+                    points=None, depositions=None, points_g4=None,
+                    depositions_g4=None, sources_label=None, sources=None):
         """Construct :class:`TruthParticle` objects from their stored versions.
 
         Parameters
         ----------
         truth_particles : List[TruthParticle]
             (P) List of partial truth particles
-        points_label : np.ndarray
+        points_label : np.ndarray, optional
             (N', 3) Set of deposition coordinates in the label image (identical
             for pixel TPCs, different if deghosting is involved)
-        depositions_label : np.ndarray
+        depositions_label : np.ndarray, optional
             (N') Set of true deposition values in MeV
         depositions_q_label : np.ndarray, optional
             (N') Set of true deposition values in ADC, if relevant
@@ -396,12 +398,13 @@ def _load_truth(self, truth_particles, points_label, depositions_label,
                     "The ordering of the stored particles is wrong.")
 
             # Update the particle with its long-form attributes
-            particle.points = points_label[particle.index]
-            particle.depositions = depositions_label[particle.index]
-            if depositions_q_label is not None:
-                particle.depositions_q = depositions_q_label[particle.index]
-            if sources_label is not None:
-                particle.sources = sources_label[particle.index]
+            if points_label is not None:
+                particle.points = points_label[particle.index]
+                particle.depositions = depositions_label[particle.index]
+                if depositions_q_label is not None:
+                    particle.depositions_q = depositions_q_label[particle.index]
+                if sources_label is not None:
+                    particle.sources = sources_label[particle.index]
 
             if points is not None:
                 particle.points_adapt = points[particle.index_adapt]

spine/driver.py

@@ -23,9 +23,10 @@
 from .io import loader_factory, reader_factory, writer_factory
 from .io.write import CSVWriter
 
+from .math import seed as numba_seed
+
 from .utils.logger import logger
 from .utils.cuda import set_visible_devices
-from .utils.numba_local import seed as numba_seed
 from .utils.unwrap import Unwrapper
 from .utils.stopwatch import StopwatchManager
 
@@ -367,7 +368,7 @@ def initialize_io(self, loader=None, reader=None, writer=None):
             # Fetch the list of previously run post-processors
             # TODO: this only works with two runs in a row, not 3 and above
             self.post_list = None
-            if self.reader.cfg is not None:
+            if self.reader.cfg is not None and 'post' in self.reader.cfg:
                 self.post_list = tuple(self.reader.cfg['post'])
 
         # Fetch an appropriate common prefix for all input files

spine/io/parse/cluster.py

@@ -12,13 +12,14 @@
 
 import numpy as np
 
+from spine.math.cluster import DBSCAN
+
 from spine.data import Meta
 
 from spine.utils.globals import DELTA_SHP, SHAPE_PREC
 from spine.utils.particles import process_particle_event
 from spine.utils.ppn import image_coordinates
 from spine.utils.conditional import larcv
-from spine.utils.numba_local import dbscan
 
 from .base import ParserBase
 from .sparse import (
@@ -185,11 +186,13 @@ def __init__(self, dtype, particle_event=None, add_particle_info=False,
         self.type_include_mpr = type_include_mpr
         self.type_include_secondary = type_include_secondary
         self.primary_include_mpr = primary_include_mpr
-        self.break_clusters = break_clusters
-        self.break_eps = break_eps
-        self.break_metric = break_metric
         self.shape_precedence = shape_precedence
 
+        # Intialize DBSCAN if the clusters are to be broken up
+        self.break_clusters = break_clusters
+        if break_clusters:
+            self.dbscan = DBSCAN(break_eps, min_samples=1, metric=break_metric)
+
         # Intialize the sparse and particle parsers
         self.sparse_parser = Sparse3DParser(dtype, sparse_event='dummy')
 
@@ -334,8 +337,7 @@ def process(self, cluster_event, particle_event=None,
 
                 # If requested, break cluster into detached pieces
                 if self.break_clusters:
-                    frag_labels = dbscan(
-                            voxels, self.break_eps, self.break_metric)
+                    frag_labels = self.dbscan.fit_predict(voxels)
                     features[1] = id_offset + frag_labels
                     id_offset += max(frag_labels) + 1