diff --git a/fuse/plugins/truth_information/cluster_tagging.py b/fuse/plugins/truth_information/cluster_tagging.py index 28e87499..ec3cd104 100644 --- a/fuse/plugins/truth_information/cluster_tagging.py +++ b/fuse/plugins/truth_information/cluster_tagging.py @@ -1,4 +1,5 @@ import strax +import straxen import numpy as np export, __all__ = strax.exporter() @@ -6,16 +7,17 @@ @export class ClusterTagging(strax.Plugin): - """Plugin to tag if clusters contribute to the main or alternative - s1/s2.""" + """Plugin to tag if clusters contribute to the main or alternative s1/s2 in + an event, or successfully reconstructed as s0/s1/s2 peaks.""" - __version__ = "0.0.3" + __version__ = "0.0.4" depends_on = ("microphysics_summary", "photon_summary", "peak_basics", "event_basics") provides = "tagged_clusters" data_kind = "interactions_in_roi" dtype = [ + # Tags by events ("in_main_s1", np.bool_), ("in_main_s2", np.bool_), ("in_alt_s1", np.bool_), @@ -24,20 +26,43 @@ class ClusterTagging(strax.Plugin): ("photons_in_main_s2", np.int32), ("photons_in_alt_s1", np.int32), ("photons_in_alt_s2", np.int32), + # Tags by S1 peaks + ("has_s1", np.bool_), + ("s1_tight_coincidence", np.int32), ] + strax.time_fields - def compute(self, interactions_in_roi, propagated_photons, peaks, events): - peaks_in_event = strax.split_by_containment(peaks, events) - photon_in_event = strax.split_touching_windows(propagated_photons, events) + photon_finding_window = straxen.URLConfig( + default=200, + type=int, + help="Time window [ns] that defines whether a photon is in a peak. " + "Peaks' start and end times are extended by this window to find photons in them.", + ) + def compute(self, interactions_in_roi, propagated_photons, peaks, events): result = np.zeros(len(interactions_in_roi), dtype=self.dtype) result["time"] = interactions_in_roi["time"] result["endtime"] = interactions_in_roi["endtime"] + # First we tag the clusters that are in a peak + s1_peaks = peaks[peaks["type"] == 1] + photons_in_peak = strax.split_touching_windows( + interactions_in_roi, s1_peaks, window=self.photon_finding_window + ) + for peak, photons in zip(s1_peaks, photons_in_peak): + mask = np.isin(interactions_in_roi["cluster_id"], photons["cluster_id"]) + result["has_s1"][mask] = True + result["s1_tight_coincidence"][mask] = peak["tight_coincidence"] + + # Then we tag the clusters that are in an event's main or alternative s1/s2 + peaks_in_event = strax.split_by_containment(peaks, events) + photon_in_event = strax.split_touching_windows(propagated_photons, events) + for i, (peaks_of_event, event_i, photon_of_event) in enumerate( zip(peaks_in_event, events, photon_in_event) ): - peak_photons = strax.split_touching_windows(photon_of_event, peaks_of_event) + peak_photons = strax.split_touching_windows( + photon_of_event, peaks_of_event, window=self.photon_finding_window + ) peak_name_dict = { "main_s1": "s1_index",