Handling vetoes and timeslides in pygrb mini followups

bin/pygrb/pycbc_pygrb_minifollowups

@@ -56,11 +56,23 @@ def add_wiki_row(outfile, cols):
 
 
 def make_timeseries_plot(workflow, trig_file, snr_type, central_time,
-                         shift_time, out_dir, ifo=None, tags=None):
+                         shift_time, out_dir, ifo=None, seg_files=None,
+                         veto_file=None, slide_id=0, tags=None):
     """Adds a node for a timeseries of PyGRB results to the workflow"""
 
     tags = [] if tags is None else tags
 
+    # If necessary, overwrite the slide-id configuration option to ensure only
+    # the correct timeslide is used in these follow-up plots
+    orig_slide_id = None
+    if workflow.cp.has_option('pygrb_plot_snr_timeseries','slide-id'):
+        orig_slide_id = workflow.cp.get('pygrb_plot_snr_timeseries','slide-id')
+        workflow.cp.add_options_to_section(
+            'pygrb_plot_snr_timeseries',
+            [('slide-id', str(slide_id))],
+            True
+        )
+
     # Initialize job node with its tags
     grb_name = workflow.cp.get('workflow', 'trigger-name')
     extra_tags = ['GRB'+grb_name]
@@ -71,6 +83,14 @@ def make_timeseries_plot(workflow, trig_file, snr_type, central_time,
                           ifos=workflow.ifos, out_dir=out_dir,
                           tags=tags+extra_tags).create_node()
     node.add_input_opt('--trig-file', trig_file)
+    # Include the onsource trial if this is a follow up on the onsource
+    if 'loudest_onsource_event' in tags:
+        node.add_opt('--onsource')
+    # Pass the segments files and veto file
+    if seg_files:
+        node.add_input_list_opt('--seg-files', seg_files)
+    if veto_file:
+        node.add_input_opt('--veto-file', veto_file)
     node.new_output_file_opt(workflow.analysis_time, '.png',
                              '--output-file', tags=extra_tags)
     # Quantity to be displayed on the y-axis of the plot
@@ -93,6 +113,14 @@ def make_timeseries_plot(workflow, trig_file, snr_type, central_time,
     # Add job node to workflow
     workflow += node
 
+    # Revert back the original slide-id option if necessary
+    if orig_slide_id:
+        workflow.cp.add_options_to_section(
+            'pygrb_plot_snr_timeseries',
+            [('slide-id', orig_slide_id)],
+            True
+        )
+
     return node.output_files
 
 
@@ -107,6 +135,11 @@ parser.add_argument('--followups-file',
                     help="HDF file with the triggers/injections to follow up")
 parser.add_argument('--wiki-file',
                     help="Name of file to save wiki-formatted table in")
+parser.add_argument("-a", "--seg-files", nargs="+", action="store",
+                    default=[], help="The location of the buffer, " +
+                    "onsource and offsource txt segment files.")
+parser.add_argument("-V", "--veto-file", action="store",
+                    help="The location of the xml veto file.")
 wf.add_workflow_command_line_group(parser)
 wf.add_workflow_settings_cli(parser, include_subdax_opts=True)
 ppu.pygrb_add_bestnr_cut_opt(parser)
@@ -142,6 +175,19 @@ trig_file = resolve_url_to_file(os.path.abspath(args.trig_file))
 ifos = ppu.extract_ifos(os.path.abspath(args.trig_file))
 num_ifos = len(ifos)
 
+# File instance of the veto file
+veto_file = args.veto_file
+start_rundir = os.getcwd()
+if veto_file:
+    veto_file = os.path.join(start_rundir, args.veto_file)
+    veto_file = wf.resolve_url_to_file(veto_file)
+
+# Convert the segments files to a FileList
+seg_files = wf.FileList([
+    wf.resolve_url_to_file(os.path.join(start_rundir, f))
+    for f in args.seg_files
+])
+
 # (Loudest) off/on-source events are on time-slid data so the
 # try will succeed, as it finds the time shift columns.
 is_injection_followup = True
@@ -151,7 +197,6 @@ try:
 except:
     pass
 
-
 # Loop over triggers/injections to be followed up
 for num_event in range(num_events):
     files = FileList([])
@@ -163,29 +208,41 @@ for num_event in range(num_events):
         for key in fp.keys():
             row.append(fp[key][num_event])
         add_wiki_row(wiki_file, row)
+    # Handle injections (which are on unslid data)
+    if is_injection_followup:
+        for snr_type in ['reweighted', 'coherent']:
+            files += make_timeseries_plot(workflow, trig_file,
+                                          snr_type, gps_time, gps_time,
+                                          args.output_dir, ifo=None,
+                                          seg_files=seg_files,
+                                          veto_file=veto_file,
+                                          slide_id=0,
+                                          tags=args.tags + [str(num_event)])
+        for ifo in ifos:
+            files += mini.make_qscan_plot(workflow, ifo, gps_time,
+                                          args.output_dir,
+                                          tags=args.tags + [str(num_event)])
     # Handle off/on-source loudest triggers follow-up (which are on slid data)
-    if not is_injection_followup:
+    else:
+        # If the file has information about slides, this is the follow up of an
+        # offsource event, otherwise it is the loudest onsource event, and
+        # zero-lag data must be used
+        slide_id = 0
+        if 'Slide Num' in fp.keys():
+            slide_id = fp['Slide Num'][num_event]
         for ifo in ifos:
             time_shift = fp[ifo+' time shift (s)'][num_event]
             ifo_time = gps_time - time_shift
             files += make_timeseries_plot(workflow, trig_file,
                                           'single', gps_time, ifo_time,
                                           args.output_dir, ifo=ifo,
+                                          seg_files=seg_files,
+                                          veto_file=veto_file,
+                                          slide_id=slide_id,
                                           tags=args.tags + [str(num_event)])
             files += mini.make_qscan_plot(workflow, ifo, ifo_time,
                                           args.output_dir,
                                           tags=args.tags + [str(num_event)])
-    # Handle injections (which are on unslid data)
-    else:
-        for snr_type in ['reweighted', 'coherent']:
-            files += make_timeseries_plot(workflow, trig_file,
-                                          snr_type, gps_time, gps_time,
-                                          args.output_dir, ifo=None,
-                                          tags=args.tags + [str(num_event)])
-        for ifo in ifos:
-            files += mini.make_qscan_plot(workflow, ifo, gps_time,
-                                          args.output_dir,
-                                          tags=args.tags + [str(num_event)])
 
     layouts += list(layout.grouper(files, 2))
 

bin/pygrb/pycbc_pygrb_page_tables

@@ -395,7 +395,7 @@ if lofft_outfile:
         d.append(bestnr)
         td.append(d)
 
-    # th: table header
+    # th: table header [pycbc_pygrb_minifollowups looks for 'Slide Num']
     th = ['Trial', 'Slide Num', 'p-value', 'GPS time',
           'Rec. m1', 'Rec. m2', 'Rec. Mc', 'Rec. spin1z', 'Rec. spin2z',
           'Rec. RA', 'Rec. Dec', 'SNR', 'Chi^2', 'Null SNR']

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,
+                     hide_onsource=True):
     """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 hide_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)