include all peak code in peak_som

amstrax/plugins/peaks/peaks_som.py

@@ -10,7 +10,43 @@
 
 export, __all__ = strax.exporter()
 
+# These are also needed in peaklets, since hitfinding is repeated
+HITFINDER_OPTIONS = tuple([
+    strax.Option(
+        'hit_min_amplitude',
+        default='pmt_commissioning_initial',
+        help='Minimum hit amplitude in ADC counts above baseline. '
+             'See straxen.hit_min_amplitude for options.'
+    )])
 
+
+@export
+@strax.takes_config(
+    strax.Option('peak_gap_threshold', default=300,
+                 help="No hits for this many ns triggers a new peak"),
+    strax.Option('peak_left_extension', default=10,
+                 help="Include this many ns left of hits in peaks"),
+    strax.Option('peak_right_extension', default=10,
+                 help="Include this many ns right of hits in peaks"),
+    strax.Option('peak_min_area', default=10,
+                 help="Minimum contributing PMTs needed to define a peak"),
+    strax.Option('peak_min_pmts', default=1,
+                 help="Minimum contributing PMTs needed to define a peak"),
+    strax.Option('single_channel_peaks', default=False,
+                 help='Whether single-channel peaks should be reported'),
+    strax.Option('peak_split_min_height', default=25,
+                 help="Minimum height in PE above a local sum waveform"
+                      "minimum, on either side, to trigger a split"),
+    strax.Option('peak_split_min_ratio', default=4,
+                 help="Minimum ratio between local sum waveform"
+                      "minimum and maxima on either side, to trigger a split"),
+    strax.Option('diagnose_sorting', track=False, default=False,
+                 help="Enable runtime checks for sorting and disjointness"),
+    strax.Option('n_tpc_pmts', track=False, default=False,
+                 help="Number of channels"),
+    strax.Option('gain_to_pe_array', default=None,
+                 help="Gain to pe array"),
+)
 @export
 class PeaksSOM(Peaks):
     """
@@ -29,12 +65,11 @@ class PeaksSOM(Peaks):
     """
 
     __version__ = "0.2.0"
-    child_plugin = True
+    depends_on = ('records',)
     data_kind = 'peaks'
     parallel = 'process'
     provides = ('peaks')
-    #depends_on = ("records",)
-    #provides = "peaks"
+    rechunk_on_save = True
 
     # This is not a great solution but needed for now
     directory = os.path.dirname(__file__)  # Get the directory where the script is located
@@ -52,7 +87,8 @@ def infer_dtype(self):
             ("loc_x_som", np.int16, "x location of the peak(let) in the SOM"),
             ("loc_y_som", np.int16, "y location of the peak(let) in the SOM"),
         ]
-        return dtype
+        #return strax.peak_dtype(n_channels=self.config['n_tpc_pmts'])
+        return dtype, strax.peak_dtype(n_channels=self.config['n_tpc_pmts'])
 
     def setup(self):
         self.som_weight_cube = self.som_files["weight_cube"]
@@ -65,8 +101,41 @@ def setup(self):
 
     def compute(self, records, start, end):
         # Current classification
-        peaks_classifcation = super().compute(records, start, end)
+        r = records
+
+        if self.config['gain_to_pe_array'] is None:
+            self.to_pe = np.ones(self.config['n_tpc_pmts'])
+        else:
+            self.to_pe = self.config['gain_to_pe_array']
+
+        hits = strax.find_hits(r)
+        hits = strax.sort_by_time(hits)
+
+        rlinks = strax.record_links(r)
+
+        # Rewrite to just peaks/hits
+        peaks = strax.find_peaks(
+            hits, self.to_pe,
+            gap_threshold=self.config['peak_gap_threshold'],
+            left_extension=self.config['peak_left_extension'],
+            right_extension=self.config['peak_right_extension'],
+            min_area=self.config['peak_min_area'],
+            min_channels=self.config['peak_min_pmts'],
+            #             min_channels=1,
+            result_dtype=strax.peak_dtype(n_channels=self.config['n_tpc_pmts'])
+            #             result_dtype=self.dtype
+        )
+
+        strax.sum_waveform(peaks, hits, r, rlinks, self.to_pe)
+
+        peaks = strax.split_peaks(
+            peaks, hits, r, rlinks, self.to_pe,
+            min_height=self.config['peak_split_min_height'],
+            min_ratio=self.config['peak_split_min_ratio'])
+
+        strax.compute_widths(peaks)
 
+        peaks_classifcation = peaks
         peaks_with_som = np.zeros(len(peaks_classifcation), dtype=self.dtype)
         strax.copy_to_buffer(peaks_classifcation, peaks_with_som, "_copy_peaklets_information")
         peaks_with_som["straxen_type"] = peaks_classifcation["type"]