Merge remote-tracking branch 'dkoh/develop' into me

analysis/classes/Interaction.py

@@ -233,7 +233,7 @@ def particles_summary(self):
         if self._particles is None: return
         for p in sorted(self._particles.values(), key=lambda x: x.is_primary, reverse=True):
             pmsg = "    {} Particle {}: PID = {}, Size = {}, Match = {} \n".format(
-                primary_str[p.is_primary], p.id, PID_LABELS[p.pid], p.size, str(p.match))
+                primary_str[p.is_primary], p.id, p.pid, p.size, str(p.match))
             self._particles_summary += pmsg
         return self._particles_summary
 

analysis/classes/Particle.py

@@ -74,8 +74,6 @@ def __init__(self,
                  volume_id: int = -1,
                  image_id: int = -1, 
                  semantic_type: int = -1, 
-                 pid: int = -1,
-                 is_primary: int = -1,
                  index: np.ndarray = np.empty(0, dtype=np.int64), 
                  points: np.ndarray = np.empty(0, dtype=np.float32),
                  depositions: np.ndarray = np.empty(0, dtype=np.float32), 
@@ -90,12 +88,12 @@ def __init__(self,
 
         # Initialize private attributes to be assigned through setters only
         self._num_fragments = None
-        self._index = np.array(index, dtype=np.int64)
-        self._size = len(self._index)
-        self._depositions = np.atleast_1d(depositions)
-        self._depositions_sum = np.sum(self._depositions)
-        self._pid_scores = pid_scores
-        self._primary_scores = primary_scores
+        self._index = None
+        self._depositions = None
+        self._depositions_sum = -1
+        self._pid = -1
+        self._size = -1
+        self._is_primary = -1
 
         # Initialize attributes
         self.id             = int(group_id)
@@ -107,15 +105,20 @@ def __init__(self,
         self.semantic_type  = int(semantic_type)
         self.points         = points
 
-        if np.all(pid_scores < 0):
-            self._pid = pid
-        else:
-            self._pid = int(np.argmax(pid_scores))
+        self.index = index
+        self.depositions = depositions
+        self.pid_scores = pid_scores
+        self.primary_scores = primary_scores
 
-        if np.all(primary_scores < 0):
-            self._is_primary = is_primary
-        else:
-            self._is_primary = int(np.argmax(primary_scores))
+        # if np.all(pid_scores < 0):
+        #     self._pid = pid
+        # else:
+        #     self._pid = int(np.argmax(pid_scores))
+
+        # if np.all(primary_scores < 0):
+        #     self._is_primary = is_primary
+        # else:
+        #     self._is_primary = int(np.argmax(primary_scores))
 
         self.start_point    = start_point
         self.end_point      = end_point
@@ -205,7 +208,7 @@ def index(self):
     @index.setter
     def index(self, index):
         # Count the number of voxels
-        self._index = index
+        self._index = np.array(index, dtype=np.int64)
         self._size = len(index)
 
     @property
@@ -214,7 +217,7 @@ def depositions_sum(self):
         Total amount of charge/energy deposited. This attribute has no setter,
         as it can only be set by providing a set of depositions.
         '''
-        return self._depositions_sum
+        return float(self._depositions_sum)
 
     @property
     def depositions(self):
@@ -240,14 +243,13 @@ def pid_scores(self):
 
     @pid_scores.setter
     def pid_scores(self, pid_scores):
+        self._pid_scores = pid_scores
         # If no PID scores are providen, the PID is unknown
         if pid_scores[0] < 0.:
-            self._pid_scores = pid_scores
             self._pid = -1
-
+        else:
         # Store the PID scores
-        self._pid_scores = pid_scores
-        self._pid = int(np.argmax(pid_scores))
+            self._pid = int(np.argmax(pid_scores))
 
     @property
     def pid(self):
@@ -266,8 +268,8 @@ def primary_scores(self, primary_scores):
         # If no primary scores are given, the primary status is unknown
         if primary_scores[0] < 0.:
             self._primary_scores = primary_scores
-            self.is_primary = -1
+            self._is_primary = -1
 
         # Store the PID scores and give a best guess
         self._primary_scores = primary_scores
-        self.is_primary = np.argmax(primary_scores)
+        self._is_primary = np.argmax(primary_scores)

analysis/classes/TruthInteraction.py

@@ -34,6 +34,8 @@ def __init__(self,
 
         # Initialize private attributes to be set by setter only
         self._particles  = None
+        self._particle_counts = np.zeros(6, dtype=np.int64)
+        self._primary_counts  = np.zeros(6, dtype=np.int64)
         # Invoke particles setter
         self.particles   = particles
 
@@ -87,6 +89,13 @@ def particles(self, value):
                 true_depositions_list.append(p.truth_depositions)
                 true_depositions_MeV_list.append(p.truth_depositions_MeV)
 
+                if p.pid >= 0:
+                    self._particle_counts[p.pid] += 1
+                    self._primary_counts[p.pid] += int(p.is_primary)
+                else:
+                    self._particle_counts[-1] += 1
+                    self._primary_counts[-1] += int(p.is_primary)
+
             self._particle_ids          = np.array(id_list, dtype=np.int64)
             self._num_particles         = len(value)
             self._num_primaries         = len([1 for p in value if p.is_primary])
@@ -122,11 +131,6 @@ def from_particles(cls, particles, verbose=False, **kwargs):
 
         _process_interaction_attributes(init_args, processed_args, **kwargs)
 
-        for i, t in enumerate(init_args['truth_depositions_MeV']):
-            if len(t.shape) == 0:
-                print(t, t.shape)
-                print(init_args['truth_index'][i])
-
         # Handle depositions_MeV for TruthParticles
         processed_args['depositions_MeV']       = np.concatenate(init_args['depositions_MeV'])
         processed_args['truth_depositions']     = np.concatenate(init_args['truth_depositions'])

analysis/classes/TruthParticle.py

@@ -2,7 +2,7 @@
 
 from typing import Counter, List, Union
 from . import Particle
-
+from mlreco.utils.globals import PDG_TO_PID
 
 class TruthParticle(Particle):
     '''
@@ -36,15 +36,21 @@ class TruthParticle(Particle):
     def __init__(self, 
                  *args, 
                  depositions_MeV: np.ndarray = np.empty(0, dtype=np.float32),
+                 pid: int = -1,
+                 is_primary: int = -1,
                  truth_index: np.ndarray = np.empty(0, dtype=np.int64), 
                  truth_points: np.ndarray = np.empty((0,3), dtype=np.float32),
                  truth_depositions: np.ndarray = np.empty(0, dtype=np.float32),
                  truth_depositions_MeV: np.ndarray = np.empty(0, dtype=np.float32),
                  momentum: np.ndarray = -np.ones(3, dtype=np.float32),
                  particle_asis: object = None, 
                  **kwargs):
+
         super(TruthParticle, self).__init__(*args, **kwargs)
 
+        self._pid = pid
+        self._is_primary = is_primary
+
         # Initialize attributes
         self.depositions_MeV        = np.atleast_1d(depositions_MeV)
         self.truth_index            = truth_index
@@ -57,6 +63,7 @@ def __init__(self,
             self.end_position      = particle_asis.end_position()
 
         self.asis = particle_asis
+        assert PDG_TO_PID[int(self.asis.pdg_code())] == self.pid
 
         self.start_point = np.array([getattr(particle_asis.first_step(), a)() \
             for a in ['x', 'y', 'z']], dtype=np.float32)
@@ -69,6 +76,15 @@ def __init__(self,
         if np.linalg.norm(self.momentum) > 0.:
             self.start_dir = self.momentum/np.linalg.norm(self.momentum)
 
+
+    @property
+    def pid(self):
+        return int(self._pid)
+
+    @property
+    def is_primary(self):
+        return self._is_primary
+
     def __repr__(self):
         msg = super(TruthParticle, self).__repr__()
         return 'Truth'+msg

analysis/classes/builders.py

@@ -353,6 +353,7 @@ def _build_truth(self,
         for i, lpart in enumerate(larcv_particles):
             id = int(lpart.id())
             pdg = PDG_TO_PID.get(lpart.pdg_code(), -1)
+            # print(pdg)
             is_primary = lpart.group_id() == lpart.parent_id()
             mask_nonghost = labels_nonghost[:, 6].astype(int) == id
             if np.count_nonzero(mask_nonghost) <= 0: