Fix peaklet area bias

straxen/plugins/peaklet_processing.py

@@ -92,7 +92,7 @@ class Peaklets(strax.Plugin):
     parallel = 'process'
     compressor = 'zstd'
 
-    __version__ = '0.3.10'
+    __version__ = '0.4.1'
 
     def infer_dtype(self):
         return dict(peaklets=strax.peak_dtype(
@@ -153,22 +153,49 @@ def compute(self, records, start, end):
         # Get hits outside peaklets, and store them separately.
         # fully_contained is OK provided gap_threshold > extension,
         # which is asserted inside strax.find_peaks.
-        lone_hits = hits[strax.fully_contained_in(hits, peaklets) == -1]
+        is_lone_hit = strax.fully_contained_in(hits, peaklets) == -1
+        lone_hits = hits[is_lone_hit]
         strax.integrate_lone_hits(
             lone_hits, records, peaklets,
             save_outside_hits=(self.config['peak_left_extension'],
                                self.config['peak_right_extension']),
             n_channels=len(self.to_pe))
 
         # Compute basic peak properties -- needed before natural breaks
-        strax.sum_waveform(peaklets, r, self.to_pe)
+        hits = hits[~is_lone_hit]
+        # Define regions outside of peaks such that _find_hit_integration_bounds
+        # is not extended beyond a peak.
+        outside_peaks = self.create_outside_peaks_region(peaklets, start, end)
+        strax.find_hit_integration_bounds(
+            hits, outside_peaks, records,
+            save_outside_hits=(self.config['peak_left_extension'],
+                               self.config['peak_right_extension']),
+            n_channels=len(self.to_pe),
+            allow_bounds_beyond_records=True,
+        )
+
+        # Transform hits to hitlets for naming conventions. A hit refers
+        # to the central part above threshold a hitlet to the entire signal
+        # including the left and right extension.
+        # (We are not going to use the actual hitlet data_type here.)
+        hitlets = hits
+        del hits
+
+        hitlet_time_shift = (hitlets['left'] - hitlets['left_integration']) * hitlets['dt']
+        hitlets['time'] = hitlets['time'] - hitlet_time_shift
+        hitlets['length'] = (hitlets['right_integration'] - hitlets['left_integration'])
+        hitlets = strax.sort_by_time(hitlets)
+        rlinks = strax.record_links(records)
+
+        strax.sum_waveform(peaklets, hitlets, r, rlinks, self.to_pe)
+
         strax.compute_widths(peaklets)
 
         # Split peaks using low-split natural breaks;
         # see https://github.com/XENONnT/straxen/pull/45
         # and https://github.com/AxFoundation/strax/pull/225
         peaklets = strax.split_peaks(
-            peaklets, r, self.to_pe,
+            peaklets, hitlets, r, rlinks, self.to_pe,
             algorithm='natural_breaks',
             threshold=self.natural_breaks_threshold,
             split_low=True,
@@ -187,7 +214,7 @@ def compute(self, records, start, end):
 
         if self.config['saturation_correction_on']:
             peak_list = peak_saturation_correction(
-                r, peaklets, self.to_pe,
+                r, rlinks, peaklets, hitlets, self.to_pe,
                 reference_length=self.config['saturation_reference_length'],
                 min_reference_length=self.config['saturation_min_reference_length'])
 
@@ -200,8 +227,10 @@ def compute(self, records, start, end):
         # (b) increase strax memory usage / max_messages,
         #     possibly due to its currently primitive scheduling.
         hit_max_times = np.sort(
-            hits['time']
-            + hits['dt'] * hit_max_sample(records, hits))
+            hitlets['time']
+            + hitlets['dt'] * hit_max_sample(records, hitlets)
+            + hitlet_time_shift  # add time shift again to get correct maximum
+        )
         peaklet_max_times = (
                 peaklets['time']
                 + np.argmax(peaklets['data'], axis=1) * peaklets['dt'])
@@ -213,12 +242,12 @@ def compute(self, records, start, end):
 
         if self.config['diagnose_sorting'] and len(r):
             assert np.diff(r['time']).min(initial=1) >= 0, "Records not sorted"
-            assert np.diff(hits['time']).min(initial=1) >= 0, "Hits not sorted"
+            assert np.diff(hitlets['time']).min(initial=1) >= 0, "Hits/Hitlets not sorted"
             assert np.all(peaklets['time'][1:]
                           >= strax.endtime(peaklets)[:-1]), "Peaks not disjoint"
 
         # Update nhits of peaklets:
-        counts = strax.touching_windows(hits, peaklets)
+        counts = strax.touching_windows(hitlets, peaklets)
         counts = np.diff(counts, axis=1).flatten()
         counts += 1
         peaklets['n_hits'] = counts
@@ -253,16 +282,43 @@ def clip_peaklet_times(peaklets, start, end):
             if strax.endtime(p) > end:
                 p['length'] = (end - p['time']) // p['dt']
 
+    @staticmethod
+    def create_outside_peaks_region(peaklets, start, end):
+        """
+        Creates time intervals which are outside peaks.
+
+        :param peaklets: Peaklets for which intervals should be computed.
+        :param start: Chunk start
+        :param end: Chunk end
+        :return: array of strax.time_fields dtype.
+        """
+        if not len(peaklets):
+            return np.zeros(0, dtype=strax.time_fields)
+
+        outside_peaks = np.zeros(len(peaklets) + 1,
+                                 dtype=strax.time_fields)
+
+        outside_peaks[0]['time'] = start
+        outside_peaks[0]['endtime'] = peaklets[0]['time']
+        outside_peaks[1:-1]['time'] = strax.endtime(peaklets[:-1])
+        outside_peaks[1:-1]['endtime'] = peaklets['time'][1:]
+        outside_peaks[-1]['time'] = strax.endtime(peaklets[-1])
+        outside_peaks[-1]['endtime'] = end
+        return outside_peaks
+
 
 @numba.jit(nopython=True, nogil=True, cache=True)
-def peak_saturation_correction(records, peaks, to_pe,
+def peak_saturation_correction(records, rlinks, peaks, hitlets, to_pe,
                                reference_length=100,
                                min_reference_length=20,
                                use_classification=False,
                                ):
     """Correct the area and per pmt area of peaks from saturation
     :param records: Records
+    :param rlinks: strax.record_links of corresponding records.
     :param peaks: Peaklets / Peaks
+    :param hitlets: Hitlets found in records to build peaks.
+        (Hitlets are hits including the left/right extension)
     :param to_pe: adc to PE conversion (length should equal number of PMTs)
     :param reference_length: Maximum number of reference sample used
     to correct saturated samples
@@ -334,7 +390,7 @@ def peak_saturation_correction(records, peaks, to_pe,
         peaks[peak_i]['length'] = p['length'] * p['dt'] / dt
         peaks[peak_i]['dt'] = dt
 
-    strax.sum_waveform(peaks, records, to_pe, peak_list)
+    strax.sum_waveform(peaks, hitlets, records, rlinks, to_pe, peak_list)
     return peak_list
 
 

straxen/plugins/veto_hitlets.py

@@ -126,7 +126,9 @@ def compute(self, records_nv, start, end):
                                               min_hitlet_sample=600)
 
         temp_hitlets = strax.split_peaks(temp_hitlets,
+                                         None,  # Only needed for peak splitting
                                          records_nv,
+                                         None,  # Only needed for peak splitting
                                          self.to_pe,
                                          data_type='hitlets',
                                          algorithm='local_minimum',

tests/test_peaklet_processing.py

@@ -0,0 +1,28 @@
+import numpy as np
+from hypothesis import given, settings
+import hypothesis.strategies as strat
+
+import strax
+import straxen
+
+
+@settings(deadline=None)
+@given(strat.lists(strat.integers(min_value=0, max_value=10),
+                   min_size=8, max_size=8, unique=True),
+       )
+def test_create_outside_peaks_region(time):
+    time = np.sort(time)
+    time_intervals = np.zeros(len(time)//2, strax.time_dt_fields)
+    time_intervals['time'] = time[::2]
+    time_intervals['length'] = time[1::2] - time[::2]
+    time_intervals['dt'] = 1
+
+    st = straxen.contexts.demo()
+    p = st.get_single_plugin('0', 'peaklets')
+    outside = p.create_outside_peaks_region(time_intervals, 0, np.max(time))
+
+    touching = strax.touching_windows(outside, time_intervals, window=0)
+
+    for tw in touching:
+        print(tw)
+        assert np.diff(tw) == 0, 'Intervals overlap although they should not!'

tests/test_several.py

@@ -40,7 +40,7 @@ def test_secret():
 # They were added on 27/11/2020 and may be outdated by now
 EXPECTED_OUTCOMES_TEST_SEVERAL = {
     'n_peaks': 128,
-    'n_s1': 6,
+    'n_s1': 8,
     'run_live_time': 0.17933107,
     'n_events': 2,
 }