Added peak tagging plugin

straxen/plugins/event_info.py

@@ -1,57 +1,113 @@
-import strax
-from straxen import pre_apply_function
-export, __all__ = strax.exporter()
-
-
-@export
-@strax.takes_config(
-    strax.Option(
-        name='event_info_function',
-        default='pre_apply_function',
-        help="Function that must be applied to all event_info data. Do not change.",
-    )
-)
-class EventInfo(strax.MergeOnlyPlugin):
-    """
-    Plugin which merges the information of all event data_kinds into a
-    single data_type.
-    """
-    depends_on = ['events',
-                  'event_basics',
-                  'event_positions',
-                  'corrected_areas',
-                  'energy_estimates',
-                  # 'event_pattern_fit', <- this will be added soon
-                  ]
-    save_when = strax.SaveWhen.ALWAYS
-    provides = 'event_info'
-    __version__ = '0.0.1'
-
-    def compute(self, **kwargs):
-        event_info_function = self.config['event_info_function']
-        event_info = super().compute(**kwargs)
-        if event_info_function != 'disabled':
-            event_info = pre_apply_function(event_info,
-                                            self.run_id,
-                                            self.provides,
-                                            event_info_function,
-                                            )
-        return event_info
-
-
-@export
-class EventInfo1T(strax.MergeOnlyPlugin):
-    """
-    Plugin which merges the information of all event data_kinds into a
-    single data_type.
-
-    This only uses 1T data-types as several event-plugins are nT only
-    """
-    depends_on = ['events',
-                  'event_basics',
-                  'event_positions',
-                  'corrected_areas',
-                  'energy_estimates']
-    provides = 'event_info'
-    save_when = strax.SaveWhen.ALWAYS
-    __version__ = '0.0.0'
+import strax
+from straxen import pre_apply_function
+
+import numpy as np
+import numba
+export, __all__ = strax.exporter()
+
+
+@export
+@strax.takes_config(
+    strax.Option(
+        name='event_info_function',
+        default='pre_apply_function',
+        help="Function that must be applied to all event_info data. Do not change.",
+    )
+)
+class EventInfo(strax.MergeOnlyPlugin):
+    """
+    Plugin which merges the information of all event data_kinds into a
+    single data_type.
+    """
+    depends_on = ['events',
+                  'event_basics',
+                  'event_positions',
+                  'corrected_areas',
+                  'energy_estimates',
+                  # 'event_pattern_fit', <- this will be added soon
+                  ]
+    save_when = strax.SaveWhen.ALWAYS
+    provides = 'event_info'
+    __version__ = '0.0.1'
+
+    def compute(self, **kwargs):
+        event_info_function = self.config['event_info_function']
+        event_info = super().compute(**kwargs)
+        if event_info_function != 'disabled':
+            event_info = pre_apply_function(event_info,
+                                            self.run_id,
+                                            self.provides,
+                                            event_info_function,
+                                            )
+        return event_info
+
+
+@export
+class EventInfo1T(strax.MergeOnlyPlugin):
+    """
+    Plugin which merges the information of all event data_kinds into a
+    single data_type.
+
+    This only uses 1T data-types as several event-plugins are nT only
+    """
+    depends_on = ['events',
+                  'event_basics',
+                  'event_positions',
+                  'corrected_areas',
+                  'energy_estimates']
+    provides = 'event_info'
+    save_when = strax.SaveWhen.ALWAYS
+    __version__ = '0.0.0'
+
+
+class EventInfoVetos(strax.Plugin):
+    """
+    Plugin which combines event_info with the tagged peaks information
+    from muon- and neutron-veto.
+    """
+    __version__ = '0.0.0'
+    depends_on = ('event_info', 'peak_veto_tags')
+    provides = 'events_tagged'
+    save_when = strax.SaveWhen.TARGET
+
+    def infer_dtype(self):
+        dtype = []
+        dtype += strax.time_fields
+
+        for i in range(1, 3):
+            for peak_type in ['', 'alt_']:
+                dtype += [((f"Veto tag for {peak_type}S{i}: unatagged: 0, nveto: 1, mveto: 2, both: 3",
+                            f"{peak_type}s{i}_veto_tag"), np.int8),
+                          ((f"Time to closest veto interval for {peak_type}S{i}",
+                            f"{peak_type}s{i}_dt_veto"), np.int64),
+                          ]
+        dtype += [(('Number of tagged peaks inside event', 'n_tagged_peaks'), np.int16)]
+        return dtype
+
+    def compute(self, events, peaks):
+        split_tags = strax.split_by_containment(peaks, events)
+        result = np.zeros(len(events), self.dtype)
+        result['time'] = events['time']
+        result['endtime'] = events['endtime']
+        get_veto_tags(events, split_tags, result)
+
+        return result
+
+
+def get_veto_tags(events, split_tags, result):
+    """
+    Loops over events and tag main/alt S1/2 according to peak tag.
+
+    :param events: Event_info data type to be tagged.
+    :param split_tags: Tags split by events.
+    """
+    for tags, e, r in zip(split_tags, events, result):
+        r['n_tagged_peaks'] = np.sum(tags['veto_tag'] > 0)
+        for i in range(1, 3):
+            for peak_type in ['', 'alt_']:
+                if e[f'{peak_type}s{i}_index'] == -1:
+                    continue
+
+                index = e[f'{peak_type}s{i}_index']
+                r[f'{peak_type}s{i}_veto_tag'] = tags[index]['veto_tag']
+                r[f'{peak_type}s{i}_dt_veto'] = tags[index]['time_to_closest_veto']

