Implement Pyrrha functionalities in Irene

invisible_cities/calib/calib_sensors_functions.py

@@ -136,15 +136,13 @@ def pmt_subtract_maw(cwfs, n_maw=100):
     return cwfs - maw
 
 
-def calibrate_sipms(sipm_wfs, adc_to_pes, thr, *, bls_mode=BlsMode.mode):
+def calibrate_sipms(sipm_wfs, adc_to_pes, *, bls_mode=BlsMode.mode):
     """
-    Subtracts the baseline, calibrates waveforms to pes
-    and suppresses values below `thr` (in pes).
+    Subtract baseline, and calibrates waveforms to pes.
     """
-    thr  = to_col_vector(np.full(sipm_wfs.shape[0], thr))
     bls  = subtract_baseline(sipm_wfs, bls_mode=bls_mode)
     cwfs = calibrate_wfs(bls, adc_to_pes)
-    return np.where(cwfs > thr, cwfs, 0)
+    return cwfs
 
 
 def subtract_mean  (wfs): return subtract_baseline(wfs, bls_mode=BlsMode.mean  )

invisible_cities/calib/calib_sensors_functions_test.py

@@ -13,6 +13,7 @@
 
 from .          import calib_sensors_functions as csf
 from .. core    import core_functions          as cf
+from .. reco    import wfm_functions           as wfm
 from .. sierpe  import fee                     as FE
 from .. sierpe  import waveform_generator      as wfg
 
@@ -207,19 +208,20 @@ def test_calibrate_sipms_stat(oscillating_waveform_with_baseline,
      baseline )              = oscillating_waveform_with_baseline
     #n_maw                    = n_samples // 500
 
-    ccwfs = csf.calibrate_sipms(wfs, adc_to_pes, nsigma * noise_sigma, bls_mode=BlsMode.mode)
+    ccwfs = csf.calibrate_sipms(wfs, adc_to_pes, bls_mode=BlsMode.mode)
+    zeroed_ccwfs = wfm.zero_wfs_below_threshold(ccwfs, zeroing_thr=nsigma * noise_sigma)
 
-    number_of_zeros = np.count_nonzero(ccwfs == 0)
-    assert number_of_zeros > fraction * ccwfs.size
+    number_of_zeros = np.count_nonzero(zeroed_ccwfs == 0)
+    assert number_of_zeros > fraction * zeroed_ccwfs.size
 
 
 def test_calibrate_sipms_common_threshold(toy_sipm_signal):
     (signal_adc, adc_to_pes,
      signal_zs_common_threshold, _,
      common_threshold, _) = toy_sipm_signal
 
-    zs_wf = csf.calibrate_sipms(signal_adc, adc_to_pes,
-                                common_threshold, bls_mode=BlsMode.mode)
+    ccwf = csf.calibrate_sipms(signal_adc, adc_to_pes, bls_mode=BlsMode.mode)
+    zs_wf = wfm.zero_wfs_below_threshold(ccwf, zeroing_thr=common_threshold)
 
     for actual, expected in zip(zs_wf, signal_zs_common_threshold):
         assert actual == approx(expected)
@@ -230,10 +232,9 @@ def test_calibrate_sipms_individual_thresholds(toy_sipm_signal):
      _, signal_zs_individual_thresholds,
      _, individual_thresholds) = toy_sipm_signal
 
+    ccwf = csf.calibrate_sipms(signal_adc, adc_to_pes, bls_mode=BlsMode.mode)
+    zs_wf = wfm.zero_wfs_below_threshold(ccwf, zeroing_thr=individual_thresholds)
 
-    zs_wf = csf.calibrate_sipms(signal_adc, adc_to_pes,
-                                individual_thresholds,
-                                bls_mode=BlsMode.mode)
     for actual, expected in zip(zs_wf, signal_zs_individual_thresholds):
         assert actual == approx(expected)
 
@@ -309,11 +310,12 @@ def test_area_of_sum_equals_sum_of_areas_pmts(square_pmt_and_sipm_waveforms):
 
 def test_area_of_sum_equals_sum_of_areas_sipms(square_pmt_and_sipm_waveforms):
     _, nsensors, _, _, _, sipms_wfm, _ = square_pmt_and_sipm_waveforms
-    adc_to_pes = np.full(nsensors, 100, dtype=float)
-    cwfs       = csf.calibrate_sipms(sipms_wfm, adc_to_pes, thr=10, bls_mode=BlsMode.mode)
-    stot       = np.sum(cwfs[0]) * nsensors
-    sums       = np.sum(cwfs, axis=1)
-    stot2      = reduce(add, sums)
+    adc_to_pes  = np.full(nsensors, 100, dtype=float)
+    cwfs        = csf.calibrate_sipms(sipms_wfm, adc_to_pes, bls_mode=BlsMode.mode)
+    zeroed_cwfs = wfm.zero_wfs_below_threshold(cwfs, zeroing_thr=10)
+    stot        = np.sum(zeroed_cwfs[0]) * nsensors
+    sums        = np.sum(zeroed_cwfs, axis=1)
+    stot2       = reduce(add, sums)
     assert stot == approx(stot2, rel=1e-3)
 
 

