Sanitized indexing, add template-based chi2 attributes

spine/build/fragment.py

@@ -220,10 +220,23 @@ def _build_truth(self, label_tensor, points_label, depositions_label,
                             ref_tensor[index_ref, PART_COL], return_counts=True)
                     part_id = int(part_ids[np.argmax(counts)])
                 if part_id > -1:
+                    # Load the MC particle information
                     assert part_id < len(particles), (
                             "Invalid particle ID found in fragment labels.")
-                    fragment = TruthFragment(**particles[part_id].as_dict())
+                    particle = particles[part_id]
+                    fragment = TruthFragment(**particle.as_dict())
+
+                    # Override the indexes of the fragment but preserve them
+                    fragment.orig_id = part_id
+                    fragment.orig_group_id = particle.group_id
+                    fragment.orig_parent_id = particle.parent_id
+                    fragment.orig_children_id = particle.children_id
+
                     fragment.id = i
+                    fragment.group_id = i
+                    fragment.parent_id = i
+                    fragment.children_id = np.empty(
+                            0, dtype=fragment.orig_children_id.dtype)
 
             # Fill long-form attributes
             if truth_only:
@@ -243,7 +256,7 @@ def _build_truth(self, label_tensor, points_label, depositions_label,
                         index_g4 = np.where(
                                 label_g4_tensor[:, CLUST_COL] == frag_id)[0]
                         fragment.index_g4 = index_g4
-                        fragment.points_g4 = poins_g4[index_g4]
+                        fragment.points_g4 = points_g4[index_g4]
                         fragment.depositions_g4 = depositions_g4[index_g4]
 
             else:

spine/build/manager.py

@@ -34,6 +34,7 @@ class BuildManager:
             ('label_adapt_tensor', ('clust_label_adapt',)),
             ('label_g4_tensor', ('clust_label_g4',)),
             ('depositions_q_label', ('charge_label',)),
+            ('graph_label', ('graph_label',)),
             ('sources', ('sources_adapt', 'sources')),
             ('sources_label', ('sources_label',)),
             ('particles', ('particles',)),

spine/build/particle.py

@@ -5,7 +5,9 @@
 from scipy.special import softmax
 
 from spine.data.out import RecoParticle, TruthParticle
+
 from spine.utils.globals import COORD_COLS, VALUE_COL, GROUP_COL, TRACK_SHP
+from spine.utils.gnn.network import filter_invalid_nodes
 
 from .base import BuilderBase
 
@@ -39,7 +41,8 @@ class ParticleBuilder(BuilderBase):
 
     # Necessary/optional data products to build a truth object
     _build_truth_keys = (
-            ('particles', False), ('truth_fragments', False),
+            ('particles', False), ('graph_label', False),
+            ('truth_fragments', False),
             *BuilderBase._build_truth_keys
     )
 
@@ -165,7 +168,8 @@ def _build_truth(self, particles, 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, sources=None, truth_fragments=None):
+                     sources_label=None, sources=None, graph_label=None,
+                     truth_fragments=None):
         """Builds :class:`TruthParticle` objects from the full chain output.
 
         Parameters
@@ -199,6 +203,8 @@ def _build_truth(self, particles, label_tensor, points_label,
             (N', 2) Tensor which contains the label module/tpc information
         sources : np.ndarray, optional
             (N, 2) Tensor which contains the module/tpc information
+        graph_label : np.ndarray, optional
+            (E, 2) Parentage relations in the set of particles
         truth_fragments : List[TruthFragment], optional
             (F) List of true fragments
 
@@ -222,9 +228,17 @@ def _build_truth(self, particles, label_tensor, points_label,
             assert particle.id == group_id, (
                     "The ordering of the true particles is wrong.")
 
-            # Override the index of the particle but preserve it
+            # Override the index of the particle and its group, but preserve it
             particle.orig_id = group_id
+            particle.orig_group_id = group_id
+            particle.orig_parent_id = particle.parent_id
+            particle.orig_children_id = particle.children_id
+
             particle.id = i
+            particle.group_id = i
+            particle.parent_id = i
+            particle.children_id = np.empty(
+                    0, dtype=particle.orig_children_id.dtype)
 
             # Update the deposited energy attribute by summing that of all
             # particles in the group (LArCV definition != SPINE definition)
@@ -268,6 +282,23 @@ def _build_truth(self, particles, label_tensor, points_label,
             # Append
             truth_particles.append(particle)
 
+        # If the parentage relations of non-empty particles are available,
+        # use them to assign parent/children IDs in the new particle set
+        if graph_label is not None:
+            # Narrow down the list of edges to those connecting visible particles
+            inval = set(np.unique(graph_label)).difference(set(valid_group_ids))
+            if len(inval) > 0:
+                graph_label = filter_invalid_nodes(graph_label, tuple(inval))
+
+            # Use the remaining edges to build parantage relations
+            mapping = {group_id: i for i, group_id in enumerate(valid_group_ids)}
+            for (source, target) in graph_label:
+                parent = truth_particles[mapping[source]]
+                child = truth_particles[mapping[target]]
+
+                child.parent_id = parent.id
+                parent.children_id = np.append(parent.children_id, child.id)
+
         return truth_particles
 
     def load_reco(self, data):

spine/data/out/fragment.py

@@ -131,7 +131,13 @@ class TruthFragment(Particle, FragmentBase, TruthBase):
     Attributes
     ----------
     orig_interaction_id : int
-        Unaltered index of the interaction in the original MC paricle list
+        Unaltered index of the interaction in the original MC particle list
+    orig_parent_id : int
+        Unaltered index of the particle parent in the original MC particle list
+    orig_group_id : int
+        Unaltered index of the particle group in the original MC particle list
+    orig_children_id : np.ndarray
+        Unaltered list of the particle children in the original MC particle list
     children_counts : np.ndarray
         (P) Number of truth child fragment of each shape
     reco_length : float
@@ -143,6 +149,9 @@ class TruthFragment(Particle, FragmentBase, TruthBase):
         to track objects)
     """
     orig_interaction_id: int = -1