straxen/plugins/peak_processing.py

@@ -324,3 +324,85 @@ def find_n_competing(peaks, windows, fraction):
             n_tot[i] = n_left[i] + np.sum(areas[i + 1:right_i] > threshold)
 
         return n_left, n_tot
+
+
+class PeakVetoTagging(strax.Plugin):
+    """
+    Plugin which tags S1 peaks according to  muon and neutron-vetos.
+    Tagging S2s is does not make sense as they occur with a delay.
+    However, we compute for both S1/S2 the time delay to the closest veto
+    region.
+
+        * untagged: 0
+        * neutron-veto: 1
+        * muon-veto: 2
+        * both vetos: 3
+    """
+    __version__ = '0.0.1'
+    depends_on = ('peak_basics', 'veto_regions_nv', 'veto_regions_mv')
+    provides = ('peak_veto_tags')
+    save_when = strax.SaveWhen.TARGET
+
+    dtype = strax.time_fields + \
+            [('veto_tag', np.int8,
+              'Veto tag for S1 peaks. unatagged: 0, nveto: 1, mveto: 2, both: 3'),
+             ('time_to_closest_veto', np.int64)]
+
+    def compute(self, peaks, veto_regions_nv, veto_regions_mv):
+        touching_mv = strax.touching_windows(peaks, veto_regions_mv)
+        touching_nv = strax.touching_windows(peaks, veto_regions_nv)
+
+        tags = np.zeros(len(peaks))
+        tags = tag_peaks(tags, touching_nv, 1)
+        tags = tag_peaks(tags, touching_mv, 2)
+
+        tags[peaks['type'] == 2] = 0
+
+        vetos = np.concatenate([veto_regions_nv, veto_regions_mv])
+        vetos = np.sort(vetos, order='time')
+        dt = get_time_to_closest_veto(peaks, vetos)
+
+        return {'time': peaks['time'],
+                'endtime': strax.endtime(peaks),
+                'veto_tag': tags,
+                'time_to_closest_veto': dt,
+                }
+
+
+@numba.njit(cache=True, nogil=True)
+def tag_peaks(tags, touching_windows, tag_number):
+    pre_tags = np.zeros(len(tags), dtype=np.int8)
+    for start, end in touching_windows:
+        pre_tags[start:end] = tag_number
+    tags += pre_tags
+    return tags
+
+
+def get_time_to_closest_veto(peaks, veto_intervals):
+    """
+    Computes time difference between peak and closest veto interval.
+    """
+    vetos = np.zeros(len(veto_intervals)+2, np.int64)
+    vetos[1:-1] = veto_intervals['time']
+    vetos[-1] = 9223372036854775807  # 64 bit infinity
+    return _get_time_to_closest_veto(peaks, vetos)
+
+
+@numba.njit(cache=True, nogil=True)
+def _get_time_to_closest_veto(peaks, vetos):
+    res = np.zeros(len(peaks), dtype=np.int64)
+    veto_index = 0
+    for ind, p in enumerate(peaks):
+        for veto_index in range(veto_index, len(vetos)):
+            dt_current_veto = np.abs(vetos[veto_index] - p['time'])
+            dt_next_veto = np.abs(vetos[veto_index+1] - p['time'])
+
+            # Next veto is closer so we have to repeat
+            if dt_current_veto >= dt_next_veto:
+                veto_index += 1
+                continue
+
+            res[ind] = dt_current_veto
+            break
+
+    return res