Implement Pyrrha functionalities in Irene

invisible_cities/calib/calib_sensors_functions.py

@@ -59,8 +59,9 @@ def modes  (wfs): return to_col_vector(mode  (wfs, axis=1))
 
 def subtract_baseline(wfs, *, bls_mode=BlsMode.mean):
     """
-    Subtract the baseline to all waveforms in the input
-    with a specific algorithm.
+    Subtract the baseline from all waveforms in the input
+    with a specific algorithm, leaving individual samples
+    that are already 0 untouched.
 
     Parameters
     ----------
@@ -75,16 +76,19 @@ def subtract_baseline(wfs, *, bls_mode=BlsMode.mean):
     Returns
     -------
     bls: np.ndarray with shape (n, m)
-        Baseline-subtracted waveforms.
+        Baseline-subtracted waveforms, with originally-zero
+        samples left as 0.
     """
 
-    if   bls_mode is BlsMode.mean     : return wfs - means     (wfs)
-    elif bls_mode is BlsMode.median   : return wfs - medians   (wfs)
-    elif bls_mode is BlsMode.mode     : return wfs - modes     (wfs)
-    elif bls_mode is BlsMode.scipymode: return wfs - scipy_mode(wfs, axis=1)
+    if   bls_mode is BlsMode.mean     : baseline = means     (wfs)
+    elif bls_mode is BlsMode.median   : baseline = medians   (wfs)
+    elif bls_mode is BlsMode.mode     : baseline = modes     (wfs)
+    elif bls_mode is BlsMode.scipymode: baseline = scipy_mode(wfs, axis=1)
     else:
         raise TypeError(f"Unrecognized baseline subtraction option: {bls_mode}")
 
+    return np.where(wfs == 0, 0, wfs - baseline) # subtracts the baseline to non-zero samples only
+
 
 def calibrate_wfs(wfs, adc_to_pes):
     """
@@ -136,15 +140,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  )
@@ -159,10 +161,10 @@ def sipm_subtract_median_and_calibrate(sipm_wfs, adc_to_pes): return calibrate_w
 
 # Dict of functions for SiPM processing
 sipm_processing = {
-    SiPMCalibMode.subtract_mode            :      subtract_mode                ,# For gain extraction
-    SiPMCalibMode.subtract_median          :      subtract_median              ,# For gain extraction
-    SiPMCalibMode.subtract_mode_calibrate  : sipm_subtract_mode_and_calibrate  ,# For PDF calculation
-    SiPMCalibMode.subtract_mean_calibrate  : sipm_subtract_mean_and_calibrate  ,# For PDF calculation
-    SiPMCalibMode.subtract_median_calibrate: sipm_subtract_median_and_calibrate,# For PDF calculation
-    SiPMCalibMode.subtract_mode_zs         : calibrate_sipms                    # For data processing
+    SiPMCalibMode.subtract_mode              :      subtract_mode                ,# For gain extraction
+    SiPMCalibMode.subtract_median            :      subtract_median              ,# For gain extraction
+    SiPMCalibMode.subtract_mode_calibrate    : sipm_subtract_mode_and_calibrate  ,# For PDF calculation
+    SiPMCalibMode.subtract_mean_calibrate    : sipm_subtract_mean_and_calibrate  ,# For PDF calculation
+    SiPMCalibMode.subtract_median_calibrate  : sipm_subtract_median_and_calibrate,# For PDF calculation
+    SiPMCalibMode.subtract_baseline_calibrate: calibrate_sipms                    # For data processing
 }

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

@@ -84,6 +84,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
@@ -762,7 +764,8 @@ 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, 
+               sipm_selection_algo):
     s1_params = dict(time        = minmax(min = s1_tmin,
                                           max = s1_tmax),
                     length       = minmax(min = s1_lmin,
@@ -783,8 +786,8 @@ def build_pmap(detector_db, run_number, pmt_samp_wid, sipm_samp_wid,
     sipm_ids = np.argwhere(datasipm.Active.values==1).flatten()
     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, sipm_ids,
-                            pmt_samp_wid, sipm_samp_wid)
+                            s1_params, s2_params, pmt_ids, sipm_ids,
+                            pmt_samp_wid, sipm_samp_wid, sipm_selection_algo)
 
     return build_pmap
 
@@ -802,20 +805,85 @@ 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)
     DataSiPM   = DataSiPM.loc[lambda df: df.Active==1]
     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 select_cutting_algorithm(algo, **cutting_params):
+    if algo is CutAlgo.threshold:
+        return threshold_sipm_selection(**cutting_params)
+    elif algo is CutAlgo.pyrrha:
+        return pyrrha_sipm_selection(**cutting_params)
+    elif algo is CutAlgo.no_cut:
+        return no_cut_sipm_selection()
+    else:
+        raise ValueError(f"Unsupported cutting algorithm: {algo!r}. Expected one of {list(CutAlgo)}")
+
+
+def threshold_sipm_selection( thr_sipm_type
+                            , thr_sipm
+                            , thr_sipm_s2
+                            , run_number
+                            , detector_db = None):
+    '''
+    Applies energy thresholds to SiPM S2 waveforms:
+    - thr_sipm:    applied per time bin. Waveform samples below this 
+                   threshold are set to zero.  
+    - thr_sipm_s2: applied to the integrated waveform charge. If the total
+                   charge is below this threshold, it is set to zero.
+    '''
+    if detector_db is None:
+        sipm_thr = thr_sipm
+    else:
+        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):
+    '''
+    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 no_cut_sipm_selection():
+    """"
+    Function that applies no cuts to the SiPM waveforms. 
+    """
+    def no_cut_sipm_selection(wfs):
+        sipm_ids = np.arange(wfs.shape[0])
+        return sipm_ids, wfs
+    return no_cut_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)
@@ -1231,8 +1299,8 @@ def integrate_wfs(wfs):
 # Compound components
 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):
+                  s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin,
+                  h5out, sipm_selection_algo, sipm_rwf_to_cal=None):
 
     # Filter events without signal over threshold
     indices_pass    = fl.map(check_nonempty_indices,
@@ -1243,7 +1311,8 @@ 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, 
+                                         sipm_selection_algo),
                               args = ("ccwfs", "s1_indices", "s2_indices", "sipm"),
                               out  = "pmap")
 

invisible_cities/cities/hypathia.py

@@ -28,6 +28,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
@@ -48,36 +49,38 @@
 from .  components import calibrate_sipms
 from .  components import get_actual_sipm_thr
 from .  components import sensor_masker
+from .  components import select_cutting_algorithm
+
+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
+            , 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
+            , 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)
 
@@ -115,9 +118,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
+    sipm_selection_algo = select_cutting_algorithm(cutting_function, **cutting_params)
+
+
     event_count_in  = fl.spy_count()
     event_count_out = fl.spy_count()
 
@@ -136,8 +143,8 @@ def hypathia( files_in        : OneOrManyFiles
         compute_pmaps, empty_indices, empty_pmaps = 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)
+                                             s2_lmax, s2_lmin, s2_rebin_stride, s2_stride, s2_tmax, s2_tmin,
+                                             h5out, sipm_selection_algo, sipm_rwf_to_cal)
 
         result = push(source = wf_from_files(files_in, WfType.mcrd),
                       pipe   = pipe(fl.slice(*event_range, close_all=True),