+    orig_parent_id: int = -1
+    orig_group_id: int = -1
+    orig_children_id: np.ndarray = -1
     children_counts: np.ndarray = None
     reco_length: float = -1.
     reco_start_dir: np.ndarray = None
@@ -157,7 +166,7 @@ class TruthFragment(Particle, FragmentBase, TruthBase):
 
     # Variable-length attributes
     _var_length_attrs = (
-            ('children_counts', np.int32),
+            ('orig_children_id', np.int64), ('children_counts', np.int32),
             *TruthBase._var_length_attrs,
             *Particle._var_length_attrs
     )

spine/data/out/particle.py

@@ -35,8 +35,13 @@ class ParticleBase:
         Semantic type (shower (0), track (1), Michel (2), delta (3),
         low energy scatter (4)) of this particle
     pid : int
-        Particle spcies (Photon (0), Electron (1), Muon (2), Charged Pion (3),
-        Proton (4)) of this particle
+        Particle species (Photon (0), Electron (1), Muon (2), Charged Pion (3),
+        Proton (4), Kaon (5)) of this particle
+    chi2_pid : int
+        Particle species as predicted by the chi2 template method (Muon (2),
+        Charged Pion (3), Proton (4), Kaon (5)) of this particle
+    chi2_per_pid : np.ndarray
+        (P) Array of chi2 values associated with each particle class
     pdg_code : int
         PDG code corresponding to the PID number
     is_primary : bool
@@ -61,11 +66,11 @@ class ParticleBase:
     csda_ke : float
         Kinetic energy reconstructed from the particle range in MeV
     csda_ke_per_pid : np.ndarray
-        Same as `csda_ke` but for every available track PID hypothesis
+        (P) Same as `csda_ke` but for every available track PID hypothesis
     mcs_ke : float
         Kinetic energy reconstructed using the MCS method in MeV
     mcs_ke_per_pid : np.ndarray
-        Same as `mcs_ke` but for every available track PID hypothesis
+        (P) Same as `mcs_ke` but for every available track PID hypothesis
     momentum : np.ndarray
         3-momentum of the particle at the production point in MeV/c
     p : float
@@ -80,6 +85,8 @@ class ParticleBase:
     interaction_id: int = -1
     shape: int = -1
     pid: int = -1
+    chi2_pid: int = -1
+    chi2_per_pid: np.ndarray = None
     pdg_code: int = -1
     is_primary: bool = False
     length: float = -1.
@@ -102,6 +109,7 @@ class ParticleBase:
     _fixed_length_attrs = (
             ('start_point', 3), ('end_point', 3), ('start_dir', 3),
             ('end_dir', 3), ('momentum', 3),
+            ('chi2_per_pid', len(PID_LABELS) - 1),
             ('csda_ke_per_pid', len(PID_LABELS) - 1),
             ('mcs_ke_per_pid', len(PID_LABELS) - 1)
     )
@@ -197,7 +205,7 @@ class RecoParticle(ParticleBase, RecoBase):
     Attributes
     ----------
     pid_scores : np.ndarray