invisible_cities/cities/components.py

@@ -83,6 +83,8 @@
 from .. types  .ic_types          import        types_dict_summary
 from .. types  .ic_types          import         types_dict_tracks
 from .. types  .symbols           import                    WfType
+from .. types  .symbols           import                   CutAlgo
+from .. types  .symbols           import       SiPMSelectionMethod
 from .. types  .symbols           import               RebinMethod
 from .. types  .symbols           import                SiPMCharge
 from .. types  .symbols           import                   BlsMode
@@ -746,7 +748,7 @@ def sensor_data(path, wf_type):
 
 def build_pmap(detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
                s1_lmax, s1_lmin, s1_rebin_stride, s1_stride, s1_tmax, s1_tmin,
-               s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, thr_sipm_s2):
+               s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, apply_cut):
     s1_params = dict(time        = minmax(min = s1_tmin,
                                           max = s1_tmax),
                     length       = minmax(min = s1_lmin,
@@ -766,8 +768,8 @@ def build_pmap(detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
 
     def build_pmap(ccwf, s1_indx, s2_indx, sipmzs): # -> PMap
         return pkf.get_pmap(ccwf, s1_indx, s2_indx, sipmzs,
-                            s1_params, s2_params, thr_sipm_s2, pmt_ids,
-                            pmt_samp_wid, sipm_samp_wid)
+                            s1_params, s2_params, pmt_ids,
+                            pmt_samp_wid, sipm_samp_wid, apply_cut)
 
     return build_pmap
 
@@ -785,19 +787,77 @@ def calibrate_pmts(cwf):# -> CCwfs:
     return calibrate_pmts
 
 
-def calibrate_sipms(dbfile, run_number, thr_sipm):
+def calibrate_sipms(dbfile, run_number):
     DataSiPM   = load_db.DataSiPM(dbfile, run_number)
     adc_to_pes = np.abs(DataSiPM.adc_to_pes.values)
 
     def calibrate_sipms(rwf):
         return csf.calibrate_sipms(rwf,
                                    adc_to_pes = adc_to_pes,
-                                   thr        = thr_sipm,
                                    bls_mode   = BlsMode.mode)
 
     return calibrate_sipms
 
 
+def apply_cutting_function(algo, **cutting_params):
+
+    if algo is CutAlgo.threshold:
+        func = threshold_sipm_selection(**cutting_params)
+    elif algo is CutAlgo.pyrrha:
+        func = pyrrha_sipm_selection(**cutting_params)
+    else:
+        raise ValueError(f"Unsupported cutting algorithm: {algo!r}. Expected one of {list(CutAlgo)}")
+
+    def apply_cutting_function(wfs):
+        return func(wfs)
+
+    return apply_cutting_function
+
+
+def threshold_sipm_selection(thr_sipm_type
+                            , thr_sipm
+                            , thr_sipm_s2
+                            , run_number
+                            , detector_db = None):
+    '''
+    Function that applies thresholding to the sipms in standard irene manner,
+    by zeroing all waveform values below a threshold.
+    '''
+    # assume that if the detector_db is None, you return sipm threshold as the number provided
+    if detector_db is None:
+        sipm_thr = thr_sipm
+    else:
+        # extract sipm threshold
+        sipm_thr = get_actual_sipm_thr(thr_sipm_type, thr_sipm, detector_db, run_number)
+
+    def threshold_sipm_selection(wfs):
+        return wfm.charge_threshold_method(wfs, zeroing_thr = sipm_thr, integration_thr=thr_sipm_s2)
+
+    return threshold_sipm_selection
+
+
+def pyrrha_sipm_selection(selection_method     : SiPMSelectionMethod
+                         , selection_kwargs    : dict
+                         , proximity_threshold : float
+                         , padding_radius      : float
+                         , run_number          : int
+                         , detector_db         : str):
+    '''
+    Function that applies a generic selection function to the sipms, which can be used to 
+    implement a spatial SiPM selection method (called Pyrrha).
+    '''
+    def pyrrha_sipm_selection(wfs):
+        return wfm.spatial_selection_method(wfs, 
+                                            selection_method, 
+                                            selection_kwargs, 
+                                            proximity_threshold,
+                                            padding_radius,
+                                            run_number,
+                                            detector_db)
+
+    return pyrrha_sipm_selection
+
+
 def calibrate_with_mean(dbfile, run_number):
     DataSiPM   = load_db.DataSiPM(dbfile, run_number)
     adc_to_pes = np.abs(DataSiPM.adc_to_pes.values)
@@ -1214,7 +1274,7 @@ def integrate_wfs(wfs):
 def compute_and_write_pmaps(detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
                   s1_lmax, s1_lmin, s1_rebin_stride, s1_stride, s1_tmax, s1_tmin,
                   s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, thr_sipm_s2,
-                  h5out, sipm_rwf_to_cal=None):
+                  h5out, apply_cut, sipm_rwf_to_cal=None):
 
     # Filter events without signal over threshold
     indices_pass    = fl.map(check_nonempty_indices,
@@ -1225,7 +1285,7 @@ def compute_and_write_pmaps(detector_db, run_number, pmt_samp_wid, sipm_samp_wid
     # Build the PMap
     compute_pmap     = fl.map(build_pmap(detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
                                          s1_lmax, s1_lmin, s1_rebin_stride, s1_stride, s1_tmax, s1_tmin,
-                                         s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, thr_sipm_s2),
+                                         s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, apply_cut),
                               args = ("ccwfs", "s1_indices", "s2_indices", "sipm"),
                               out  = "pmap")
 

invisible_cities/cities/hypathia.py

@@ -30,6 +30,7 @@
 from .. io   .run_and_event_io import run_and_event_writer
 from .. io   .      trigger_io import       trigger_writer
 from .. types.symbols          import WfType
+from .. types.symbols          import CutAlgo
 from .. types.symbols          import SiPMThreshold
 
 from .. dataflow            import dataflow as fl
@@ -49,36 +50,40 @@
 from .  components import simulate_sipm_response
 from .  components import calibrate_sipms
 from .  components import get_actual_sipm_thr
+from .  components import apply_cutting_function
+
+from typing import Dict
+from typing import Any
 
 
 @city
-def hypathia( files_in        : OneOrManyFiles
-            , file_out        : str
-            , compression     : str
-            , event_range     : EventRangeType
-            , print_mod       : int
-            , detector_db     : str
-            , run_number      : int
-            , sipm_noise_cut  : float
-            , filter_padding  : int
-            , thr_sipm        : float
-            , thr_sipm_type   : SiPMThreshold
-            , pmt_wfs_rebin   : int
-            , pmt_pe_rms      : float
-            , s1_lmin         : int  , s1_lmax     : int
-            , s1_tmin         : float, s1_tmax     : float
-            , s1_rebin_stride : int  , s1_stride   : int
-            , thr_csum_s1     : float
-            , s2_lmin         : int  , s2_lmax     : int
-            , s2_tmin         : float, s2_tmax     : float
-            , s2_rebin_stride : int  , s2_stride   : int
-            , thr_csum_s2     : float, thr_sipm_s2 : float
-            , pmt_samp_wid    : float
-            , sipm_samp_wid   : float
+def hypathia( files_in         : OneOrManyFiles
+            , file_out         : str
+            , compression      : str
+            , event_range      : EventRangeType
+            , print_mod        : int
+            , detector_db      : str
+            , run_number       : int
+            , sipm_noise_cut   : float
+            , filter_padding   : int
+            , thr_sipm         : float
+            , thr_sipm_type    : SiPMThreshold
+            , pmt_wfs_rebin    : int
+            , pmt_pe_rms       : float
+            , s1_lmin          : int  , s1_lmax     : int
+            , s1_tmin          : float, s1_tmax     : float
+            , s1_rebin_stride  : int  , s1_stride   : int
+            , thr_csum_s1      : float
+            , s2_lmin          : int  , s2_lmax     : int
+            , s2_tmin          : float, s2_tmax     : float
+            , s2_rebin_stride  : int  , s2_stride   : int
+            , thr_csum_s2      : float, thr_sipm_s2 : float
+            , pmt_samp_wid     : float
+            , sipm_samp_wid    : float
+            , cutting_function : CutAlgo
+            , cutting_params   : Dict[str, Any]
             ):
 
-    sipm_thr = get_actual_sipm_thr(thr_sipm_type, thr_sipm, detector_db, run_number)
-
     #### Define data transformations
     sd = sensor_data(files_in[0], WfType.mcrd)
 
@@ -112,9 +117,13 @@ def hypathia( files_in        : OneOrManyFiles
                               item="sipm")
 
     # SiPMs calibration
-    sipm_rwf_to_cal  = fl.map(calibrate_sipms(detector_db, run_number, sipm_thr),
+    sipm_rwf_to_cal  = fl.map(calibrate_sipms(detector_db, run_number),
                               item = "sipm")
 
+    # apply function depending on user input, from provided list of functions
+    apply_cut        = apply_cutting_function(cutting_function, **cutting_params)
+
+
     event_count_in  = fl.spy_count()
     event_count_out = fl.spy_count()
 
@@ -134,7 +143,7 @@ def hypathia( files_in        : OneOrManyFiles
                                              detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
                                              s1_lmax, s1_lmin, s1_rebin_stride, s1_stride, s1_tmax, s1_tmin,
                                              s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin, thr_sipm_s2,
-                                             h5out, sipm_rwf_to_cal)
+                                             h5out, apply_cut, sipm_rwf_to_cal)
 
         result = push(source = wf_from_files(files_in, WfType.mcrd),
                       pipe   = pipe(fl.slice(*event_range, close_all=True),