Merge pull request #81 from francois-drielsma/develop

Refactored particle point assignment in the full chain

Author: francois-drielsma

spine/build/manager.py

@@ -70,7 +70,7 @@ def __init__(self, fragments, particles, interactions,
         assert units in self._units, (
                 f"Units not recognized: {units}. Must be one {self._units}")
         self.units = units
-        
+
         # If custom sources are provided, update the tuple
         if sources is not None:
             sources_dict = dict(self._sources)
@@ -92,7 +92,7 @@ def __init__(self, fragments, particles, interactions,
         if particles:
             self.builders['particle'] = ParticleBuilder(mode, units)
         if interactions:
-            assert particles is not None, (
+            assert particles, (
                     "Interactions are built from particles. If `interactions` "
                     "is True, so must `particles` be.")
             self.builders['interaction'] = InteractionBuilder(mode, units)
@@ -178,7 +178,7 @@ def build_sources(self, data, entry=None):
 
             if 'sources' in sources:
                 update['sources'] = sources['sources'].astype(int)
-        
+
         if self.mode != 'reco':
             update['label_tensor'] = sources['label_tensor']
             update['points_label'] = sources['label_tensor'][:, COORD_COLS]

spine/io/read/base.py

@@ -184,9 +184,10 @@ def process_entry_list(self, n_entry=None, n_skip=None, entry_list=None,
                    entry_list is not None, skip_entry_list is not None,
                    run_event_list is not None,
                    skip_run_event_list is not None]):
-            assert (bool(n_entry or n_skip) ^
-                    bool(entry_list or skip_entry_list) ^
-                    bool(run_event_list or skip_run_event_list)), (
+            assert ((n_entry is not None or n_skip is not None) ^
+                    (entry_list is not None or skip_entry_list is not None) ^
+                    (run_event_list is not None
+                     or skip_run_event_list is not None)), (
                     "Cannot specify `n_entry` or `n_skip` at the same time "
                     "as `entry_list` or `skip_entry_list` or at the same time "
                     "as `run_event_list` or `skip_run_event_list`.")

spine/model/full_chain.py

@@ -17,12 +17,12 @@
 # TODO: raname it something more generic like ParticleClusterImageClassifier?
 
 from spine.data import TensorBatch, IndexBatch, RunInfo
+from spine.utils.logger import logger
 from spine.utils.globals import (
         COORD_COLS, VALUE_COL, CLUST_COL, SHAPE_COL, SHOWR_SHP, TRACK_SHP,
         MICHL_SHP, DELTA_SHP, GHOST_SHP)
 from spine.utils.calib import CalibrationManager
-from spine.utils.logger import logger
-from spine.utils.ppn import get_particle_points
+from spine.utils.ppn import ParticlePointPredictor
 from spine.utils.ghost import compute_rescaled_charge_batch
 from spine.utils.cluster.label import ClusterLabelAdapter
 from spine.utils.gnn.cluster import (
@@ -115,10 +115,10 @@ class FullChain(torch.nn.Module):
     )
 
     def __init__(self, chain, uresnet_deghost=None, uresnet=None,
-                 uresnet_ppn=None, adapt_labels=None, graph_spice=None,
-                 dbscan=None, grappa_shower=None, grappa_track=None,
-                 grappa_particle=None, grappa_inter=None, calibration=None,
-                 uresnet_deghost_loss=None, uresnet_loss=None,
+                 uresnet_ppn=None, adapt_labels=None, predict_points=None,
+                 graph_spice=None, dbscan=None, grappa_shower=None,
+                 grappa_track=None, grappa_particle=None, grappa_inter=None,
+                 calibration=None, uresnet_deghost_loss=None, uresnet_loss=None,
                  uresnet_ppn_loss=None, graph_spice_loss=None,
                  grappa_shower_loss=None, grappa_track_loss=None,
                  grappa_particle_loss=None, grappa_inter_loss=None):
@@ -134,6 +134,8 @@ def __init__(self, chain, uresnet_deghost=None, uresnet=None,
             Segmentation and point proposal model configuration
         adapt_labels : dict, optional
             Parameters for the cluster label adaptation (if non-standard)
+        predict_points : dict, optional
+            Parameters for the particle point predictor (if non-standard)
         dbscan : dict, optional
             Connected component clustering configuration
         graph_spice : dict, optional
@@ -175,6 +177,9 @@ def __init__(self, chain, uresnet_deghost=None, uresnet=None,
         # Initialize the relabeling process (adapt to the semantic predictions)
         self.label_adapter = ClusterLabelAdapter(**(adapt_labels or {}))
 
+        # Initialize the point predictor (for fragment/particle clusters)
+        self.point_predictor = ParticlePointPredictor(**(predict_points or {}))
+
         # Initialize the dense clustering model
         self.fragment_shapes = []
         if (self.fragmentation is not None and
@@ -1017,7 +1022,7 @@ def prepare_grappa_input(self, model, data, clusts, clust_shapes,
                 ref_clusts = clust_primaries
 
             # Get and store the points
-            points = get_particle_points(
+            points = self.point_predictor(
                     data, ref_clusts, clust_shapes, self.result['ppn_points'])
 
             grappa_input['points'] = points

spine/utils/globals.py

@@ -85,7 +85,6 @@
 KAON_PID = 5
 
 # Mapping between particle PDG code and particle ID labels
-PHOT_PID = 0
 PDG_TO_PID = defaultdict(lambda: -1)
 PDG_TO_PID.update({
     22:   PHOT_PID,

spine/utils/numba_local.py

@@ -642,9 +642,9 @@ def principal_components(x: nb.float32[:,:]) -> nb.float32[:,:]:
 
 @nb.njit(cache=True)
 def farthest_pair(x: nb.float32[:,:],
-                  algorithm: bool = 'brute') -> (nb.int32, nb.int32, nb.float32):
-    """Algorithm which finds the two points which are
-    farthest from each other in a set.
+                  algorithm: str = 'brute') -> (nb.int32, nb.int32, nb.float32):
+    """Algorithm which finds the two points which are farthest from each other
+    in a set.
 
     Two algorithms:
     - `brute`: compute pdist, use argmax
@@ -675,7 +675,9 @@ def farthest_pair(x: nb.float32[:,:],
         dist = dist_mat[idxs[0], idxs[1]]
 
     elif algorithm == 'recursive':
-        idxs, subidx, dist, tempdist = [0, 0], 0, 0., -1.
+        centroid = mean(x, 0)
+        start_idx = np.argmax(cdist(centroid.reshape(1, -1), x))
+        idxs, subidx, dist, tempdist = [start_idx, start_idx], 0, 0., -1.
         while dist > tempdist:
             tempdist = dist
             dists = cdist(np.ascontiguousarray(x[idxs[subidx]]).reshape(1,-1), x).flatten()