-        (P) Array of softmax scores associated with each of particle class
+        (P) Array of softmax scores associated with each particle class
     primary_scores : np.ndarray
         (2) Array of softmax scores associated with secondary and primary
     ppn_ids : np.ndarray
@@ -411,7 +419,13 @@ class TruthParticle(Particle, ParticleBase, TruthBase):
     Attributes
     ----------
     orig_interaction_id : int
-        Unaltered index of the interaction in the original MC paricle list
+        Unaltered index of the interaction in the original MC particle list
+    orig_parent_id : int
+        Unaltered index of the particle parent in the original MC particle list
+    orig_group_id : int
+        Unaltered index of the particle group in the original MC particle list
+    orig_children_id : np.ndarray
+        Unaltered list of the particle children in the original MC particle list
     children_counts : np.ndarray
         (P) Number of truth child particle of each shape
     reco_length : float
@@ -427,6 +441,9 @@ class TruthParticle(Particle, ParticleBase, TruthBase):
         Best-guess reconstructed momentum of the particle
     """
     orig_interaction_id: int = -1
+    orig_parent_id: int = -1
+    orig_group_id: int = -1
+    orig_children_id: np.ndarray = -1
     children_counts: np.ndarray = None
     reco_length: float = -1.
     reco_start_dir: np.ndarray = None
@@ -443,7 +460,7 @@ class TruthParticle(Particle, ParticleBase, TruthBase):
 
     # Variable-length attributes
     _var_length_attrs = (
-            ('children_counts', np.int32),
+            ('orig_children_id', np.int64), ('children_counts', np.int32),
             *TruthBase._var_length_attrs,
             *ParticleBase._var_length_attrs,
             *Particle._var_length_attrs

spine/data/particle.py

@@ -27,8 +27,6 @@ class Particle(PosDataBase):
         Index in the original MCTruth array from whence it came
     mcst_index : int
         Index in the original MCTrack/MCShower array from whence it came
-    gen_id : int
-        Index of the particle at the generator level
     group_id : int
         Index of the group the particle belongs to
     interaction_id : int
@@ -110,14 +108,13 @@ class Particle(PosDataBase):
     id: int = -1
     mct_index: int = -1
     mcst_index: int = -1
-    gen_id: int = -1
     group_id: int = -1
     interaction_id: int = -1
     nu_id: int = -1
     interaction_primary: int = -1
     group_primary: int = -1
     parent_id: int = -1
-    children_id: int = None
+    children_id: np.ndarray = None
     track_id: int = -1
     parent_track_id: int = -1
     ancestor_track_id: int = -1
@@ -248,7 +245,7 @@ def from_larcv(cls, particle):
         for prefix in ('', 'parent_', 'ancestor_'):
             for key in ('track_id', 'pdg_code', 'creation_process', 't'):
                 obj_dict[prefix+key] = getattr(particle, prefix+key)()
-        for key in ('id', 'gen_id', 'group_id', 'interaction_id', 'parent_id',
+        for key in ('id', 'group_id', 'interaction_id', 'parent_id',
                     'mct_index', 'mcst_index', 'num_voxels', 'shape',
                     'energy_init', 'energy_deposit', 'distance_travel'):
             if not hasattr(particle, key):

spine/io/collate.py

@@ -169,12 +169,18 @@ def __call__(self, batch):
                         has_batch_col=True, coord_cols=coord_cols)
 
             elif isinstance(ref_obj, tuple) and len(ref_obj) == 2:
-                # Case where an index and an offset is provided per entry.
+                # Case where an index and a count is provided per entry.
+                # Start by computing the necessary node ID offsets to apply
+                total_counts = [sample[key][1] for sample in batch]
+                offsets = np.zeros(len(total_counts), dtype=int)
+                offsets[1:] = np.cumsum(total_counts)[:-1]
+
                 # Stack the indexes, do not add a batch column
-                tensor  = np.concatenate(
-                        [sample[key][0] for sample in batch], axis=1)
+                tensor_list = []
+                for i, sample in enumerate(batch):
+                    tensor_list.append(sample[key][0] + offsets[i])
+                tensor = np.concatenate(tensor_list, axis=1)
                 counts  = [sample[key][0].shape[-1] for sample in batch]
-                offsets = [sample[key][1] for sample in batch]
 
                 if len(tensor.shape) == 1:
                     data[key] = IndexBatch(tensor, counts, offsets)

spine/io/parse/cluster.py

@@ -3,6 +3,8 @@
 Contains the following parsers:
 - :class:`Cluster2DParser`
 - :class:`Cluster3DParser`
+- :class:`Cluster3DAggregateParser`
+- :class:`Cluster3DChargeRescaledParser`
 """
 
 from warnings import warn