Cathode crosser PR: Bug fixes for cdist and already matched particles

spine/data/neutrino.py

@@ -34,7 +34,7 @@ class Neutrino(PosDataBase):
     pdg_code : int
         PDG code of the neutrino
     lepton_pdg_code : int
-        PDF code of the outgoing lepton
+        PDG code of the outgoing lepton
     current_type : int
         Enumerated current type of the neutrino interaction
     interaction_mode : int

spine/data/optical.py

@@ -21,7 +21,7 @@ class Flash(PosDataBase):
     id : int
         Index of the flash in the list
     volume_id : int
-        Index of the optical volume in which the flahs was recorded
+        Index of the optical volume in which the flash was recorded
     time : float
         Time with respect to the trigger in microseconds
     time_width : float
@@ -111,3 +111,66 @@ def from_larcv(cls, flash):
                    time_abs=flash.absTime(), time_width=flash.timeWidth(),
                    total_pe=flash.TotalPE(), pe_per_ch=pe_per_ch,
                    center=center, width=width)
+
+    def merge(self, other):
+        """Merge another flash into this one.
+
+        The merging strategy proceeds as follows:
+        - The earlier flash takes precedence over the later flash as far as
+          all timing-related information is concerned (flash time, etc.)
+        - The combined flash centroid is produced by taking the weighted average
+          of the two existing flash centroids
+        - The PE values in each light collection system are added together, so
+          is the total PE value of the combined flash
+
+        Parameters
+        ----------
+        other : Flash
+            Flash to merge into this one
+        """
+        # Check that the position units are the same
+        assert self.units == other.units, (
+                "The units of the flash to be merged do not match.")
+
+        # Determine the flash window end points (to merge time widths later)
+        end_i, end_j = self.time + self.time_width, other.time + other.time_width
+
+        # If the other flash happened first, update the timing information
+        if self.time > other.time:
+            self.time = other.time
+            self.time_abs = other.time_abs
+            self.on_beam_time = other.on_beam_time
+            self.frame = other.frame
+            self.in_beam_frame = other.in_beam_frame
+
+        # Take the union of the two time widths as the new combined width
+        self.time_width = max(end_i, end_j) - self.time
+
+        # Take the weighted average of the centroids to compute the new one
+        valid_mask = (self.width > 0.) & (other.width > 0.)
+
+        w_i, w_j = 1./self.width**2, 1./other.width**2
+        self.center = (w_i*self.center + w_j*other.center)/(w_i + w_j)
+
+        self.width = 1./np.sqrt(w_i + w_j)
+        self.width[~valid_mask] = -1.
+
+        # Compute the new total PE and fast light component to total ratio
+        t_i, t_j = self.total_pe, other.total_pe
+        self.total_pe = t_i + t_j
+
+        r_i, r_j = self.fast_to_total, other.fast_to_total
+        self.fast_to_total = (r_i*t_i + r_j*t_j)/(t_i + t_j)
+
+        # Merge the PE count in each PMT
+        pe_per_ch = np.zeros(
+                max(len(self.pe_per_ch), len(other.pe_per_ch)),
+                dtype=self.pe_per_ch.dtype)
+        pe_per_ch[:len(self.pe_per_ch)] += self.pe_per_ch
+        pe_per_ch[:len(other.pe_per_ch)] += other.pe_per_ch
+
+        self.pe_per_ch = pe_per_ch
+
+        # The new volume ID is invalid if the two original volumes differ
+        if self.volume_id != other.volume_id:
+            self.volume_id = -1

spine/data/out/base.py

@@ -89,6 +89,11 @@ class OutBase(PosDataBase):
     # Attributes that must not be stored to file when storing lite files
     _lite_skip_attrs = ('index',)
 
+    def reset_match(self):
+        """Resets the reco/truth matching information for the object."""
+        self.is_matched = False
+        self.match_ids = np.empty(0, dtype=np.int64)
+
     @property
     def size(self):
         """Total number of voxels that make up the object.

spine/data/out/interaction.py

@@ -120,6 +120,28 @@ def __str__(self):
 
         return info
 
+    def reset_flash_match(self, typed=True):
+        """Reset all the flash matching attributes.
+
+        Parameters
+        ----------
+        typed : bool, default True
+            If `True`, the underlying arrays are reset to typed empty arrays
+        """
+        self.is_flash_matched = False
+        self.flash_total_pe = -1.
+        self.flash_type_pe = -1.
+        if typed:
+            self.flash_ids = np.empty(0, dtype=np.int32)
+            self.flash_volume_ids = np.empty(0, dtype=np.int32)
+            self.flash_times = np.empty(0, dtype=np.float32)
+            self.flash_scores = np.empty(0, dtype=np.float32)
+        else:
+            self.flash_ids = []
+            self.flash_volume_ids = []
+            self.flash_times = []
+            self.flash_scores = []
+
     @property
     def primary_particles(self):
         """List of primary particles associated with this interaction.

spine/data/out/particle.py

@@ -196,13 +196,6 @@ def p(self):
     def p(self, p):
         pass
 
-    def unmatch(self):
-        """
-        Unmatch the particle from its reco or truth particle match.
-        """
-        self.match_ids = []
-        self.is_matched = False
-
 
 @dataclass(eq=False)
 @inherit_docstring(RecoBase, ParticleBase)
@@ -224,6 +217,8 @@ class RecoParticle(ParticleBase, RecoBase):
         interaction vertex position in cm
     start_dedx : float
         dE/dx around a user-defined neighborhood of the start point in MeV/cm
+    end_dedx : float
+        dE/dx around a user-defined neighborhood of the end point in MeV/cm
     start_straightness : float
         Explained variance ratio of the beginning of the particle
     directional_spread : float
@@ -238,6 +233,7 @@ class RecoParticle(ParticleBase, RecoBase):
     ppn_points: np.ndarray = None
     vertex_distance: float = -1.
     start_dedx: float = -1.
+    end_dedx: float = -1.
     start_straightness: float = -1.
     directional_spread: float = -1.
     axial_spread: float = -np.inf

spine/io/collate.py

@@ -253,7 +253,7 @@ def stack_feat_tensors(self, batch, key):
         # Dispatch
         if not self.split or sources is None:
             tensor = np.concatenate([sample[key].features for sample in batch])
-            counts = [len(sample[key]) for sample in batch]
+            counts = [len(sample[key].features) for sample in batch]
 
         else:
             batch_size = len(batch)

spine/io/parse/misc.py

@@ -4,7 +4,6 @@
 - :class:`Meta2DParser`
 - :class:`Meta3DParser`
 - :class:`RunInfoParser`
-- :class:`OpFlashParser`
 - :class:`CRTHitParser`
 - :class:`TriggerParser`
 """
@@ -13,6 +12,7 @@
 from spine.data import Meta, RunInfo, Flash, CRTHit, Trigger
 
 from spine.utils.conditional import larcv
+from spine.utils.optical import FlashMerger
 
 from .base import ParserBase
 from .data import ParserObjectList
@@ -67,7 +67,7 @@ def __init__(self, projection_id=None, **kwargs):
         projection_id : int, optional
             Projection ID to get the 2D image from (if fetching from 2D)
         **kwargs : dict, optional
-            Data product arguments to be passed to the `process` function
+            data product arguments to be passed to the `process` function
         """
         # Initialize the parent class
         super().__init__(**kwargs)
@@ -186,6 +186,24 @@ class FlashParser(ParserBase):
     # Type of object(s) returned by the parser
     returns = 'object_list'
 
+    def __init__(self, merge=None, **kwargs):
+        """Initialize the flash parser.
+
+        Parameters
+        ----------
+        merge : dict, optional
+            Flash merging configuration
+        **kwargs : dict, optional
+            data product arguments to be passed to the `process` function
+        """
+        # Initialize the parent class
+        super().__init__(**kwargs)
+
+        # Initialize the flash merging class, if needed
+        self.merger = None
+        if merge is not None:
+            self.merger = FlashMerger(**merge)
+
     def __call__(self, trees):
         """Parse one entry.
 
@@ -238,6 +256,10 @@ def process(self, flash_event=None, flash_event_list=None):
                     flashes.append(flash)
                     idx += 1
 
+        # If requested, merge flashes which match in time
+        if self.merger is not None:
+            flashes, _ = merger(flashes)
+
         return ParserObjectList(flashes, Flash())
 
 

spine/io/parse/sparse.py

@@ -13,7 +13,7 @@
 from spine.data import Meta
 
 from spine.utils.globals import GHOST_SHP, SHAPE_PREC
-from spine.utils.ghost import compute_rescaled_charge
+from spine.utils.ghost import ChargeRescaler
 from spine.utils.conditional import larcv
 
 from .base import ParserBase
@@ -457,9 +457,8 @@ def __init__(self, dtype, collection_only=False, collection_id=2, **kwargs):
         # Initialize the parent class
         super().__init__(dtype, **kwargs)
 
-        # Store the revelant attributes
-        self.collection_only = collection_only
-        self.collection_id = collection_id
+        # Initialize the charge rescaler
+        self.rescaler = ChargeRescaler(collection_only, collection_id)
 
     def __call__(self, trees):
         """Parse one entry.
@@ -497,10 +496,7 @@ def process_rescale(self, sparse_event_list):
 
         # Use individual hit informations to compute a rescaled charge
         deghost_mask = np.where(tensor.features[:, -1] < GHOST_SHP)[0]
-        charges = compute_rescaled_charge(
-                tensor.features[deghost_mask, :-1],
-                collection_only=self.collection_only,
-                collection_id=self.collection_id)
+        charges = self.rescaler.process_single(tensor.features[deghost_mask, :-1])
 
         tensor.features = charges[:, None]
 

spine/model/full_chain.py

@@ -17,13 +17,14 @@
 # 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.ghost import ChargeRescaler
 from spine.utils.calib import CalibrationManager
 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 (
         form_clusters_batch, get_cluster_label_batch)
@@ -164,6 +165,11 @@ def __init__(self, chain, uresnet_deghost=None, uresnet=None,
                     "`uresnet_deghost` configuration block.")
             self.uresnet_deghost = UResNetSegmentation(uresnet_deghost)
 
+        # Initialize the charge rescaling process (adapt to ghost predictions)
+        if self.charge_rescaling is not None:
+            self.charge_rescaler = ChargeRescaler(
+                    collection_only=self.charge_rescaling == 'collection')
+
         # Initialize the semantic segmentation model (+ point proposal)
         if self.segmentation is not None and self.segmentation == 'uresnet':
             assert (uresnet is not None) ^ (uresnet_ppn is not None), (
@@ -363,8 +369,7 @@ def run_deghosting(self, data, sources=None, seg_label=None,
 
             # Rescale the charge, if requested
             if self.charge_rescaling is not None:
-                charges = compute_rescaled_charge_batch(
-                        data_adapt, self.charge_rescaling == 'collection')
+                charges = self.charge_rescaler(data_adapt)
                 tensor_deghost = data_adapt.tensor[:, :-6]
                 tensor_deghost[:, VALUE_COL] = charges
                 data_adapt.data = tensor_deghost