Updated construct_trials and generalized extract_basic_trig_propertie…

…s to extract_trig_properties

pycbc/results/pygrb_postprocessing_utils.py

@@ -460,38 +460,37 @@ def sort_trigs(trial_dict, trigs, slide_dict, seg_dict):
 # =============================================================================
 # Extract trigger properties and store them as dictionaries
 # =============================================================================
-def extract_basic_trig_properties(trial_dict, trigs, slide_dict, seg_dict,
-                                  opts):
-    """Extract and store as dictionaries time, SNR, and BestNR of
-    time-slid triggers"""
+def extract_trig_properties(trial_dict, trigs, slide_dict, seg_dict, keys):
+    """Extract and store as dictionaries specific keys of time-slid
+    triggers (trigs) compatibly with the trials dictionary (trial_dict)"""
 
     # Sort the triggers into each slide
     sorted_trigs = sort_trigs(trial_dict, trigs, slide_dict, seg_dict)
-    logger.info("Triggers sorted.")
+    n_surviving_trigs = sum([len(sorted_trigs[slide_id])
+                             for slide_id in sorted_trigs.keys()])
+    msg = f"{n_surviving_trigs} triggers found within the trials dictionary "
+    msg += "and sorted."
+    logging.info(msg)
+
+    found_trigs = {}
+    for key in keys:
+        found_trigs[key] = {}
 
-    # Build the 3 dictionaries
-    trig_time = {}
-    trig_snr = {}
-    trig_bestnr = {}
     for slide_id in slide_dict:
         slide_trigs = sorted_trigs[slide_id]
         indices = numpy.nonzero(
             numpy.isin(trigs['network/event_id'], slide_trigs))[0]
-        if slide_trigs:
-            trig_time[slide_id] = trigs['network/end_time_gc'][
-                indices]
-            trig_snr[slide_id] = trigs['network/coherent_snr'][
-                indices]
-        else:
-            trig_time[slide_id] = numpy.asarray([])
-            trig_snr[slide_id] = numpy.asarray([])
-        trig_bestnr[slide_id] = reweightedsnr_cut(
-            trigs['network/reweighted_snr'][indices],
-            opts.newsnr_threshold)
+        for key in keys:
+            if slide_trigs:
+                found_trigs[key][slide_id] = get_coinc_snr(trigs)[indices] \
+                    if key == 'network/coincident_snr' else trigs[key][indices]
+            else:
+                found_trigs[key][slide_id] = numpy.asarray([])
+
+    logging.info("Triggers information extracted.")
 
-    logger.info("Time, SNR, and BestNR of triggers extracted.")
+    return found_trigs
 
-    return trig_time, trig_snr, trig_bestnr
 
 
 # =============================================================================
@@ -582,9 +581,11 @@ def load_segment_dict(hdf_file_path):
 # =============================================================================
 # Construct the trials from the timeslides, segments, and vetoes
 # =============================================================================
-def construct_trials(seg_files, seg_dict, ifos, slide_dict, vetoes):
+def construct_trials(seg_files, seg_dict, ifos, slide_dict, veto_file,
+                     show_onsource=False):
     """Constructs trials from triggers, timeslides, segments and vetoes"""
 
+    logging.info("Constructing trials.")
     trial_dict = {}
 
     # Get segments
@@ -593,19 +594,23 @@ def construct_trials(seg_files, seg_dict, ifos, slide_dict, vetoes):
     # Separate segments
     trial_time = abs(segs['on'])
 
+    # Determine the veto segments
+    vetoes = _extract_vetoes(veto_file, ifos, segs['off'])
+
+    # Slide vetoes over trials: this can only *reduce* the analysis time
     for slide_id in slide_dict:
-        # These can only *reduce* the analysis time
         curr_seg_list = seg_dict[slide_id]
 
-        # Construct the buffer segment list
+        # Fill in the buffer segment list if the onsource data must be hidden
         seg_buffer = segments.segmentlist()
-        for ifo in ifos:
-            slide_offset = slide_dict[slide_id][ifo]
-            seg_buffer.append(segments.segment(segs['buffer'][0] -
-                                               slide_offset,
-                                               segs['buffer'][1] -
-                                               slide_offset))
-        seg_buffer.coalesce()
+        if show_onsource:
+            for ifo in ifos:
+                slide_offset = slide_dict[slide_id][ifo]
+                seg_buffer.append(segments.segment(segs['buffer'][0] -
+                                                   slide_offset,
+                                                   segs['buffer'][1] -
+                                                   slide_offset))
+            seg_buffer.coalesce()
 
         # Construct the ifo-indexed dictionary of slid veteoes
         slid_vetoes = _slide_vetoes(vetoes, slide_dict, slide_id, ifos)